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.

Quick start

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

1
import torch
2
from nemo_automodel import NeMoAutoModelForCausalLM
3
from nemo_automodel.components.distributed import (
4
    DistributedSetup,
5
    FSDP2Config,
6
    ParallelismSizes,
7
    initialize_distributed,
8
)
9
10
dist_env = initialize_distributed("nccl")
11
12
distributed_setup = DistributedSetup.build(
13
    strategy=FSDP2Config(sequence_parallel=True),
14
    parallelism_sizes=ParallelismSizes(tp_size=2),
15
    activation_checkpointing=True,
16
    world_size=dist_env.world_size,
17
)
18
19
model = NeMoAutoModelForCausalLM.from_pretrained(
20
    "meta-llama/Llama-3.2-1B",
21
    distributed_setup=distributed_setup,
22
)

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 config object, and the pipeline / MoE configs accept either a dataclass or a plain dict.

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:

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

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

Migrating from the per-keyword API

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

Before:

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

After:

1
distributed_setup = DistributedSetup.build(
2
    strategy=FSDP2Config(),
3
    parallelism_sizes=ParallelismSizes(tp_size=2),
4
    pipeline_config=pp_cfg,
5
    moe_parallel_config=moe_cfg,
6
    activation_checkpointing=True,
7
)
8
9
model = NeMoAutoModelForCausalLM.from_pretrained(
10
    "meta-llama/Llama-3.2-1B",
11
    distributed_setup=distributed_setup,
12
)

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.