Aceleradores do PyTorch

O LeRobot suporta várias opções de aceleração de hardware para treinamento e inferência.

Essas opções incluem:

• CPU: a CPU executa todos os cálculos, sem o uso de um acelerador dedicado
• CUDA: aceleração com GPUs NVIDIA e com GPUs AMD por meio do ROCm (veja GPUs AMD)
• MPS: aceleração com GPUs Apple Silicon
• XPU: aceleração com GPUs integradas e discretas da Intel

Começando

Para usar um acelerador específico, uma versão adequada do PyTorch deve estar instalada.

Para os backends CPU, CUDA e MPS, siga as instruções fornecidas na página de instalação do PyTorch. Para o backend XPU, siga as instruções da documentação do PyTorch.

GPUs AMD (ROCm)

As compilações ROCm do PyTorch expõem o runtime HIP por meio do namespace torch.cuda, portanto --policy.device=cuda seleciona uma GPU AMD e nenhuma alteração no LeRobot é necessária. HIP_VISIBLE_DEVICES é preferível a CUDA_VISIBLE_DEVICES ao selecionar placas específicas.

Instalar o LeRobot sobre um PyTorch ROCm existente exige cuidado: os limites de torch, torchvision e numpy do LeRobot podem diferir do que uma imagem ROCm fornece, e uma instalação normal busca o torch do PyPI, cujos wheels para Linux são builds CUDA, de modo que uma instalação sem restrições substitui o build ROCm e deixa as GPUs inutilizáveis. Fixe esses pacotes com uv pip install --override para que o restante da árvore seja resolvido normalmente. docker/Dockerfile.rocm faz isso e é a maneira mais rápida de obter um ambiente funcional; consulte docker/README.md para as flags de execução que o ROCm requer.

Verificando a instalação

Após a instalação, a disponibilidade do acelerador pode ser verificada executando

import torch
print(torch.<backend_name>.is_available())  # <backend_name> is cuda, mps, or xpu

Como executar o treinamento ou a avaliação

Para selecionar o acelerador desejado, use a flag --policy.device ao executar lerobot-train ou lerobot-eval. Por exemplo, para usar MPS no Apple Silicon, execute:

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

lerobot-eval \
    --policy.device=mps ...

No entanto, na maioria dos casos, a presença de um acelerador é detectada automaticamente e o parâmetro policy.device pode ser omitido dos comandos de CLI.

Precisão mista

A precisão do treinamento é controlada por --accelerator.mixed_precision, que aceita no (padrão), bf16 e fp16:

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

fp16 requer um acelerador: é rejeitado em --policy.device=cpu, onde o accelerate não constrói um loss scaler e a execução cairia silenciosamente para precisão completa. bf16 também faz autocast em CPU, portanto não é restrito.

bf16 requer suporte de hardware (Ampere ou mais recente em GPUs NVIDIA). Em aceleradores mais antigos, como V100 ou T4, use fp16.

fp16 e o escalonador de perda

fp16 tem uma faixa de expoente muito menor que a de bf16, portanto os gradientes são ampliados antes da passagem reversa e reduzidos novamente antes do passo do otimizador. O LeRobot delega isso ao GradScaler do accelerate, que também detecta gradientes com estouro e pula essas atualizações do otimizador. No treinamento particionado, os gradientes são DTensors e o PyTorch reduz a flag de estouro em toda a malha de dispositivos, de modo que um estouro em um rank faz com que a atualização seja pulada em todos eles.

Atualizações ignoradas são normais, especialmente nos primeiros passos, enquanto a escala se calibra a partir do valor inicial. Elas são visíveis nos logs como a métrica scale diminuindo. A escala é salva com cada checkpoint (training_state/scaler_state.json) e restaurada ao retomar, de modo que uma execução retomada não precisa se recalibrar.

O scaler pode ser ajustado se os padrões não forem adequados a uma política:

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

TF32

Em GPUs NVIDIA Ampere e mais recentes, o LeRobot faz com que as multiplicações de matrizes float32 e as convoluções cuDNN usem TF32 por padrão (--tf32=true). Use --tf32=false para mantê-las em precisão float32 completa, e --tf32=null preserva as configurações padrão do PyTorch.