nemo_automodel.components.distributed

Enable activation checkpointing. True or "full" keeps the existing full activation checkpointing behavior. "selective" wraps transformer blocks with PyTorch selective activation checkpointing.

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.

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).

Synchronize module buffers before each forward.

DDP gradient bucket size in MiB. None uses PyTorch’s default.

Forwarded to PyTorch DDP for models with conditionally unused trainable parameters.

Make gradients views into DDP buckets after the first iteration.

Tell DDP the used/unused parameter set is stable.

The backend used for torch.distributed (e.g., ‘nccl’).

The device assigned to the current process.

True if the process is the main process (rank 0).

The rank of the current process.

The total number of processes.

Enable activation checkpointing. True or "full" keeps the existing full activation checkpointing behavior. "selective" wraps transformer blocks with PyTorch selective activation checkpointing.

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.

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 gradient sync to final micro-batch.

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.

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 explicit forward/backward prefetch chains between FSDP2 sharded layers. Default True.

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.

Number of FSDP units to prefetch during forward pass. Default 1.

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.

Policies for resolved multimodal modules that use the text model’s distributed axes.

CPUOffloadPolicy for CPU offloading.

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.

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.

Enable sequence parallelism in TP plan.

Custom TP plan. If None, auto-selected based on model type.

Enable activation checkpointing for transformer MLP layers to save memory.

Average in collective if True.

Calculate per token loss if 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 if True.

Use double buffer if True.

Reduce gradients in fp32 if True.

Initialize MegatronFSDP with meta device if True.

Keep fp8 transpose cache when using custom FSDP if True.

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).

Use NCCL UBs if True.

Overlap gradient reduction if True.

Overlap parameter gathering if True.

Preserve fp32 weights if True.

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.

Data parallel sharding strategy.

Resolved activation-checkpointing mode.

Ring-attention factor of the context-parallel axis.

Context-parallel degree (from device_mesh, default 1).

Ulysses factor of the context-parallel axis. None lets the consumer use the full context-parallel axis.

Device mesh for distributed training.

HSDP replication degree (from device_mesh, default None).

DP shard degree (from device_mesh, default 1).

Data-parallel degree (from device_mesh, default None).

Expert-parallel degree (from moe_mesh, default 1).

MoE-specific device mesh.

Optional MoE-specific parallelization policy.

True when pp_size > 1.

Pipeline-parallel degree (from device_mesh, default 1).

Optional model-local group for recipe-level collectives that must not involve ranks outside this mesh.

Optional callback run after model surgery and before distributed wrapping.

Resolved DDP, FSDP2, or Megatron-FSDP policy.

Tensor-parallel degree (from device_mesh, default 1).

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.

Policies that apply specifically to resolved vision modules.

Data type for pipeline computation. If None, uses the model’s default dtype.

Number of transformer layers per pipeline stage. If None, layers are split evenly across stages.

Loss function used for pipeline training. Required when pipeline is enabled. The function signature should be compatible with the model’s output format.

Explicit specification of which module FQNs belong to each model part/stage. If provided, overrides automatic layer splitting.

Apply pipeline patches to the CausalLM wrapper model. Defaults to True.

Apply pipeline patches to the inner model (e.g., the base transformer in a CausalLM wrapper). Defaults to True.

Patch stage backward to use no_sync context for gradient accumulation efficiency. Useful when combining PP with FSDP.

Total batch size per pipeline stage. Must be divisible by pp_microbatch_size.

Size of each microbatch for pipeline execution. pp_batch_size must be divisible by pp_microbatch_size.

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.

Extra buffer sets beyond the 1F1B in-flight depth when pp_recv_buffer_pool is enabled. Defaults to 2.

Pipeline schedule type. Supported values: “1f1b” (one-forward-one-backward), “gpipe”, “interleaved_1f1b”, “looped_bfs”, “dfs”, “v_schedule”, “zero_bubble”. Defaults to “1f1b”.

Path to a CSV file defining a custom pipeline schedule. If provided, overrides pp_schedule.

Sequence length hint for pipeline parallel shape inference. If None, it will be inferred from the dataset config.

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 gradients within the pipeline schedule (by 1/n_microbatches). If False, gradients must be scaled externally. Defaults to False.

The backend to use for torch.distributed (e.g., ‘nccl’).

Timeout (in minutes) for distributed initialization. Defaults to 1.