nemo_automodel.components.distributed.model_parallelizer

Resolve and execute model-owned parallelization sidecars.

Module Contents

Functions

Name	Description
`_apply_model_parallelizer`	Execute shared strategy dispatch for one model-owned parallelizer.
`_parallelize_ddp`	-
`_parallelize_fsdp2`	-
`_parallelize_megatron_fsdp`	-
`_parallelize_moe`	-
`_parallelize_unsharded_fsdp2`	-
`compile_parallelized_model`	Compile FSDP2 layers after parallelization when requested.
`get_model_parallelizer`	Return the class-owned sidecar, or the shared default implementation.
`parallelize_model`	Apply all requested parallelisms through the model-owned contract.

Data

_DEFAULT_PARALLELIZER

__all__

API

nemo_automodel.components.distributed.model_parallelizer._apply_model_parallelizer(
    parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer,
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

Execute shared strategy dispatch for one model-owned parallelizer.

nemo_automodel.components.distributed.model_parallelizer._parallelize_ddp(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer._parallelize_fsdp2(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext,
    parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer._parallelize_megatron_fsdp(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer._parallelize_moe(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext,
    parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer._parallelize_unsharded_fsdp2(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext,
    parallelizer: nemo_automodel.components.distributed.parallelizer.ModelParallelizer
) -> torch.nn.Module

nemo_automodel.components.distributed.model_parallelizer.compile_parallelized_model(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> None

Compile FSDP2 layers after parallelization when requested.

nemo_automodel.components.distributed.model_parallelizer.get_model_parallelizer(
    model: torch.nn.Module
) -> nemo_automodel.components.distributed.parallelizer.ModelParallelizer

Return the class-owned sidecar, or the shared default implementation.

nemo_automodel.components.distributed.model_parallelizer.parallelize_model(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> torch.nn.Module

Apply all requested parallelisms through the model-owned contract.

nemo_automodel.components.distributed.model_parallelizer._DEFAULT_PARALLELIZER = ModelParallelizer()

nemo_automodel.components.distributed.model_parallelizer.__all__ = ['ModelParallelizer', 'compile_parallelized_model', 'get_model_parallelizer', 'p...