Ctrl+K

Distributed Training#

Scaling a training run across multiple GPUs or nodes in ALCHEMI comes down to one object: DistributedManager. It gathers the distributed runtime state a job needs — process rank, local rank, world size, device selection, process groups, and DistributedDataParallel defaults — behind a single handle, re-exported from PhysicsNeMo as nvalchemi.distributed.DistributedManager.

Routing that state through one object buys a shared context: passing a manager to TrainingStrategy gives every ALCHEMI hook the same view of the distributed runtime, so no hook has to read environment variables or initialize communication on its own. The same script then runs unchanged whether launched on one process or many. Advanced workflows can still drive torch.distributed directly, but the manager is the recommended entry point.

Basic pattern#

A distributed script differs from a single-process one in only two places: you bring up the runtime, and you hand the strategy a manager alongside a DDPHook. Call DistributedManager.initialize() once to start the distributed runtime, construct a manager from it, and pass that instance to the strategy. From there DDPHook does the wiring during setup: it selects the rank-local device, wraps the optimized models in torch.nn.parallel.DistributedDataParallel, and installs a distributed sampler for supported dataloaders.

from nvalchemi.distributed import DistributedManager
from nvalchemi.training import TrainingStrategy
from nvalchemi.training.hooks import DDPHook

DistributedManager.initialize()
manager = DistributedManager()

strategy = TrainingStrategy(
    ...,
    distributed_manager=manager,
    hooks=[
        DDPHook(),
    ],
)

strategy.run(train_loader)

Launch the script with the process launcher for your environment. For a simple single-node PyTorch launch:

$ torchrun --nproc_per_node=4 train.py

Note

DistributedManager.initialize() also supports single-process execution. When the world size is one, DDPHook becomes a no-op, so the same script runs unchanged locally or under a distributed launcher.

For a complete single-node dummy training script, see Distributed Training: DDPHook with a Dummy MLP. It can be launched with:

$ uv run --extra cu12 torchrun --standalone --nproc_per_node=2 \
    examples/intermediate/06_ddp_mlp_training.py --backend auto

Data loaders and samplers#

The one part of distributed training that needs care beyond the manager is data loading: each data-parallel rank must see a different slice of the training data, or the ranks would redundantly train on the same samples. How that sharding is arranged depends on the sampler, and the subsections below walk the cases in increasing order of control — the automatic default, a custom distributed sampler, and multi-dataset batch sampling.

Automatic configuration via DDPHook#

For regular nvalchemi data pipes, DDPHook installs the distributed sampler during strategy setup with no extra configuration; what it installs depends on the dataset type. For a single Dataset, the hook wraps a DistributedSampler in a BatchSampler so the loader keeps emitting complete batches. For MultiDataset, it installs MultiDatasetBatchSampler so the per-dataset batch composition and rank sharding are handled together. Either way, the hook infers num_replicas, rank, shuffle, and drop_last from the distributed manager and dataloader, and uses seed=0 unless overridden.

from nvalchemi.data.datapipes import DataLoader, Dataset
from nvalchemi.distributed import DistributedManager
from nvalchemi.training import TrainingStrategy
from nvalchemi.training.hooks import DDPHook

DistributedManager.initialize()
manager = DistributedManager()

dataset = Dataset(reader, device=manager.device)
train_loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    pin_memory=True,
)

strategy = TrainingStrategy(
    ...,
    distributed_manager=manager,
    hooks=[DDPHook()],
)
strategy.run(train_loader)

This is the preferred starting point for a single dataset, and like the basic pattern it stays single-process friendly: when manager.world_size == 1, DDPHook leaves the loader unchanged.

Use sampler_kwargs to override arguments passed to the default sampler:

DDPHook(
    sampler_kwargs={
        "shuffle": False,
        "seed": 1234,
    },
)

Custom distributed sampler#

When the default sampler does not fit, you can supply your own, and DDPHook gets out of the way accordingly. If a dataloader already has a distributed-aware sampler, the hook preserves it instead of replacing it. A sampler counts as distributed-aware when it satisfies DistributedSamplerProtocol: it exposes num_replicas, rank, and set_epoch(epoch). Native PyTorch DistributedSampler satisfies this protocol.

For a sampler class or factory that accepts PyTorch-style distributed sampler arguments, pass it to DDPHook. The hook supplies num_replicas, rank, shuffle, seed, and drop_last defaults before applying your sampler_kwargs.

DDPHook(
    sampler_cls=MyDistributedSampler,
    sampler_kwargs={
        "seed": 1234,
    },
)

If your sampler uses different constructor names, pass those names explicitly in sampler_kwargs.

DDPHook(
    sampler_cls=MyDistributedSampler,
    sampler_kwargs={
        "replicas": manager.world_size,
        "worker_rank": manager.rank,
    },
)

Multidataset batch sampling#

Training on several datasets at once raises a second question on top of sharding: how batches are composed across the child datasets. When DDPHook sees a nvalchemi DataLoader backed by MultiDataset and no custom sampler class was supplied, it installs MultiDatasetBatchSampler automatically, keeping per-dataset batch composition and distributed sharding in the same sampler.

Pass MultiDatasetBatchSampler options through DDPHook.sampler_kwargs when you need a specific allocation policy, such as balanced batches or a fixed number of samples per child dataset. The hook still supplies num_replicas, rank, shuffle, and drop_last defaults before applying your overrides. As with the single-dataset case, DDPHook preserves an existing distributed-aware batch sampler instead of replacing it, so manual construction remains available for fully custom samplers.

from nvalchemi.data.datapipes import (
    AtomicDataZarrReader,
    DataLoader,
    Dataset,
    MultiDataset,
)
from nvalchemi.distributed import DistributedManager
from nvalchemi.training import TrainingStrategy
from nvalchemi.training.hooks import DDPHook

DistributedManager.initialize()
manager = DistributedManager()

dataset = MultiDataset(
    Dataset(AtomicDataZarrReader("dataset_a.zarr"), device=manager.device),
    Dataset(AtomicDataZarrReader("dataset_b.zarr"), device=manager.device),
)

train_loader = DataLoader(
    dataset,
    batch_size=64,
    prefetch_factor=16,
    pin_memory=True,
)

strategy = TrainingStrategy(
    ...,
    distributed_manager=manager,
    hooks=[
        DDPHook(
            sampler_kwargs={
                "epoch_policy": "max_size",
                "replacement": True,
                "seed": 1234,
            },
        ),
    ],
)
strategy.run(train_loader)

Internally, MultiDatasetBatchSampler first builds the global batch order according to its per-dataset allocation policy, then splits that order across data-parallel ranks. With drop_last=False, it pads the batch order so each rank emits the same number of batches, matching PyTorch DistributedSampler behavior; with drop_last=True, it truncates the uneven tail instead.

Whichever sampler is in play, call set_epoch() yourself, or let TrainingStrategy call it during training, so distributed samplers reshuffle deterministically from epoch to epoch.

API details#

This guide covers the training-facing surface of the manager. For the complete API, including process-group methods and the distributed configuration knobs, see the PhysicsNeMo DistributedManager API.

On this page
Show Source