# Multi-GPU Training

LeRobot trains on multiple GPUs through [Hugging Face Accelerate](https://huggingface.co/docs/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:

```bash
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`:

```bash
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 \
  --trackio.enable=true
```

With `accelerate launch` (as a plain launcher):

```bash
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 \
  --trackio.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`).

> [!WARNING]
> 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](https://huggingface.co/docs/accelerate/usage_guides/fsdp) 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:

```bash
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:

```bash
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.

`--accelerator.mixed_precision` support both `bf16` and `fp16`. Details in [mixed precision](./torch_accelerators#mixed-precision)

### 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:

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

FSDP can't split tied weights (e.g. a language head tied to the input embedding) across wrap units. torch 2.12 and newer stop at the first training step and name the parameter, while older versions silently untie them. Wrapping the repeated block class usually keeps them together in the root unit.

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.
- `--accelerator.fsdp.mixed_precision_policy.{param_dtype,reduce_dtype,output_dtype,cast_forward_inputs}`: the fields of PyTorch's `MixedPrecisionPolicy`, applied to every FSDP unit. Once any field is set, the policy is built from these fields alone and replaces the one derived from `--accelerator.mixed_precision`, so set `param_dtype` too. Before torch 2.15, an unset `reduce_dtype` follows `param_dtype`. From torch 2.15, it follows the parameters' own dtype (float32 under mixed precision). `float16` needs `--accelerator.mixed_precision=fp16`, which builds the loss scaler.

### 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:

```bash
# 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 ...
```

### PEFT

[PEFT fine-tuning](./peft_training) (`--peft.*`) works under FSDP2 and HSDP. The frozen base model is sharded together with the adapters, so each rank only holds its shard of the base weights:

```bash
torchrun --nproc-per-node=4 $(which lerobot-train) \
  --policy.path=${HF_USER}/my_base_policy \
  --peft.method_type=LORA --peft.r=16 \
  --parallelism.dp_shard=4 \
  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp_lora
```

- Wrap units still come from the policy's `_fsdp_wrap_modules` declaration (or `--accelerator.fsdp.wrap_modules`).
- Use mixed precision when the base model is stored in bf16 or fp16. PEFT keeps its adapters in fp32 while the copies of `--peft.full_training_modules` keep the base model's dtype, and FSDP2 needs a single dtype for the trainable parameters it shards together. With mixed precision on, accelerate upcasts every trainable parameter to fp32.
- Checkpoints store the adapter instead of the full weights, in every `--checkpoint_format` (see [Checkpoints](#checkpoints)).
- `--accelerator.fsdp.ignored_modules` is matched against the module names of the PEFT model, which start with `base_model.model.`. Begin the pattern with `.*` to match them.

## Checkpoints

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

```text
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/
    ├── optimizer_param_groups.json  # optimizer param groups (non-sharded runs)
    ├── optimizer_state.safetensors  # optimizer state (non-sharded runs)
    ├── optimizer_0/  # DCP optimizer shards (sharded runs)
    ├── rng_state.safetensors  # rng states
    ├── scheduler_state.json  # scheduler state (when the run has a scheduler)
    ├── scaler_state.json  # fp16 loss-scaler state (mixed_precision=fp16 only)
    └── training_step.json  # step + dp_world_size/batch_size/grad_accum/precision + topology
```

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           |

A PEFT run stores its adapter instead of the full weights: `adapter_model.safetensors` and `adapter_config.json` replace `model.safetensors`, and the DCP shards hold the adapter only. The base model loads from the path recorded in `adapter_config.json`.

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:

```bash
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. The shards of a PEFT run merge into `adapter_model.safetensors`. Like any `dcp` checkpoint, a PEFT checkpoint saved with `--checkpoint_format=dcp` loads with `--policy.path` only after this conversion.

### 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).

> [!NOTE]
> 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](https://huggingface.co/docs/accelerate/usage_guides/fsdp). To go deeper on large-scale training, check out the [Ultrascale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook).