> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # nemo_automodel.components.distributed ## Subpackages * **[`nemo_automodel.components.distributed.blockdiag_cp`](/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp)** * **[`nemo_automodel.components.distributed.context_parallel`](/nemo-automodel/nemo_automodel/components/distributed/context_parallel)** * **[`nemo_automodel.components.distributed.pipelining`](/nemo-automodel/nemo_automodel/components/distributed/pipelining)** ## Submodules * **[`nemo_automodel.components.distributed.activation_checkpointing`](/nemo-automodel/nemo_automodel/components/distributed/activation_checkpointing)** * **[`nemo_automodel.components.distributed.config`](/nemo-automodel/nemo_automodel/components/distributed/config)** * **[`nemo_automodel.components.distributed.cp_vision_frame_shard`](/nemo-automodel/nemo_automodel/components/distributed/cp_vision_frame_shard)** * **[`nemo_automodel.components.distributed.ddp`](/nemo-automodel/nemo_automodel/components/distributed/ddp)** * **[`nemo_automodel.components.distributed.fsdp2`](/nemo-automodel/nemo_automodel/components/distributed/fsdp2)** * **[`nemo_automodel.components.distributed.fsdp_patches`](/nemo-automodel/nemo_automodel/components/distributed/fsdp_patches)** * **[`nemo_automodel.components.distributed.grad_utils`](/nemo-automodel/nemo_automodel/components/distributed/grad_utils)** * **[`nemo_automodel.components.distributed.init_utils`](/nemo-automodel/nemo_automodel/components/distributed/init_utils)** * **[`nemo_automodel.components.distributed.megatron_fsdp`](/nemo-automodel/nemo_automodel/components/distributed/megatron_fsdp)** * **[`nemo_automodel.components.distributed.mesh`](/nemo-automodel/nemo_automodel/components/distributed/mesh)** * **[`nemo_automodel.components.distributed.mesh_utils`](/nemo-automodel/nemo_automodel/components/distributed/mesh_utils)** * **[`nemo_automodel.components.distributed.model_parallelizer`](/nemo-automodel/nemo_automodel/components/distributed/model_parallelizer)** * **[`nemo_automodel.components.distributed.multimodal_fsdp`](/nemo-automodel/nemo_automodel/components/distributed/multimodal_fsdp)** * **[`nemo_automodel.components.distributed.optimized_tp_plans`](/nemo-automodel/nemo_automodel/components/distributed/optimized_tp_plans)** * **[`nemo_automodel.components.distributed.parallel_styles`](/nemo-automodel/nemo_automodel/components/distributed/parallel_styles)** * **[`nemo_automodel.components.distributed.parallelizer`](/nemo-automodel/nemo_automodel/components/distributed/parallelizer)** * **[`nemo_automodel.components.distributed.parallelizer_utils`](/nemo-automodel/nemo_automodel/components/distributed/parallelizer_utils)** * **[`nemo_automodel.components.distributed.recompute_replay`](/nemo-automodel/nemo_automodel/components/distributed/recompute_replay)** * **[`nemo_automodel.components.distributed.tensor_utils`](/nemo-automodel/nemo_automodel/components/distributed/tensor_utils)** * **[`nemo_automodel.components.distributed.thd_utils`](/nemo-automodel/nemo_automodel/components/distributed/thd_utils)** * **[`nemo_automodel.components.distributed.tp_replicas`](/nemo-automodel/nemo_automodel/components/distributed/tp_replicas)** * **[`nemo_automodel.components.distributed.utils`](/nemo-automodel/nemo_automodel/components/distributed/utils)** ## Package Contents ### Classes | Name | Description | | ---------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------- | | [`DDPConfig`](#nemo_automodel-components-distributed-config-DDPConfig) | Additional configuration for DDP distributed training. | | [`DistInfo`](#nemo_automodel-components-distributed-init_utils-DistInfo) | Holds information about the distributed training environment. | | [`DistributedSetup`](#nemo_automodel-components-distributed-config-DistributedSetup) | Resolved distributed topology and execution policies. | | [`FSDP2Config`](#nemo_automodel-components-distributed-config-FSDP2Config) | Additional configuration for FSDP2 distributed training. | | [`MegatronFSDPConfig`](#nemo_automodel-components-distributed-config-MegatronFSDPConfig) | Additional configuration for MegatronFSDP distributed training. | | [`MeshContext`](#nemo_automodel-components-distributed-mesh-MeshContext) | Resolved runtime context for model parallelization. | | [`MoEParallelizerConfig`](#nemo_automodel-components-distributed-config-MoEParallelizerConfig) | Configuration for MoE model parallelization (EP + FSDP settings). | | [`ModelParallelizer`](#nemo_automodel-components-distributed-parallelizer-ModelParallelizer) | Single model-owned parallelization sidecar contract. | | [`MultimodalDistributedConfig`](#nemo_automodel-components-distributed-config-MultimodalDistributedConfig) | Distributed policies for resolved multimodal modules. | | [`ParallelismSizes`](#nemo_automodel-components-distributed-mesh-ParallelismSizes) | Build-time requested parallelism sizes. | | [`PipelineConfig`](#nemo_automodel-components-distributed-pipelining-config-PipelineConfig) | Configuration for pipeline parallel training. | ### Functions | Name | Description | | ---------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------ | | [`initialize_distributed`](#nemo_automodel-components-distributed-init_utils-initialize_distributed) | Initialize the torch.distributed environment and core model parallel infrastructure. | ### API ```python class nemo_automodel.components.distributed.config.DDPConfig( activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False, activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope = 'all', broadcast_buffers: bool = False, find_unused_parameters: bool = False, static_graph: bool = False, bucket_cap_mb: float | None = None, gradient_as_bucket_view: bool = False, autocast_dtype: torch.dtype | None = None ) ``` Dataclass Additional configuration for DDP distributed training. Note: DDP does not support tensor parallelism, pipeline parallelism, or expert parallelism. Only dp\_size is relevant (inferred from world\_size). **`activation_checkpointing`** `ActivationCheckpointingMode = False` Enable activation checkpointing. `True` or `"full"` keeps the existing full activation checkpointing behavior. `"selective"` wraps transformer blocks with PyTorch selective activation checkpointing. --- **`activation_checkpointing_scope`** `ActivationCheckpointingScope = 'all'` Which extracted layer groups activation checkpointing should wrap. `"all"` selects every extracted group. Scoped values such as `"language"`, `"vision"`, and `"multimodal"` are filtered to trainable layers before generic wrapping. --- **`autocast_dtype`** `dtype | None = None` If set, recipes can wrap the forward pass in `torch.autocast(device_type="cuda", dtype=autocast_dtype)`. Set to `None` to disable. Can be set from YAML as a string (e.g. `autocast_dtype: bfloat16`). --- **`broadcast_buffers`** `bool = False` Synchronize module buffers before each forward. --- **`bucket_cap_mb`** `float | None = None` DDP gradient bucket size in MiB. `None` uses PyTorch's default. --- **`find_unused_parameters`** `bool = False` Forwarded to PyTorch DDP for models with conditionally unused trainable parameters. --- **`gradient_as_bucket_view`** `bool = False` Make gradients views into DDP buckets after the first iteration. --- **`static_graph`** `bool = False` Tell DDP the used/unused parameter set is stable. --- ```python nemo_automodel.components.distributed.config.DDPConfig.__post_init__() ``` ```python nemo_automodel.components.distributed.config.DDPConfig.to_dict() -> typing.Dict[str, typing.Any] ``` Convert config to dictionary. ```python class nemo_automodel.components.distributed.init_utils.DistInfo( backend: str, rank: int, world_size: int, device: torch.device, is_main: bool ) ``` Dataclass Holds information about the distributed training environment. **`backend`** `str` The backend used for torch.distributed (e.g., 'nccl'). --- **`device`** `device` The device assigned to the current process. --- **`is_main`** `bool` True if the process is the main process (rank 0). --- **`rank`** `int` The rank of the current process. --- **`world_size`** `int` The total number of processes. --- ```python class nemo_automodel.components.distributed.config.DistributedSetup( mesh_context: 'MeshContext', strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig | None = None, pipeline_config: 'PipelineConfig | None' = None, moe_parallel_config: 'MoEParallelizerConfig | None' = None, activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False ) ``` Dataclass Resolved distributed topology and execution policies. **`activation_checkpointing`** `ActivationCheckpointingMode = False` --- **`mesh_context`** `'MeshContext'` --- **`moe_parallel_config`** `'MoEParallelizerConfig | None' = None` --- **`pipeline_config`** `'PipelineConfig | None' = None` --- **`strategy_config`** `DistributedStrategyConfig | None = None` --- ```python nemo_automodel.components.distributed.config.DistributedSetup.__post_init__() -> None ``` Keep the compatibility bundle and its MeshContext policy in sync. ```python nemo_automodel.components.distributed.config.DistributedSetup.build( strategy: str | nemo_automodel.components.distributed.config.DistributedStrategyConfig = 'fsdp2', parallelism_sizes: 'ParallelismSizes | None' = None, pipeline_config: 'PipelineConfig | dict | None' = None, moe_parallel_config: 'MoEParallelizerConfig | dict | None' = None, activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False, world_size: int | None = None, timeout_minutes: int | None = None, ranks: list[int] | tuple[int, ...] | None = None ) -> 'DistributedSetup' ``` classmethod Create a resolved distributed setup from sizes and policy configs. Intentionally, this function is forgiving wrt the input types, allowing strings for the strategy and dicts for the pipeline and MoE configs. ```python class nemo_automodel.components.distributed.config.FSDP2Config( sequence_parallel: bool = False, tp_plan: dict | str | None = None, patch_is_packed_sequence: bool = False, mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None = (lambda: MixedPrecisionPoli..., offload_policy: torch.distributed.fsdp.CPUOffloadPolicy | None = None, autocast_dtype: torch.dtype | None = None, activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False, activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope = 'all', defer_fsdp_grad_sync: bool = True, reshard_after_forward: bool | None = None, enable_async_tensor_parallel: bool = False, enable_compile: bool = False, enable_fsdp2_prefetch: bool = False, fsdp2_backward_prefetch_depth: int = 2, fsdp2_forward_prefetch_depth: int = 1, multimodal: nemo_automodel.components.distributed.config.MultimodalDistributedConfig = MultimodalDistributedConfig() ) ``` Dataclass Additional configuration for FSDP2 distributed training. Note: Size parameters (dp\_size, dp\_replicate\_size, tp\_size, pp\_size, cp\_size, ep\_size) are grouped separately in `ParallelismSizes`. **`activation_checkpointing`** `ActivationCheckpointingMode = False` Enable activation checkpointing. `True` or `"full"` keeps the existing full activation checkpointing behavior. `"selective"` wraps transformer blocks with PyTorch selective activation checkpointing. --- **`activation_checkpointing_scope`** `ActivationCheckpointingScope = 'all'` Which extracted layer groups activation checkpointing should wrap. `"all"` selects every extracted group. Scoped values such as `"language"`, `"vision"`, and `"multimodal"` are filtered to trainable layers before generic wrapping. --- **`autocast_dtype`** `dtype | None = None` If set, wraps the forward pass in `torch.autocast(device_type="cuda", dtype=autocast_dtype)`. Use with `output_dtype=float32` in mp\_policy to keep the residual stream in fp32 while running matmuls in lower precision. Set to `None` to disable. Can be set from YAML as a string (e.g. `autocast_dtype: bfloat16`). --- **`defer_fsdp_grad_sync`** `bool = True` Defer FSDP gradient sync to final micro-batch. --- **`enable_async_tensor_parallel`** `bool = False` Enable async tensor parallelism via `torch._inductor.config._micro_pipeline_tp`. Overlaps ReduceScatter with compute in row-parallel layers. Requires `sequence_parallel=True` (forced automatically with a warning if not set). Also enables symmetric memory for the TP group. --- **`enable_compile`** `bool = False` Apply per-layer `torch.compile` to transformer decoder layers (with NO\_REENTRANT activation checkpointing inside each compiled layer). Skips whole-model compile so that checkpoint loading does not produce `_orig_mod` key-prefix mismatches. --- **`enable_fsdp2_prefetch`** `bool = False` Enable explicit forward/backward prefetch chains between FSDP2 sharded layers. Default `True`. --- **`fsdp2_backward_prefetch_depth`** `int = 2` Number of FSDP units to prefetch during backward pass. `2` hides AllGather behind compute; `1` reduces peak memory at a small throughput cost. Default `2`. --- **`fsdp2_forward_prefetch_depth`** `int = 1` Number of FSDP units to prefetch during forward pass. Default `1`. --- **`mp_policy`** `MixedPrecisionPolicy | None` MixedPrecisionPolicy for FSDP2. If `None` (default), uses bf16 forward/backward compute with fp32 gradient reduction, preserving the dtypes of module inputs and outputs. Pair this with `model.torch_dtype: float32` for the Megatron-style fp32 master-weights pattern. Override from YAML using the `_target_` pattern:: mp\_policy: *target*: torch.distributed.fsdp.MixedPrecisionPolicy param\_dtype: bfloat16 reduce\_dtype: float32 output\_dtype: null cast\_forward\_inputs: false See `docs/guides/mixed-precision-training.md` for the full set of recommended patterns and the bf16-storage trap. --- **`multimodal`** `MultimodalDistributedConfig = field(default_factory=MultimodalDistributedConfig)` Policies for resolved multimodal modules that use the text model's distributed axes. --- **`offload_policy`** `CPUOffloadPolicy | None = None` CPUOffloadPolicy for CPU offloading. --- **`patch_is_packed_sequence`** `bool = False` Patch `transformers._is_packed_sequence` to always return Python `False`. This does two things: (1) removes a CPU-GPU sync per attention layer (`aten::is_nonzero` triggered by HF when batch\_size==1), and (2) ensures static attention shapes for `torch.compile`. Safe for standard (non-packed) training only. Disable if using packed-sequence training (position\_ids that reset to 0 mid-sequence). Default `False`. --- **`reshard_after_forward`** `bool | None = None` Override layer-level FSDP2 resharding. `None` preserves AutoModel's heuristic: pipeline-parallel layers do not reshard after forward, while non-pipeline layers reshard all but the last layer. Set `False` for a ZeRO-2-like benchmark where gathered parameters stay resident after forward. Set `True` to force resharding everywhere, including pipeline-parallel layers, which may reduce throughput by adding per-microbatch all-gathers. --- **`sequence_parallel`** `bool = False` Enable sequence parallelism in TP plan. --- **`tp_plan`** `dict | str | None = None` Custom TP plan. If None, auto-selected based on model type. --- ```python nemo_automodel.components.distributed.config.FSDP2Config.__post_init__() -> None ``` ```python nemo_automodel.components.distributed.config.FSDP2Config.to_dict() -> typing.Dict[str, typing.Any] ``` Convert config to dictionary (shallow, preserves policy objects). ```python class nemo_automodel.components.distributed.config.MegatronFSDPConfig( megatron_fsdp_unit_modules: list[str] | None = None, zero_dp_strategy: int = 3, init_fsdp_with_meta_device: bool = False, grad_reduce_in_fp32: bool = False, preserve_fp32_weights: bool = False, overlap_grad_reduce: bool = True, overlap_param_gather: bool = True, check_for_nan_in_grad: bool = True, report_nan_in_param_grad: bool = False, average_in_collective: bool = False, disable_bucketing: bool = False, calculate_per_token_loss: bool = False, keep_fp8_transpose_cache: bool = False, nccl_ub: bool = False, fsdp_double_buffer: bool = False, activation_checkpointing: bool = False ) ``` Dataclass Additional configuration for MegatronFSDP distributed training. Note: Size parameters (dp\_size, tp\_size, cp\_size) are grouped separately in `ParallelismSizes`. MegatronFSDP does not support pp\_size, dp\_replicate\_size, or ep\_size. **`activation_checkpointing`** `bool = False` Enable activation checkpointing for transformer MLP layers to save memory. --- **`average_in_collective`** `bool = False` Average in collective if True. --- **`calculate_per_token_loss`** `bool = False` Calculate per token loss if True. --- **`check_for_nan_in_grad`** `bool = True` Legacy buffer-level gradient NaN check. BREAKING CHANGE on megatron-fsdp 0.5.0: this flag is a no-op, preserved only for config compatibility. 0.5.0 removed the buffer-level NaN check entirely, so gradient NaN checking is now OFF regardless of this value; a truthy value is dropped with a one-time warning. The default is kept True for config compatibility, but has no effect. To restore gradient NaN checking, enable report\_nan\_in\_param\_grad instead. --- **`disable_bucketing`** `bool = False` Disable bucketing if True. --- **`fsdp_double_buffer`** `bool = False` Use double buffer if True. --- **`grad_reduce_in_fp32`** `bool = False` Reduce gradients in fp32 if True. --- **`init_fsdp_with_meta_device`** `bool = False` Initialize MegatronFSDP with meta device if True. --- **`keep_fp8_transpose_cache`** `bool = False` Keep fp8 transpose cache when using custom FSDP if True. --- **`megatron_fsdp_unit_modules`** `list[str] | None = None` Class paths of the submodules to wrap as individual MegatronFSDP units. When `None` (the default), the wrap classes are auto-derived from the model's `_no_split_modules` so the real instantiated block classes are used regardless of backend (HF or NeMo-custom). --- **`nccl_ub`** `bool = False` Use NCCL UBs if True. --- **`overlap_grad_reduce`** `bool = True` Overlap gradient reduction if True. --- **`overlap_param_gather`** `bool = True` Overlap parameter gathering if True. --- **`preserve_fp32_weights`** `bool = False` Preserve fp32 weights if True. --- **`report_nan_in_param_grad`** `bool = False` Enable megatron-fsdp 0.5.0's precise per-parameter gradient NaN check. This is the replacement for the removed check\_for\_nan\_in\_grad and is OFF by default; enabling it can significantly reduce training throughput. --- **`zero_dp_strategy`** `int = 3` Data parallel sharding strategy. --- ```python nemo_automodel.components.distributed.config.MegatronFSDPConfig.to_dict() -> typing.Dict[str, typing.Any] ``` Convert config to dictionary (shallow, preserves objects). ```python class nemo_automodel.components.distributed.mesh.MeshContext( device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None, moe_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None, process_group: torch.distributed.ProcessGroup | None = None, strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig | None = None, moe_parallel_config: nemo_automodel.components.distributed.config.MoEParallelizerConfig | None = None, activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False, reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None, cp_ring_degree: int = 1, cp_ulysses_degree: int | None = None ) ``` Dataclass Resolved runtime context for model parallelization. Parallelism sizes (`tp_size`, `pp_size`, etc.) are **not** stored as fields; they are `@property` accessors that read directly from the attached `DeviceMesh` / `moe_mesh`. When no mesh is present the properties return safe defaults (`1` for sizes, `None` for dp / hsdp). All `DeviceMesh` objects passed in must use axis names from `MeshAxisName`; a `ValueError` is raised on construction if any unknown name is encountered. ## Lifecycle 1. Recipes parse YAML to obtain sizes and strategy configs. 2. Sizes are passed to `build` to build `DeviceMesh` objects. 3. `MeshContext` is created with those meshes; axis names are validated automatically in `__post_init__`. Alternatively, `from_meshes` constructs an instance directly from `DeviceMesh` objects (used by `NeMoAutoModel.from_pretrained`). **`activation_checkpointing`** `ActivationCheckpointingMode = False` Resolved activation-checkpointing mode. --- **`cp_ring_degree`** `int = 1` Ring-attention factor of the context-parallel axis. --- **`cp_size`** `int` Context-parallel degree (from `device_mesh`, default `1`). --- **`cp_ulysses_degree`** `int | None = None` Ulysses factor of the context-parallel axis. `None` lets the consumer use the full context-parallel axis. --- **`device_mesh`** `Optional[DeviceMesh] = field(default=None, repr=False)` Device mesh for distributed training. --- **`dp_replicate_size`** `int | None` HSDP replication degree (from `device_mesh`, default `None`). --- **`dp_shard_size`** `int` DP shard degree (from `device_mesh`, default `1`). --- **`dp_size`** `int | None` Data-parallel degree (from `device_mesh`, default `None`). --- **`ep_size`** `int` Expert-parallel degree (from `moe_mesh`, default `1`). --- **`moe_mesh`** `Optional[DeviceMesh] = field(default=None, repr=False)` MoE-specific device mesh. --- **`moe_parallel_config`** `MoEParallelizerConfig | None = None` Optional MoE-specific parallelization policy. --- **`pp_enabled`** `bool` `True` when `pp_size > 1`. --- **`pp_size`** `int` Pipeline-parallel degree (from `device_mesh`, default `1`). --- **`process_group`** `ProcessGroup | None = field(default=None, repr=False)` Optional model-local group for recipe-level collectives that must not involve ranks outside this mesh. --- **`reapply_trainability`** `Callable[[nn.Module], None] | None = field(default=None, repr=False)` Optional callback run after model surgery and before distributed wrapping. --- **`strategy_config`** `DistributedStrategyConfig | None = None` Resolved DDP, FSDP2, or Megatron-FSDP policy. --- **`tp_size`** `int` Tensor-parallel degree (from `device_mesh`, default `1`). --- ```python nemo_automodel.components.distributed.mesh.MeshContext.__post_init__() -> None ``` ```python nemo_automodel.components.distributed.mesh.MeshContext._dp_axis_names() -> typing.Tuple[str, ...] ``` DP axis names for FSDP mesh slicing. ```python nemo_automodel.components.distributed.mesh.MeshContext.build( strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig, parallelism_sizes: nemo_automodel.components.distributed.mesh.ParallelismSizes | None = None, moe_parallel_config: nemo_automodel.components.distributed.config.MoEParallelizerConfig | None = None, activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False, reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None, cp_ring_degree: int = 1, cp_ulysses_degree: int | None = None, world_size: int | None = None, timeout_minutes: int | None = None, ranks: typing.Sequence[int] | None = None ) -> nemo_automodel.components.distributed.mesh.MeshContext ``` classmethod Build a resolved `MeshContext` from sizes and execution policy. **Parameters:** **`strategy_config`** `DistributedStrategyConfig` Already-instantiated distributed strategy config. --- **`parallelism_sizes`** `ParallelismSizes | None` — default: None Requested data, tensor, pipeline, context, and expert parallelism sizes. If `None`, defaults to no parallelism with DP inferred from `world_size`. --- **`moe_parallel_config`** `MoEParallelizerConfig | None` — default: None Optional MoE-specific parallelization policy. --- **`activation_checkpointing`** `ActivationCheckpointingMode` — default: False Activation-checkpointing mode applied by the model sidecar. --- **`reapply_trainability`** `Callable[[nn.Module], None] | None` — default: None Optional callback run after model surgery and before distributed wrapping. --- **`cp_ring_degree`** `int` — default: 1 Ring-attention factor of the context-parallel axis. --- **`cp_ulysses_degree`** `int | None` — default: None Ulysses factor of the context-parallel axis. --- **`world_size`** `int | None` — default: None Total process count. If `None`, inferred from the distributed environment. --- **`timeout_minutes`** `int | None` — default: None Optional timeout for process groups created by `DeviceMesh` axes. --- **`ranks`** `Sequence[int] | None` — default: None Optional ordered global ranks to use for the mesh. When omitted, the mesh uses every rank in the default process group. --- ```python nemo_automodel.components.distributed.mesh.MeshContext.from_meshes( device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh], moe_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None, strategy_config: nemo_automodel.components.distributed.config.DistributedStrategyConfig | None = None, moe_parallel_config: nemo_automodel.components.distributed.config.MoEParallelizerConfig | None = None, activation_checkpointing: nemo_automodel.components.distributed.config.ActivationCheckpointingMode = False, reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None, cp_ring_degree: int = 1, cp_ulysses_degree: int | None = None ) -> nemo_automodel.components.distributed.mesh.MeshContext ``` classmethod Build a `MeshContext` from `DeviceMesh` objects. This is the entry-point used by `NeMoAutoModel.from_pretrained` / `from_config` where the caller has raw meshes rather than a parsed YAML config. **Parameters:** **`device_mesh`** `Optional[DeviceMesh]` Primary device mesh. --- **`moe_mesh`** `Optional[DeviceMesh]` — default: None Optional MoE-specific device mesh. --- **`strategy_config`** `DistributedStrategyConfig | None` — default: None Resolved DDP, FSDP2, or Megatron-FSDP policy. --- **`moe_parallel_config`** `MoEParallelizerConfig | None` — default: None Optional MoE-specific parallelization policy. --- **`activation_checkpointing`** `ActivationCheckpointingMode` — default: False Activation-checkpointing mode applied by the model sidecar. --- **`reapply_trainability`** `Callable[[nn.Module], None] | None` — default: None Optional callback run after model surgery and before distributed wrapping. --- **`cp_ring_degree`** `int` — default: 1 Ring-attention factor of the context-parallel axis. --- **`cp_ulysses_degree`** `int | None` — default: None Ulysses factor of the context-parallel axis. --- ```python nemo_automodel.components.distributed.mesh.MeshContext.parallelize_axis_kwargs() -> typing.Dict[str, object] ``` Axis-name kwargs for `parallelize_fn` (EP/FSDP, no `pp_axis_name`). ```python nemo_automodel.components.distributed.mesh.MeshContext.pipeline_axis_kwargs() -> typing.Dict[str, object] ``` Axis-name kwargs for `AutoPipeline`. ```python class nemo_automodel.components.distributed.config.MoEParallelizerConfig( ignore_router_for_ac: bool = True, reshard_after_forward: bool = False, lm_head_precision: typing.Union[str, torch.dtype] | None = None, wrap_outer_model: bool = True, mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None = None ) ``` Dataclass Configuration for MoE model parallelization (EP + FSDP settings). **`ignore_router_for_ac`** `bool = True` --- **`lm_head_precision`** `Union[str, dtype] | None = None` --- **`mp_policy`** `MixedPrecisionPolicy | None = None` --- **`reshard_after_forward`** `bool = False` --- **`wrap_outer_model`** `bool = True` --- ```python nemo_automodel.components.distributed.config.MoEParallelizerConfig.to_dict() -> typing.Dict[str, typing.Any] ``` ```python class nemo_automodel.components.distributed.parallelizer.ModelParallelizer() ``` Single model-owned parallelization sidecar contract. ```python nemo_automodel.components.distributed.parallelizer.ModelParallelizer._apply( model: torch.nn.Module, device_mesh: torch.distributed.device_mesh.DeviceMesh, mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None = None, offload_policy: torch.distributed.fsdp.OffloadPolicy | None = None, sequence_parallel: bool = False, activation_checkpointing: bool = False, tp_shard_plan: typing.Union[typing.Dict[str, torch.distributed.tensor.parallel.ParallelStyle], str] | None = None, dp_replicate_mesh_name: str = 'dp_replicate', dp_shard_cp_mesh_name: str = 'dp_shard_cp', tp_mesh_name: str = 'tp', enable_async_tensor_parallel: bool = False, enable_compile: bool = False, enable_fsdp2_prefetch: bool = True, fsdp2_backward_prefetch_depth: int = 2, fsdp2_forward_prefetch_depth: int = 1, reshard_after_forward: bool | None = None, activation_checkpointing_scope: nemo_automodel.components.distributed.config.ActivationCheckpointingScope | None = 'all', frozen_multimodal_sharding: nemo_automodel.components.distributed.multimodal_fsdp.FrozenMultimodalSharding = 'root', reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None ) -> torch.nn.Module ``` Apply the shared dense FSDP2 implementation. ```python nemo_automodel.components.distributed.parallelizer.ModelParallelizer._apply_fsdp_sharding( module: torch.nn.Module, mesh: torch.distributed.device_mesh.DeviceMesh, mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None, offload_policy: torch.distributed.fsdp.OffloadPolicy | None = None, enable_fsdp2_prefetch: bool = True, fsdp2_backward_prefetch_depth: int = 2, fsdp2_forward_prefetch_depth: int = 1, reshard_after_forward: bool | None = None, frozen_multimodal_sharding: nemo_automodel.components.distributed.multimodal_fsdp.FrozenMultimodalSharding = 'root', ignored_multimodal_params: set[torch.nn.Parameter] | None = None ) -> None ``` Wrap the model's submodules into FSDP2 units. Model-owned sidecar strategies deriving from this class override this hook to change how parameters are grouped into FSDP units without reimplementing the surrounding TP/AC/mixed-precision flow. Override `_fully_shard_module` when a model needs a different primitive. ```python nemo_automodel.components.distributed.parallelizer.ModelParallelizer._fully_shard_module( module: torch.nn.Module, kwargs = {} ) -> torch.nn.Module ``` Apply the FSDP2 primitive used by this model sidecar. ```python nemo_automodel.components.distributed.parallelizer.ModelParallelizer._fully_shard_untied_input_output_embeddings( model: torch.nn.Module, mesh: torch.distributed.device_mesh.DeviceMesh, mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None, offload_policy: torch.distributed.fsdp.OffloadPolicy | None, input_reshard_after_forward: bool ) -> None ``` Give large trainable untied embedding tables independent FSDP buffers. The generic dense path otherwise leaves both tables in the root FSDP unit. With fp32 gradient reduction, that unit allocates one contiguous reduce-scatter input containing both gradients. Keeping the two trainable leaf modules in separate FSDP units bounds that allocation by the larger table instead of their sum. This pass skips tied tables to preserve aliasing, and frozen tables because they have no gradient communication buffer to split. Tables already wrapped as FSDP units retain their existing reshard policy. With an explicit `reshard_after_forward=True`, a container-hosted head reshards while a newly split top-level head stays gathered. FusedLinearCrossEntropy handles the sharded weight with correct gradients, at the cost of an extra all-gather. **Parameters:** **`model`** `nn.Module` Model whose input and output embedding modules may be sharded. --- **`mesh`** `DeviceMesh` Device mesh that owns the FSDP shards. --- **`mp_policy`** `MixedPrecisionPolicy | None` Mixed-precision policy used by the surrounding FSDP units. --- **`offload_policy`** `OffloadPolicy | None` Optional offload policy used by the surrounding FSDP units. --- **`input_reshard_after_forward`** `bool` Whether the input embedding unit reshards its parameters after forward. --- ```python nemo_automodel.components.distributed.parallelizer.ModelParallelizer._use_full_layer_activation_checkpointing( model: torch.nn.Module ) -> bool ``` Return whether this model safely opts into whole-layer checkpointing. ```python nemo_automodel.components.distributed.parallelizer.ModelParallelizer._validate_tp_mesh( model: torch.nn.Module, tp_mesh: torch.distributed.device_mesh.DeviceMesh ) -> None ``` Validate the model's attention topology against its TP mesh. ```python nemo_automodel.components.distributed.parallelizer.ModelParallelizer.parallelize( model: torch.nn.Module, mesh_context: nemo_automodel.components.distributed.mesh.MeshContext ) -> torch.nn.Module ``` Apply every requested parallelism and return the parallelized model. ```python class nemo_automodel.components.distributed.config.MultimodalDistributedConfig( frozen_sharding: nemo_automodel.components.distributed.multimodal_fsdp.FrozenMultimodalSharding = 'root', vision: nemo_automodel.components.distributed.config.MultimodalVisionConfig = MultimodalVisionConfig() ) ``` Dataclass Distributed policies for resolved multimodal modules. **`frozen_sharding`** `FrozenMultimodalSharding = 'root'` Controls fully frozen multimodal modules such as vision/audio towers and projectors. `"root"` (default) keeps their parameters in an always-run outer FSDP root, which is safe when modality execution differs across ranks. `"per_layer"` uses normal layer/container FSDP units and requires every rank in the FSDP group to execute or skip the module identically on every microbatch. `"replicate"` excludes the frozen parameters from FSDP roots so each rank keeps a full copy. Modules with any trainable parameters use normal layer/container sharding regardless of this setting. --- **`vision`** `MultimodalVisionConfig = field(default_factory=MultimodalVisionConfig)` Policies that apply specifically to resolved vision modules. --- ```python nemo_automodel.components.distributed.config.MultimodalDistributedConfig.__post_init__() -> None ``` ```python class nemo_automodel.components.distributed.mesh.ParallelismSizes( dp_size: int | None = None, dp_replicate_size: int | None = None, tp_size: int = 1, pp_size: int = 1, cp_size: int = 1, ep_size: int = 1 ) ``` Dataclass Build-time requested parallelism sizes. This is durable user intent, not runtime topology. `MeshContext` derives its size properties from live `DeviceMesh` objects after build. **`cp_size`** `int = 1` --- **`dp_replicate_size`** `int | None = None` --- **`dp_size`** `int | None = None` --- **`ep_size`** `int = 1` --- **`pp_size`** `int = 1` --- **`tp_size`** `int = 1` --- ```python class nemo_automodel.components.distributed.pipelining.config.PipelineConfig( pp_schedule: str | None = '1f1b', pp_schedule_csv: str | None = None, pp_microbatch_size: int = 1, pp_batch_size: int = 1, layers_per_stage: int | None = None, round_virtual_stages_to_pp_multiple: typing.Literal['up', 'down'] | None = None, module_fqns_per_model_part: typing.List[typing.List[str]] | None = None, patch_inner_model: bool = True, patch_causal_lm_model: bool = True, patch_stage_backward_maybe_with_nosync: bool = False, dtype: torch.dtype | None = None, scale_grads_in_schedule: bool = False, loss_fn: typing.Callable | None = None, pp_seq_len: int | None = None, pp_recv_buffer_pool: bool = False, pp_recv_buffer_pool_slack: int = 2 ) ``` Dataclass Configuration for pipeline parallel training. Note: Device mesh (world\_mesh, moe\_mesh) is passed separately on the from\_pretrained/from\_config method signature. Pipeline parallelism is enabled when pp\_size > 1. Axis names are inferred automatically from the device mesh structure. **`dtype`** `dtype | None = None` Data type for pipeline computation. If None, uses the model's default dtype. --- **`layers_per_stage`** `int | None = None` Number of transformer layers per pipeline stage. If None, layers are split evenly across stages. --- **`loss_fn`** `Callable | None = None` Loss function used for pipeline training. Required when pipeline is enabled. The function signature should be compatible with the model's output format. --- **`module_fqns_per_model_part`** `List[List[str]] | None = None` Explicit specification of which module FQNs belong to each model part/stage. If provided, overrides automatic layer splitting. --- **`patch_causal_lm_model`** `bool = True` Apply pipeline patches to the CausalLM wrapper model. Defaults to True. --- **`patch_inner_model`** `bool = True` Apply pipeline patches to the inner model (e.g., the base transformer in a CausalLM wrapper). Defaults to True. --- **`patch_stage_backward_maybe_with_nosync`** `bool = False` Patch stage backward to use no\_sync context for gradient accumulation efficiency. Useful when combining PP with FSDP. --- **`pp_batch_size`** `int = 1` Total batch size per pipeline stage. Must be divisible by pp\_microbatch\_size. --- **`pp_microbatch_size`** `int = 1` Size of each microbatch for pipeline execution. pp\_batch\_size must be divisible by pp\_microbatch\_size. --- **`pp_recv_buffer_pool`** `bool = False` Pool the per-microbatch P2P recv buffers into a ring sized to the 1F1B in-flight depth instead of pre-allocating one buffer set per microbatch and direction. Cuts pipeline-stage buffer memory roughly by num\_microbatches / (num\_stages - stage\_index + slack) on early/middle stages. For native interleaved 1F1B, receives are pooled using the lowered schedule's actual lifetimes. Inputs are copied into independent leaves before forward, and gradients are copied before backward to protect checkpoint recomputation and escaping aliases. Other schedules are ignored with a warning. Defaults to False. --- **`pp_recv_buffer_pool_slack`** `int = 2` Extra buffer sets beyond the 1F1B in-flight depth when pp\_recv\_buffer\_pool is enabled. Defaults to 2. --- **`pp_schedule`** `str | None = '1f1b'` Pipeline schedule type. Supported values: "1f1b" (one-forward-one-backward), "gpipe", "interleaved\_1f1b", "looped\_bfs", "dfs", "v\_schedule", "zero\_bubble". Defaults to "1f1b". --- **`pp_schedule_csv`** `str | None = None` Path to a CSV file defining a custom pipeline schedule. If provided, overrides pp\_schedule. --- **`pp_seq_len`** `int | None = None` Sequence length hint for pipeline parallel shape inference. If None, it will be inferred from the dataset config. --- **`round_virtual_stages_to_pp_multiple`** `Literal['up', 'down'] | None = None` When using virtual stages (interleaved schedules), round the number of virtual stages to a multiple of pp\_size. "up" rounds up, "down" rounds down. If None, no rounding is applied. --- **`scale_grads_in_schedule`** `bool = False` Scale gradients within the pipeline schedule (by 1/n\_microbatches). If False, gradients must be scaled externally. Defaults to False. --- ```python nemo_automodel.components.distributed.pipelining.config.PipelineConfig.__post_init__() -> None ``` Validate the recv-buffer-pool knobs. ```python nemo_automodel.components.distributed.pipelining.config.PipelineConfig.to_dict() -> typing.Dict[str, typing.Any] ``` Convert config to dictionary. ```python nemo_automodel.components.distributed.init_utils.initialize_distributed( backend, timeout_minutes = 1 ) ``` Initialize the torch.distributed environment and core model parallel infrastructure. This function sets the device based on the local rank, configures the process group, and calls torch.distributed.init\_process\_group with the appropriate parameters. It also registers a cleanup function to properly destroy the process group at exit. **Parameters:** **`backend`** `str` The backend to use for torch.distributed (e.g., 'nccl'). --- **`timeout_minutes`** `int` — default: 1 Timeout (in minutes) for distributed initialization. Defaults to 1. --- **Returns:** An instance containing the distributed environment configuration. ## Docs - [nemo_automodel.components.distributed.activation_checkpointing](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/activation_checkpointing.md) - [nemo_automodel.components.distributed.blockdiag_cp](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp.md) - [nemo_automodel.components.distributed.blockdiag_cp.batch](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/batch.md) - [nemo_automodel.components.distributed.blockdiag_cp.exchange](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/exchange.md) - [nemo_automodel.components.distributed.blockdiag_cp.kernels](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/kernels.md) - [nemo_automodel.components.distributed.blockdiag_cp.packed](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/packed.md) - [nemo_automodel.components.distributed.blockdiag_cp.runtime](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/runtime.md) - [nemo_automodel.components.distributed.blockdiag_cp.state](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/blockdiag_cp/state.md) - [nemo_automodel.components.distributed.config](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/config.md) - [nemo_automodel.components.distributed.context_parallel](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel.md) - [nemo_automodel.components.distributed.context_parallel.magi](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/magi.md) - [nemo_automodel.components.distributed.context_parallel.mamba](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/mamba.md) - [nemo_automodel.components.distributed.context_parallel.sharder](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/sharder.md) - [nemo_automodel.components.distributed.context_parallel.utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/context_parallel/utils.md) - [nemo_automodel.components.distributed.cp_vision_frame_shard](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/cp_vision_frame_shard.md) - [nemo_automodel.components.distributed.ddp](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/ddp.md) - [nemo_automodel.components.distributed.fsdp_patches](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/fsdp_patches.md) - [nemo_automodel.components.distributed.fsdp2](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/fsdp2.md) - [nemo_automodel.components.distributed.grad_utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/grad_utils.md) - [nemo_automodel.components.distributed.init_utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/init_utils.md) - [nemo_automodel.components.distributed.megatron_fsdp](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/megatron_fsdp.md) - [nemo_automodel.components.distributed.mesh](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/mesh.md) - [nemo_automodel.components.distributed.mesh_utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/mesh_utils.md) - [nemo_automodel.components.distributed.model_parallelizer](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/model_parallelizer.md) - [nemo_automodel.components.distributed.multimodal_fsdp](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/multimodal_fsdp.md) - [nemo_automodel.components.distributed.optimized_tp_plans](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/optimized_tp_plans.md) - [nemo_automodel.components.distributed.parallel_styles](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/parallel_styles.md) - [nemo_automodel.components.distributed.parallelizer](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/parallelizer.md) - [nemo_automodel.components.distributed.parallelizer_utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/parallelizer_utils.md) - [nemo_automodel.components.distributed.pipelining](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining.md) - [nemo_automodel.components.distributed.pipelining.autopipeline](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/autopipeline.md) - [nemo_automodel.components.distributed.pipelining.config](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/config.md) - [nemo_automodel.components.distributed.pipelining.functional](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/functional.md) - [nemo_automodel.components.distributed.pipelining.hf_utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/hf_utils.md) - [nemo_automodel.components.distributed.pipelining.recv_buffer_pool](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/recv_buffer_pool.md) - [nemo_automodel.components.distributed.pipelining.runtime](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/pipelining/runtime.md) - [nemo_automodel.components.distributed.recompute_replay](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/recompute_replay.md) - [nemo_automodel.components.distributed.tensor_utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/tensor_utils.md) - [nemo_automodel.components.distributed.thd_utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/thd_utils.md) - [nemo_automodel.components.distributed.tp_replicas](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/tp_replicas.md) - [nemo_automodel.components.distributed.utils](https://docs.nvidia.com/nemo/automodel/v0.4/nemo-automodel/nemo_automodel/components/distributed/utils.md)