flex_model package

Subpackages

• Core: Wrapper and HookFunction
  • Modules
    • flex_model.core.FlexModel
      • FlexModel
    • flex_model.core.HookFunction
      • HookFunction
    • FlexModel
      • FlexModel.forward()
      • FlexModel.get_hook_function_groups()
      • FlexModel.get_group_hook_functions()
      • FlexModel.update_hook_groups()
      • FlexModel.remove_hook_groups()
      • FlexModel.create_hook_group()
      • FlexModel.register_forward_hook()
      • FlexModel.register_full_backward_hook()
      • FlexModel.register_hook()
      • FlexModel.register_forward_pre_hook()
      • FlexModel.register_full_backward_pre_hook()
      • FlexModel.register_trainable_module()
      • FlexModel.get_module_parameter()
      • FlexModel.wrapped_module_names
      • FlexModel.trainable_modules_names
      • FlexModel.all_modules_names
      • FlexModel.named_parameters()
      • FlexModel.named_buffers()
      • FlexModel.restore()
    • HookFunction
  • Miscellaneous
    • DummyModule
• Distributed: Backend, mappings and strategies
  • Distributed API
    • flex_model.distributed.initialize_distributed_state
      • initialize_distributed_state()
  • Mappings
    • flex_model.distributed.broadcast_tensor_parallel
      • broadcast_tensor_parallel()
    • flex_model.distributed.broadcast_data_parallel
      • broadcast_data_parallel()
    • flex_model.distributed.all_gather_tensor_parallel
      • all_gather_tensor_parallel()
    • flex_model.distributed.all_gather_data_parallel
      • all_gather_data_parallel()
    • flex_model.distributed.scatter_tensor_parallel
      • scatter_tensor_parallel()
    • flex_model.distributed.scatter_data_parallel
      • scatter_data_parallel()
    • flex_model.distributed.gather_pipeline_parallel_tensor_dicts
      • gather_pipeline_parallel_tensor_dicts()
  • Strategies
    • flex_model.distributed.BaseRoutingStrategy
      • BaseRoutingStrategy
    • flex_model.distributed.ParameterTensorParallelRoutingStrategy
      • ParameterTensorParallelRoutingStrategy
    • flex_model.distributed.ActivationTensorAllToAllRoutingStrategy
      • ActivationTensorAllToAllRoutingStrategy
    • flex_model.distributed.BaseOffloadStrategy
      • BaseOffloadStrategy
    • flex_model.distributed.NullMemoryOffloadStrategy
      • NullMemoryOffloadStrategy
    • flex_model.distributed.CPUPinnedMemoryOffloadStrategy
      • CPUPinnedMemoryOffloadStrategy
    • flex_model.distributed.CPUPagedMemoryOffloadStrategy
      • CPUPagedMemoryOffloadStrategy
    • flex_model.distributed.GPUMemoryOffloadStrategy
      • GPUMemoryOffloadStrategy
    • flex_model.distributed.BaseFunctionStrategy
      • BaseFunctionStrategy
    • flex_model.distributed.NonValidatedFunctionStrategy
      • NonValidatedFunctionStrategy
• Traverse: Pytree traversal utility
  • Ops
    • flex_model.traverse.flatten
      • flatten()
    • flex_model.traverse.unflatten
      • unflatten()
  • Nodes
    • flex_model.traverse.InternalObject
      • InternalObject
    • flex_model.traverse.LeafObject
      • LeafObject
    • flex_model.traverse.ScalarObject
      • ScalarObject
    • flex_model.traverse.InternalNode
      • InternalNode
    • flex_model.traverse.LeafNode
      • LeafNode
    • flex_model.traverse.ScalarNode
      • ScalarNode

Submodules

flex_model.utils module

flex_model.utils.setup_logger(level)

Module contents