Skip to content

Multi-GPU Training & Precision

echelon3 runs multi-GPU training with built-in DistributedDataParallel (DDP) — one process per GPU. You do not need torchrun: name the GPUs and echelon3 spawns the workers itself.

Multi-GPU in one command

echelon3-train --config-dir configs --config-name my_experiment gpus=[0,1,2,3]

gpus is a root config key. Set it in the config, or override it on the CLI. Leave it out and echelon3 uses every visible GPU on the node:

gpus: [0, 1, 2, 3]   # optional; default = all GPUs on this node

Under the hood echelon3 calls PyTorch's own launcher (elastic_launch) to start one worker per GPU, wiring up RANK / LOCAL_RANK / WORLD_SIZE / MASTER_* and the process group — exactly what torchrun does, without you typing it. With a single GPU (or on CPU) nothing is spawned and training runs in-process.

No DataParallel

DataParallel has been removed (0.5.0). Multiple GPUs always run as DDP, one process each; a single process only ever drives one GPU. The old device_ids key no longer selects multiple GPUs — use gpus.

torchrun / multi-node still works

The environment-variable path is unchanged, so torchrun (and SLURM srun) remain available for multi-node or elastic jobs. If echelon3 finds RANK in the environment it assumes it is already a worker and does not spawn again:

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 \
    $(which echelon3-train) --config-dir configs --config-name my_experiment

For a single node, gpus=[...] and torchrun are equivalent — prefer gpus for the shorter command; reach for torchrun when you span nodes.

Batch size is global

dataloaders.train.config.batch_size is the global batch size — the total across all processes. Under DDP, create_dataloaders divides it by the world size and installs a DistributedSampler:

dataloaders:
  train:
    module: torch.utils.data
    type: DataLoader
    config:
      batch_size: 40        # global; 40 / 4 = 10 samples per GPU
      shuffle: true         # handed to DistributedSampler; loader shuffle disabled
      num_workers: 6
      drop_last: true

It must divide evenly

If batch_size is not divisible by the number of GPUs, the run fails fast with a ValueError. Pick a global batch that is a multiple of len(gpus).

The test loader is treated the same way: its batch_size is global and divided per rank.

What each rank does

The network is wrapped in DistributedDataParallel (with find_unused_parameters=True by default; set trainer.config.ddp_find_unused_parameters: false when your graph uses every output every step). Each rank trains on its shard and DDP all-reduces gradients.

Validation is symmetric and sharded. Every rank evaluates its own DistributedSampler shard through the unwrapped network, and the metrics aggregate their state across ranks inside compute(). The aggregated values are identical on every rank, so the keep-best decision matches everywhere; only rank 0 writes the file.

Logs and checkpoints: rank 0 only

  • Rank 0 prints, shows progress bars, and owns the mlops logger; other ranks redirect stdout to /dev/null, disable tqdm, and get a no-op logger.
  • save_checkpoint returns immediately on non-main ranks — only rank 0 writes.

Checkpoints store the unwrapped state_dict (no module. prefix), so a file is identical whether it came from a single-GPU or a multi-GPU run and resumes under either. Older checkpoints that still carry a module. prefix load fine — it is stripped automatically.

Mixed precision (AMP) & TF32

Training uses bf16 automatic mixed precision by default on GPUs that support it (Hopper / Ampere and newer); on CPU or unsupported GPUs it stays fp32. On modern hardware this is a large speedup at negligible quality cost.

Control it in trainer.config:

trainer:
  config:
    precision: bf16        # auto (default) | bf16 | fp16 | fp32
    tf32: true             # TF32 matmul on Ampere+ (default true)
    cudnn_benchmark: true  # autotune conv algorithms for fixed input sizes (default true)
  • auto / unset → bf16 when supported, else fp32.
  • fp32 → disables autocast (bit-for-bit the pre-0.5.0 behavior). This is how you keep training in plain fp32.
  • fp16 → autocast + GradScaler. Not supported with closure optimizers (SAMOptimizer / LBFGS, which do a double backward) — those fall back to bf16 automatically.

The default flipped to bf16 in 0.5.0

Results differ from old fp32 runs (usually negligibly, and faster). Set precision: fp32 to reproduce fp32 exactly.

echelon3-evaluate and echelon3-run autocast the same way (default bf16); set precision: fp32 at the config root to force fp32 for those.

torch.compile (experimental)

bf16 only speeds up compute-bound work. A small network on a big GPU is often launch-bound instead — dominated by per-kernel launch overhead, with the GPU idle between many tiny kernels — and there bf16 buys nothing. The lever is torch.compile, which fuses kernels and cuts the launch count:

trainer:
  config:
    compile: true             # off by default
    compile_mode: null        # null | "reduce-overhead" | "max-autotune"

The network is compiled before the DDP wrapper; ddp.unwrap() and checkpoints strip the resulting _orig_mod. prefix, so checkpoints stay interchangeable with uncompiled runs. The first few steps recompile (warmup) — measure steady-state, not iteration 1.

Experimental

Verified single-GPU and on 4×H200 DDP (trains + checkpoints round-trip). The actual speedup is model-dependent — confirm on your model (watch for shape-driven recompiles, and closure optimizers like SAM). Whether a workload is launch-bound is worth checking first (nvidia-smi dmon -s pu: power well under TDP ⇒ the GPU is starved, and compile is the right lever).

Next