Index

_

__init__() (flex_model.core.FlexModel method) (flex_model.core.HookFunction method) (flex_model.distributed.ActivationTensorAllToAllRoutingStrategy method) (flex_model.distributed.BaseFunctionStrategy method) (flex_model.distributed.BaseOffloadStrategy method) (flex_model.distributed.BaseRoutingStrategy method) (flex_model.distributed.CPUPagedMemoryOffloadStrategy method) (flex_model.distributed.CPUPinnedMemoryOffloadStrategy method) (flex_model.distributed.GPUMemoryOffloadStrategy method) (flex_model.distributed.NonValidatedFunctionStrategy method) (flex_model.distributed.NullMemoryOffloadStrategy method) (flex_model.distributed.ParameterTensorParallelRoutingStrategy method) (flex_model.traverse.InternalNode method) (flex_model.traverse.LeafNode method)

A

ActivationTensorAllToAllRoutingStrategy (class in flex_model.distributed) all_gather_data_parallel() (in module flex_model.distributed)

all_gather_tensor_parallel() (in module flex_model.distributed) all_modules_names (flex_model.core.FlexModel property)

B

BaseFunctionStrategy (class in flex_model.distributed) BaseOffloadStrategy (class in flex_model.distributed)

BaseRoutingStrategy (class in flex_model.distributed) broadcast_data_parallel() (in module flex_model.distributed) broadcast_tensor_parallel() (in module flex_model.distributed)

C

CPUPagedMemoryOffloadStrategy (class in flex_model.distributed)

CPUPinnedMemoryOffloadStrategy (class in flex_model.distributed) create_hook_group() (flex_model.core.FlexModel method)

D

DummyModule (class in flex_model.core)

F

flatten() (in module flex_model.traverse) flex_model module

flex_model.utils module FlexModel (class in flex_model.core), [1] forward() (flex_model.core.FlexModel method)

G

gather_pipeline_parallel_tensor_dicts() (in module flex_model.distributed) get_group_hook_functions() (flex_model.core.FlexModel method)

get_hook_function_groups() (flex_model.core.FlexModel method) get_module_parameter() (flex_model.core.FlexModel method) GPUMemoryOffloadStrategy (class in flex_model.distributed)

H

HookFunction (class in flex_model.core), [1]

I

initialize_distributed_state() (in module flex_model.distributed)

InternalNode (class in flex_model.traverse) InternalObject (in module flex_model.traverse)

L

LeafNode (class in flex_model.traverse)

LeafObject (in module flex_model.traverse)

M

module flex_model flex_model.utils

N

named_buffers() (flex_model.core.FlexModel method) named_parameters() (flex_model.core.FlexModel method)

NonValidatedFunctionStrategy (class in flex_model.distributed) NullMemoryOffloadStrategy (class in flex_model.distributed)

P

ParameterTensorParallelRoutingStrategy (class in flex_model.distributed)

R

register_forward_hook() (flex_model.core.FlexModel method) register_forward_pre_hook() (flex_model.core.FlexModel method) register_full_backward_hook() (flex_model.core.FlexModel method) register_full_backward_pre_hook() (flex_model.core.FlexModel method)

register_hook() (flex_model.core.FlexModel method) register_trainable_module() (flex_model.core.FlexModel method) remove_hook_groups() (flex_model.core.FlexModel method) restore() (flex_model.core.FlexModel method)

S

ScalarNode (in module flex_model.traverse) ScalarObject (in module flex_model.traverse)

scatter_data_parallel() (in module flex_model.distributed) scatter_tensor_parallel() (in module flex_model.distributed) setup_logger() (in module flex_model.utils)

T

trainable_modules_names (flex_model.core.FlexModel property)

U

unflatten() (in module flex_model.traverse)

update_hook_groups() (flex_model.core.FlexModel method)

W

wrapped_module_names (flex_model.core.FlexModel property)