ls_mlkit.util.offload.forward_backward_offload module¶

class ls_mlkit.util.offload.forward_backward_offload.ForwardBackwardOffloadHookContext(model, offload_proportion=0.5, device='cuda', no_split_module_classes=None, with_backward_hook=False, enable=True, num_block: int = 2, strategy='block')[source]¶

Bases: ForwardHookForDevice

static get_backward_hook(pre=True)[source]¶

static get_backward_hook_by_block(info: dict, pre=True, device='cuda')[source]¶

static get_forward_hook_by_block(info: dict, pre=True, device='cuda', with_kwargs=True)[source]¶

static get_partition_block(module_list: list, num_block: int) → dict[source]¶

mode = ('release',)¶

register_forward_hook_by_module(module: Module, parent_name='')[source]¶

register_hook_by_block(module: Module, parent_name='')[source]¶