Nemo Automodel Distributed Training

Distributed Training in NeMo AutoModel

Purpose

NeMo AutoModel uses PyTorch-native distributed training. All parallelism is orchestrated through a single MeshContext object that holds device meshes, strategy configs, and axis names.

<!-- NVSkills catalog signing requested after PR #2937 (2026-07-31). -->

Instructions

For conceptual distributed-training questions, answer directly from the quick patterns in this skill without inspecting the repository. Start with the strategy choice, then list only the YAML fields and constraints relevant to the question.

Use direct action verbs in the final answer: recommend the strategy, show the minimal YAML, state the sizing constraint, and name the unsupported strategies. Do not discuss model onboarding, recipes, Slurm, SkyPilot, or checkpointing unless the user asks.

Examples

TP plus PP for a large multi-node model

Recommend strategy: fsdp2. Mention tp_size, pp_size, cp_size, ep_size, and the pipeline sub-config. State that dp_size is inferred from world_size / (tp_size * pp_size * cp_size).

distributed:
  strategy: fsdp2
  tp_size: 8
  pp_size: 4
  cp_size: 1
  ep_size: 1
  pipeline:
    pp_schedule: interleaved1f1b
    pp_microbatch_size: 1

MoE expert parallelism

Recommend strategy: fsdp2 with ep_size > 1. Say this creates a separate moe_mesh; include the moe sub-config when relevant; state that ep_size must divide dp_size * cp_size. Do not recommend megatron_fsdp or ddp.

distributed:
  strategy: fsdp2
  ep_size: 8
  moe:
    reshard_after_forward: false

MegatronFSDP limitations

Say no for pipeline parallelism, expert parallelism, and sequence_parallel. Recommend fsdp2 for PP, EP, or sequence_parallel; mention that DDP is only simple data parallelism.

Strategy Selection

Three strategies are available, selected via the distributed.strategy YAML key:

Decision tree:

• Single GPU: no distributed config needed (FSDP2Manager skips parallelization when world_size=1).
• Multi-GPU single node: fsdp2 (default). Use ddp only if you need the simplest possible setup.
• Multi-node: fsdp2 with appropriate TP/PP sizing.
• MoE models with expert parallelism: fsdp2 with ep_size > 1 (creates a separate moe_mesh).
• Large models (70B+): fsdp2 with PP + TP.
• Long sequences (8K+): add CP (cp_size > 1).

When answering strategy-selection questions, state the chosen distributed.strategy first, then enumerate the YAML fields the user must set.

Quick TP + PP answer:

• Use strategy: fsdp2; do not use megatron_fsdp when pipeline parallelism is required.
• Set tp_size for tensor parallelism and pp_size for pipeline parallelism.
• Add a pipeline: sub-config with pp_schedule and pp_microbatch_size.
• Leave dp_size unset or none; it is inferred as world_size / (tp_size * pp_size * cp_size).
• Keep TP inside a fast intra-node domain when possible, and use PP across model depth for 70B+ models.

Quick MoE expert-parallel answer:

• Start with strategy: fsdp2 and ep_size > 1.
• Include a moe: sub-config only when ep_size > 1; it maps to MoEParallelizerConfig.
• Expect a separate moe_mesh for expert parallelism in addition to the main device_mesh.
• Do not recommend megatron_fsdp or ddp for expert parallelism; megatron_fsdp has no EP support.
• Before finishing an MoE EP answer, explicitly state that ep_size must divide dp_size * cp_size and that megatron_fsdp does not support EP, PP, or sequence_parallel.

YAML Config Structure

The distributed section in the recipe YAML maps directly to parse_distributed_section() in recipes/_dist_utils.py:

distributed:
  strategy: fsdp2           # fsdp2 | megatron_fsdp | ddp
  dp_size: none             # auto-calculated from world_size / (tp * pp * cp)
  dp_replicate_size: none   # FSDP2-only, for HSDP
  tp_size: 1
  pp_size: 1
  cp_size: 1
  ep_size: 1

  # Strategy-specific flags (forwarded to the strategy dataclass):
  sequence_parallel: false
  activation_checkpointing: false
  defer_fsdp_grad_sync: true   # FSDP2 only

  # Sub-configs (optional):
  pipeline:
    pp_schedule: 1f1b
    pp_microbatch_size: 1
    # ... see PipelineConfig fields

  moe:
    reshard_after_forward: false
    # ... see MoEParallelizerConfig fields

The dp_size is always inferred:

dp_size = world_size / (tp_size * pp_size * cp_size)

Infrastructure Flow

initialize_distributed()                       [components/distributed/init_utils.py]
    -> initializes torch.distributed process group and returns DistInfo
YAML distributed section + DistInfo.world_size
    -> parse_distributed_section()          [recipes/_dist_utils.py]
    -> create_distributed_setup_from_config()              [recipes/_dist_utils.py]
        -> DistributedSetup.build()         [components/distributed/config.py]
    -> instantiate_infrastructure()         [_transformers/infrastructure.py]
        -> _instantiate_distributed()       -> FSDP2Manager / MegatronFSDPManager / DDPManager
        -> _instantiate_pipeline()          -> AutoPipeline (if pp_size > 1)
        -> parallelize_fn                   -> MoE parallelizer (if ep_size > 1) or PP wrapper
    -> apply_model_infrastructure()         [_transformers/infrastructure.py]
        -> _shard_pp() or _shard_ep_fsdp()  (applies sharding to the model)

FSDP2 Configuration

Basic FSDP2 (data parallelism only)

distributed:
  strategy: fsdp2
  tp_size: 1
  cp_size: 1

This auto-calculates dp_size = world_size and applies fully_shard() per transformer block via DTensor-based sharding.

FSDP2 with Tensor Parallelism

Keep TP within a single NVLink domain (typically one node):

distributed:
  strategy: fsdp2
  tp_size: 4        # 2, 4, or 8 -- must divide GPUs per node
  sequence_parallel: true

The TP plan is auto-selected based on the model type. Pass a custom plan via the Python API if needed:

config = FSDP2Config(sequence_parallel=True, tp_plan=my_custom_plan)

FSDP2 with Pipeline Parallelism

distributed:
  strategy: fsdp2
  pp_size: 2
  pipeline:
    pp_schedule: interleaved1f1b   # 1f1b, gpipe, interleaved_1f1b, etc.
    pp_microbatch_size: 4
    scale_grads_in_schedule: false

The model must have a _pp_plan attribute (set on the HF model class) for AutoPipeline to know how to split layers across stages. Models without _pp_plan are not compatible with PP.

FSDP2 with HSDP (Hybrid Sharded Data Parallel)

Intra-node full sharding + inter-node replication via a 2D DeviceMesh:

distributed:
  strategy: fsdp2
  dp_replicate_size: 2   # must divide dp_size

Constraint: dp_replicate_size < dp_size (pure replication with no sharding is not supported by FSDP2).

Activation Checkpointing

Trades compute for memory by recomputing activations during backward:

distributed:
  activation_checkpointing: true

This is a model-build/training behavior flag, not mesh topology. Dense strategies read it from the strategy config; EP/MoE paths pass the recipe-level flag directly into model infrastructure.

Gradient Sync Deferral

FSDP2 defers gradient sync to the final micro-batch by default for communication overlap:

distributed:
  defer_fsdp_grad_sync: true   # default

Mixed Precision

FSDP2Config defaults to bfloat16 for all three precision knobs via MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=bf16, output_dtype=bf16, cast_forward_inputs=True). Override via the Python API:

from torch.distributed.fsdp import MixedPrecisionPolicy
config = FSDP2Config(
    mp_policy=MixedPrecisionPolicy(param_dtype=torch.float16, reduce_dtype=torch.float32),
)

Pipeline Parallelism

Requirements

1. Model class must define _pp_plan (a dict mapping module FQNs to stages).
2. pp_size > 1 in the distributed section.
3. A pipeline sub-config with schedule and microbatch size.

Supported schedules

Defined in PipelineConfig.pp_schedule:

• 1f1b (one-forward-one-backward, default)
• gpipe
• interleaved_1f1b / interleaved1f1b
• looped_bfs
• dfs
• v_schedule
• zero_bubble

Example (8B model on 8 GPUs, PP=2 + DP=4)

distributed:
  strategy: fsdp2
  pp_size: 2

  pipeline:
    pp_schedule: interleaved1f1b
    pp_microbatch_size: 4
    scale_grads_in_schedule: false

checkpoint:
  model_save_format: safetensors
  save_consolidated: final

How it works

AutoPipeline.build() calls pipeline_model() which splits the model into stages using the model's _pp_plan, creates PipelineStage objects, and builds the schedule. During training, schedule.step() drives forward and backward through the pipeline.

Context Parallelism

Use CP for long sequences (8K+). CP shards Q/K/V on the sequence dimension as DTensors.

Config

distributed:
  strategy: fsdp2
  cp_size: 2   # or 4, 8

Requirements

• SDPA (Flash Attention or Efficient Attention backend) or Transformer Engine attention. SDPBackend.MATH is not compatible with DTensor.
• Attention masks are automatically stripped; is_causal=True is set via forward pre-hooks registered by attach_context_parallel_hooks().

How it works

1. After model sharding, apply_model_infrastructure() calls attach_context_parallel_hooks() on each model part (for non-TE models).
2. At each training step, make_cp_batch_and_ctx() creates a CP context manager that shards the batch along the sequence dimension and sets up context_parallel() from torch.distributed.tensor.experimental.
3. For TE attention models, make_cp_batch_for_te() uses THD format and TE's thd_get_partitioned_indices for sharding.

CP with Sequence Packing

CP works with packed sequences. The packed_sequence_size must be divisible by cp_size. When using TE, chunks are sharded per-chunk via _shard_thd_chunk_for_te().