lerobot/pr_3613/en/multi_gpu_training.md

Multi-GPU Training

LeRobot trains on multiple GPUs through Hugging Face Accelerate. Three data-parallel layouts are supported:

Layout	What it does	Config
DDP	Replicates the full model on every GPU	default on any multi-GPU launch
FSDP	Shards parameters, gradients, and optimizer state across GPUs	--parallelism.dp_shard=N
HSDP	Shards within groups of GPUs, replicates across groups	--parallelism.dp_replicate=R --parallelism.dp_shard=S

Installation

accelerate is included in the training extra:

pip install 'lerobot[training]'

Launching

Distributed training can be launched through both torchrun and accelerate launch. Accelerate is used as a plain launcher: it does not manage the training configuration, and every distributed training setting lives in LeRobot's own config system.

With torchrun:

torchrun --nproc-per-node=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --wandb.enable=true

With accelerate launch (as a plain launcher):

accelerate launch --num_processes=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --wandb.enable=true

With no --parallelism.* flags, a multi-process launch runs plain DDP. Multi-node runs use the standard torchrun --nnodes/--node-rank/--rdzv-endpoint flags (or accelerate launch --num_machines/--machine_rank/--main_process_ip).

Accelerate's YAML config files (accelerate launch --config_file some.yaml, accelerate config) are not supported. They configure the engine through environment variables, bypassing LeRobot's configuration system, so train_config.json would no longer describe the settings a run actually used. lerobot-train therefore refuses to start when accelerate environment variables are set. Put the settings in --parallelism.* / --accelerator.* flags instead, or set LEROBOT_ALLOW_ACCELERATE_ENV=1 to acknowledge the override and proceed anyway.

Batch semantics, learning rate, and steps

Each of the dp_replicate × dp_shard data-parallel workers loads its own --batch_size micro-batch every step, so one training step consumes batch_size × dp_world_size samples, and × gradient_accumulation_steps of those go into each optimizer update:

effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps

The training banner prints this factorization at startup. --steps counts loop steps (micro-batches per worker), not optimizer updates.

Gradient accumulation is a first-class flag:

torchrun --nproc-per-node=2 $(which lerobot-train) \
  --batch_size=8 --accelerator.gradient_accumulation.steps=4 ...

LeRobot does not auto-scale the learning rate or the number of steps when the effective batch size grows. If you scale out and want equivalent training, please adjust manually, e.g. with 2 GPUs: double --optimizer.lr (linear scaling), or halve --steps.

Sharded training (FSDP)

If a model is too large to train with DDP, shard it with FSDP2:

torchrun --nproc-per-node=4 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=<your_policy> \
  --parallelism.dp_shard=4 \
  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp

--parallelism.dp_shard=-1 shards over however many processes the launcher started.

Wrap units

FSDP shards the model in units (typically the repeated transformer block) and gathers one unit at a time during forward/backward. Policies declare their wrap units via _fsdp_wrap_modules on the policy class. For example, ACT declares ["ACTEncoderLayer", "ACTDecoderLayer"] and FastWAM declares ["MoTLayer"]. For a policy without a _fsdp_wrap_modules declaration, pass one of the flags below. You can specify the module class name explicitly, or use a size-based policy instead:

--accelerator.fsdp.wrap_modules='["MyTransformerBlock"]'   # explicit class names
--accelerator.fsdp.min_num_params=1000000                  # or: wrap every submodule above 1M params

If a policy doesn't declare _fsdp_wrap_modules and no --accelerator.fsdp.wrap_modules or --accelerator.fsdp.min_num_params is passed, the run fails at startup rather than silently wrapping only the root module (which would forfeit all sharding memory savings).

Other sharding settings:

• --accelerator.fsdp.reshard_after_forward: whether to keep each unit's parameters resident after forward.
• --accelerator.fsdp.cpu_offload: keeps parameters, gradients and optimizer states on CPU.
• --accelerator.fsdp.ignored_modules: a regex of module paths to keep unsharded.

HSDP

Hybrid Sharded Data Parallel: parameters, gradients and optimizer states are sharded across dp_shard ranks, and that sharding is replicated dp_replicate times. Parameter all-gathers and gradient reduce-scatters stay inside a shard group; only the all-reduce that synchronizes the replicas crosses between groups. The two degrees must multiply to the world size:

# 16 GPUs = 2 nodes × 8: shard within each node, replicate across nodes
torchrun --nnodes=2 --nproc-per-node=8 ... $(which lerobot-train) \
  --parallelism.dp_replicate=2 --parallelism.dp_shard=8 ...

Checkpoints

Every checkpoint contains a pretrained_model/ directory and a training_state/ directory:

005000/  # the training step at that checkpoint
├── pretrained_model/
│   ├── config.json  # policy config
│   ├── train_config.json  # the full training config
│   ├── model.safetensors  # full weights (checkpoint_format ∈ {safetensors, safetensors_dcp}, or any non-sharded run)
│   ├── pytorch_model_fsdp_0/  # DCP weight shards (checkpoint_format ∈ {dcp, safetensors_dcp})
│   ├── policy_preprocessor.json  # preprocessor config (when the run has a preprocessor)
│   ├── policy_preprocessor_step_*.safetensors  # state of the stateful preprocessor steps
│   ├── policy_postprocessor.json  # postprocessor config (when the run has a postprocessor)
│   └── policy_postprocessor_step_*.safetensors  # state of the stateful postprocessor steps
└── training_state/
    ├── training_step.json  # step counter, topology, and batch semantics
    ├── rng_state.safetensors  # rng states
    ├── scheduler_state.json  # scheduler state (when the run has a scheduler)
    ├── optimizer_state.safetensors  # full optimizer state (non-sharded runs)
    ├── optimizer_param_groups.json  # optimizer param groups (non-sharded runs)
    └── optimizer_0/  # DCP optimizer shards (sharded runs)

During single-GPU or DDP training, the pipeline serializes each state dict into a single file: model.safetensors for the model and optimizer_state.safetensors for the optimizer.

During sharded training, the optimizer state is saved as DCP shards under training_state/optimizer_0/, and the layout of the model under pretrained_model/ can be configured through --checkpoint_format:

--checkpoint_format	Weights artifact	Use when
safetensors (default)	single model.safetensors only	you want every checkpoint immediately loadable with from_pretrained
dcp	pytorch_model_fsdp_0/ shard directory only	gathering the full weights makes saves and resumes too slow
safetensors_dcp	both	you want fast resume and immediately loadable checkpoints

Two things to know about gathered (safetensors) checkpoints from sharded runs:

• They store fp32 weights. Under mixed precision training, FSDP keeps an fp32 master copy, and the checkpoint saves the master copy to make sure training resumes consistently.
• The gather is collective (all ranks participate) but only the main process writes.

Converting DCP checkpoints

lerobot-convert-dcp merges a DCP shard directory into a regular model.safetensors, offline and without GPUs:

lerobot-convert-dcp --checkpoint_dir=outputs/train/run/checkpoints/005000
lerobot-convert-dcp --checkpoint_dir=... --delete_dcp=true --push_to_hub=${HF_USER}/my_policy

--push_to_hub publishes the converted directory as a model repo.

Resuming

Resume with --resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json. Resuming from a DCP checkpoint supports resharding the model and optimizer state to the current topology, which means you can resume with a different dp_replicate/dp_shard split. The data sampler can always resume at the right epoch and offset, but is only sample-exact when the world size and batch size match the original run (a warning is logged otherwise).

FSDP checkpoints written by LeRobot 0.6.x and earlier used a different on-disk layout (a gathered full optimizer state) and cannot be resumed.

Notes

• Checkpoint saves and end-of-training publishes are collective (every rank enters them). Gathered weights, sidecar files and Hub uploads are written by the main process alone.
• Metrics are reduced across ranks before logging: losses are averaged, and samples/s reports cluster-wide throughput.
• Learning-rate scheduling is stepped once per training step regardless of the number of processes (step_scheduler_with_optimizer=False is baked in).

For background on the underlying machinery, see the Accelerate FSDP guide. To go deeper on large-scale training, check out the Ultrascale Playbook.