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. For XPU backend, follow instructions from PyTorch documentation.

Verifying the installation

After installation, accelerator availability can be verified by running

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:

lerobot-train
    --policy.device=mps ...

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:

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

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