# 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 GPUs, and with AMD GPUs through ROCm (see [AMD GPUs](#amd-gpus-rocm))
- **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).

### AMD GPUs (ROCm)

ROCm builds of PyTorch expose the HIP runtime through the `torch.cuda` namespace, so
`--policy.device=cuda` selects an AMD GPU and no LeRobot-side change is needed. `HIP_VISIBLE_DEVICES` is preferred over `CUDA_VISIBLE_DEVICES` when selecting specific cards.

Installing LeRobot on top of an existing ROCm PyTorch needs care: LeRobot pins
`torch<2.12`, `torchvision<0.27` and `numpy<2.3`, all below what current ROCm images
ship, and `[tool.uv.sources]` resolves torch from the CUDA `cu128` index, so an
unconstrained install replaces the ROCm build and leaves the GPUs unusable. Pin those
packages with `uv pip install --override` so the rest of the tree resolves normally.
[`docker/Dockerfile.rocm`](https://github.com/huggingface/lerobot/blob/main/docker/Dockerfile.rocm)
does this and is the quickest way to a working environment; see
[docker/README.md](https://github.com/huggingface/lerobot/blob/main/docker/README.md)
for the run flags ROCm requires.

### 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), `bf16` and `fp16`:

```bash
lerobot-train \
    --policy.type=act \
    --accelerator.mixed_precision=bf16 ...
```

`fp16` requires an accelerator: it is rejected on `--policy.device=cpu`, where accelerate builds no loss scaler and the run would silently fall back to full precision. `bf16` autocasts on CPU too, so it is not restricted.

`bf16` requires hardware support (Ampere or newer on NVIDIA GPUs). On older accelerators such as V100 or T4, use `fp16`.

### fp16 and the loss scaler

`fp16` has a much smaller exponent range than `bf16`, so gradients are scaled up before the backward pass and scaled back down before the optimizer step. LeRobot delegates this to accelerate's `GradScaler`, which also detects overflowing gradients and skips those optimizer updates. Under sharded training the gradients are `DTensor`s and PyTorch reduces the overflow flag across the device mesh, so an overflow on one rank skips the update on all of them.

Skipped updates are normal, especially in the first few steps while the scale calibrates down from its initial value. They are visible in the logs as the `scale` metric dropping. The scale is saved with every checkpoint (`training_state/scaler_state.json`) and restored on resume, so a resumed run does not have to recalibrate.

The scaler can be tuned if the defaults do not suit a policy:

```bash
lerobot-train \
    --policy.type=act \
    --accelerator.mixed_precision=fp16 \
    --accelerator.grad_scaler.init_scale=1024 \
    --accelerator.grad_scaler.growth_interval=2000 ...
```