TopolearnModels that do not fit, or runs that take days on one GPU, need sharded or parallel training.
Trainer runs distributed with no code changes; you change how you launch it.
- DDP (DistributedDataParallel):
torchrun --nproc_per_node=4 train.py. Each GPU holds a full model copy, processes a different slice of each batch, and gradients are all-reduced. Simple and fast, but the model plus optimizer state must fit on one GPU. - FSDP (Fully Sharded Data Parallel): shards parameters, gradients and optimizer state across GPUs and gathers each layer just in time. Enable it with
TrainingArguments(fsdp="full_shard auto_wrap", fsdp_config={...})or anaccelerate configfile. Memory per GPU drops roughly by the number of GPUs. - DeepSpeed ZeRO:
TrainingArguments(deepspeed="ds_config.json"). Stage 1 shards optimizer state, stage 2 also gradients, stage 3 also parameters; it can offload to CPU or NVMe.
Things that change in distributed runs:
per_device_train_batch_sizeis per GPU; the global batch multiplies by the world size.- Only rank 0 should write files or print; Trainer handles this, your own code should check
trainer.is_world_process_zero(). - Saving a sharded model needs a full state dict gathered on rank 0 (FSDP's state dict settings handle this).
- Data loading must be sharded; Trainer's sampler does it.
accelerate launch reads a saved config and works for both single-node and multi-node jobs. Start with DDP, move to FSDP or ZeRO-3 only when memory requires it.
Example
# train.py stays the same Trainer script; the launch and args change.
from transformers import TrainingArguments
fsdp_args = TrainingArguments(
output_dir="out-fsdp",
per_device_train_batch_size=2, # global batch = 2 x world size x accumulation
gradient_accumulation_steps=4,
bf16=True,
fsdp="full_shard auto_wrap",
fsdp_config={"transformer_layer_cls_to_wrap": ["Qwen2DecoderLayer"]},
report_to="none",
)
# DeepSpeed alternative: TrainingArguments(..., deepspeed="ds_zero3.json")
# Launch commands (shell):
# torchrun --nproc_per_node=4 train.py # DDP or FSDP on 4 GPUs
# accelerate config && accelerate launch train.py # saved launch config
# torchrun --nnodes=2 --nproc_per_node=8 --rdzv_backend=c10d \
# --rdzv_endpoint=host0:29500 train.py # 2 nodes x 8 GPUsRemember
- Trainer code is unchanged; launch with torchrun or accelerate
- DDP replicates the model; FSDP and ZeRO-3 shard parameters, gradients and optimizer state
- Batch size arguments are per device; the global batch scales with world size
Watch out for
- Writing files or logs from every rank, which produces corrupted or duplicated outputs.