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.
55 lines
1.6 KiB
Plaintext
55 lines
1.6 KiB
Plaintext
# PyTorch accelerators
|
|
|
|
LeRobot supports multiple hardware acceleration options for both training and inference.
|
|
|
|
These options include:
|
|
|
|
- **CPU**: CPU executes all computations, no dedicated accelerator is used
|
|
- **CUDA**: acceleration with NVIDIA & AMD GPUs
|
|
- **MPS**: acceleration with Apple Silicon GPUs
|
|
- **XPU**: acceleration with Intel integrated and discrete GPUs
|
|
|
|
## Getting Started
|
|
|
|
To use particular accelerator, a suitable version of PyTorch should be installed.
|
|
|
|
For CPU, CUDA, and MPS backends follow instructions provided on [PyTorch installation page](https://pytorch.org/get-started/locally).
|
|
For XPU backend, follow instructions from [PyTorch documentation](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html).
|
|
|
|
### Verifying the installation
|
|
|
|
After installation, accelerator availability can be verified by running
|
|
|
|
```python
|
|
import torch
|
|
print(torch.<backend_name>.is_available()) # <backend_name> is cuda, mps, or xpu
|
|
```
|
|
|
|
## How to run training or evaluation
|
|
|
|
To select the desired accelerator, use the `--policy.device` flag when running `lerobot-train` or `lerobot-eval`. For example, to use MPS on Apple Silicon, run:
|
|
|
|
```bash
|
|
lerobot-train
|
|
--policy.device=mps ...
|
|
```
|
|
|
|
```bash
|
|
lerobot-eval \
|
|
--policy.device=mps ...
|
|
```
|
|
|
|
However, in most cases, presence of an accelerator is detected automatically and `policy.device` parameter can be omitted from CLI commands.
|
|
|
|
## Mixed precision
|
|
|
|
Training precision is owned by `--accelerator.mixed_precision`, which accepts `no` (default) and `bf16`:
|
|
|
|
```bash
|
|
lerobot-train \
|
|
--policy.type=act \
|
|
--accelerator.mixed_precision=bf16 ...
|
|
```
|
|
|
|
`bf16` requires an accelerator that supports it.
|