mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
Haoming Song
ef88d4e52b
* feat(train): parallel training engine with FSDP2, HSDP, and DCP checkpoints Replace the FSDP1 training path with a config-owned parallel-training engine: - Topology and runtime configs (--parallelism.*, --accelerator.*): dp_replicate x dp_shard degrees select single-process, DDP (unchanged default), FSDP2, or HSDP; mixed precision, first-class gradient accumulation, and FSDP/DDP tuning knobs are mirrored as plain dataclasses that build the accelerate objects at runtime, so every run is reproducible from its train_config.json alone. Accelerate env vars are guarded against configuring the engine behind the config system's back. - Declarative policy surface: policies declare FSDP2 wrap units (_fsdp_wrap_modules) and non-forward entry points (_fsdp_forward_methods); a shared engine resolves them around accelerator.prepare(). Context-parallel fields are reserved and validated to 1. - Checkpoints: selectable --checkpoint_format (safetensors | dcp | safetensors_dcp); the sharded optimizer channel is always DCP; two-phase resume (step+RNG before prepare, DCP model/optimizer after) reshards across GPU-topology changes; lerobot-convert-dcp merges DCP shards into a distributable model.safetensors offline. - Publishing: PreTrainedPolicy.push_model_to_hub is replaced by the free publish_trained_model (model + processors + card + train config, all-ranks gather with main-rank writes); PreTrainedPolicy._save_pretrained gathers state dicts internally, removing the state_dict= threading from save_pretrained. - lerobot_train is restructured around the engine: optimizer built before the single prepare() call, deferred weight load on DCP resumes, collective save_checkpoint with no call-site rank branches, dp-world-size-based sample accounting. Breaking changes: FSDP checkpoints from lerobot <= 0.6.x are not resumable (weights stay loadable via from_pretrained; pin lerobot==0.6.x to finish old runs); the `accelerate launch --config_file` yaml flow is superseded by the config flags; training autocast is owned exclusively by --accelerator.mixed_precision (policy.dtype only casts parameters). Also fixes: reward-model hub publishing crash (TypeError on extra kwargs). Verified by ~200 new CPU tests (config round-trips, checkpoint round-trips per format, two-phase resume, publisher contracts, converter equivalence, accelerate canaries), a 5-test 4-GPU suite (FSDP2 save/resume bit-exactness, HSDP/DDP loss parity, changed-topology resume, all-ranks save_pretrained, grad-accum equivalence), and end-to-end ACT (1/4/8 GPUs) + FastWAM 6B (FSDP2 + HSDP) training runs.
177 lines
10 KiB
Plaintext
177 lines
10 KiB
Plaintext
# 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 \
|
||
--wandb.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 \
|
||
--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`).
|
||
|
||
> [!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.
|
||
|
||
### 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
|
||
```
|
||
|
||
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:
|
||
|
||
```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 ...
|
||
```
|
||
|
||
## 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/
|
||
├── 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:
|
||
|
||
```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.
|
||
|
||
### 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).
|