Distributed Setup (Python API)

When you use the YAML recipes or the automodel CLI, distributed training is configured from the distributed: section of your config (see Configuration). When you call the NeMoAutoModel* loaders directly from Python, you describe the same topology and execution policies with a single typed object: DistributedSetup.

Quickstart

Build a DistributedSetup and pass it to from_pretrained (or from_config):

import torch
from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.components.distributed import (
    DistributedSetup,
    FSDP2Config,
    ParallelismSizes,
    initialize_distributed,
)
dist_env = initialize_distributed("nccl")
distributed_setup = DistributedSetup.build(
    strategy=FSDP2Config(sequence_parallel=True),
    parallelism_sizes=ParallelismSizes(tp_size=2),
    activation_checkpointing=True,
    world_size=dist_env.world_size,
)
model = NeMoAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    distributed_setup=distributed_setup,
)

The same distributed_setup= keyword works on every NeMo AutoModel loader, including NeMoAutoModelForImageTextToText, NeMoAutoModelForSequenceClassification, and NeMoAutoModelForTokenClassification.

DistributedSetup.build

DistributedSetup.build resolves a device mesh and the execution policies from your requested parallelism sizes. It is intentionally forgiving about input types: strategy accepts a string or a configuration object, and the pipeline or MoE configurations accept either a dataclass or a plain dictionary.

YAML distributed: sections also forward strategy-specific fields to FSDP2 or DDP.

ParallelismSizes is durable user intent (what you requested). The resolved runtime topology lives on DistributedSetup.mesh_context (MeshContext), which derives its sizes from the live DeviceMesh after build.

Plain Device Mesh Shortcut

If you only need a topology and no NeMo-specific policies, you can pass a pre-created Hugging Face-style DeviceMesh directly as device_mesh=. NeMo wraps it in a topology-only DistributedSetup internally:

from torch.distributed.device_mesh import init_device_mesh
mesh = init_device_mesh("cuda", mesh_shape=(2,), mesh_dim_names=("tp",))
model = NeMoAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    device_mesh=mesh,
)

Pass either distributed_setup or device_mesh, not both. Use distributed_setup whenever you need strategy, pipeline, MoE, or activation-checkpointing policies.

Migrate From the Per-Keyword API

Earlier releases accepted a flat set of distributed keywords on from_pretrained or from_config. These are now consolidated into the single distributed_setup object.

Before:

model = NeMoAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    distributed_config=FSDP2Config(activation_checkpointing=True),
    tp_size=2,
    pipeline_config=pp_cfg,
    moe_config=moe_cfg,
    moe_mesh=moe_mesh,
    activation_checkpointing=True,
)

After:

distributed_setup = DistributedSetup.build(
    strategy=FSDP2Config(),
    parallelism_sizes=ParallelismSizes(tp_size=2),
    pipeline_config=pp_cfg,
    moe_parallel_config=moe_cfg,
    activation_checkpointing=True,
)
model = NeMoAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    distributed_setup=distributed_setup,
)

See Also

• Configuration: The equivalent distributed: YAML section used by recipes and the CLI.
• Pipeline Parallelism: The lower-level AutoPipeline building blocks for custom training loops.
• Gradient Checkpointing: Full and selective activation checkpointing.
• DistributedSetup source code.