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.

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. 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: bfloat16

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.

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.

Controls frame-level encoder compute sharding over the selected text-model mesh dimensions.