Aceleradores de PyTorch

LeRobot admite múltiples opciones de aceleración de hardware para el entrenamiento y la inferencia.

Estas opciones incluyen:

• CPU: la CPU ejecuta todos los cálculos, no se utiliza un acelerador dedicado
• CUDA: aceleración con GPUs de NVIDIA y con GPUs de AMD a través de ROCm (ver GPUs de AMD)
• MPS: aceleración con GPUs de Apple Silicon
• XPU: aceleración con GPUs integradas y discretas de Intel

Primeros pasos

Para usar un acelerador en particular, debe estar instalada una versión adecuada de PyTorch.

Para los backends CPU, CUDA y MPS, siga las instrucciones proporcionadas en la página de instalación de PyTorch. Para el backend XPU, siga las instrucciones de la documentación de PyTorch.

GPUs de AMD (ROCm)

Las compilaciones de PyTorch para ROCm exponen el tiempo de ejecución HIP a través del espacio de nombres torch.cuda, por lo que --policy.device=cuda selecciona una GPU de AMD y no se necesita ningún cambio en LeRobot. Se prefiere HIP_VISIBLE_DEVICES sobre CUDA_VISIBLE_DEVICES al seleccionar tarjetas específicas.

Instalar LeRobot sobre una instalación existente de PyTorch ROCm requiere precaución: los límites de torch, torchvision y numpy de LeRobot pueden diferir de los que incluye una imagen ROCm, y una instalación normal toma torch desde PyPI, cuyos wheels para Linux son compilaciones CUDA, por lo que una instalación sin restricciones reemplaza la compilación ROCm y deja las GPUs inutilizables. Fija esos paquetes con uv pip install --override para que el resto del árbol se resuelva de forma normal. docker/Dockerfile.rocm hace esto y es la forma más rápida de obtener un entorno funcional; consulta docker/README.md para las banderas de ejecución que requiere ROCm.

Verificando la instalación

Después de la instalación, la disponibilidad del acelerador se puede verificar ejecutando

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

Cómo ejecutar el entrenamiento o la evaluación

Para seleccionar el acelerador deseado, use la bandera --policy.device al ejecutar lerobot-train o lerobot-eval. Por ejemplo, para usar MPS en Apple Silicon, ejecute:

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

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

Sin embargo, en la mayoría de los casos, la presencia de un acelerador se detecta automáticamente y el parámetro policy.device puede omitirse de los comandos de CLI.

Precisión mixta

La precisión del entrenamiento la controla --accelerator.mixed_precision, que acepta no (predeterminado), bf16 y fp16:

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

fp16 requiere un acelerador: se rechaza con --policy.device=cpu, donde accelerate no construye un escalador de pérdida y la ejecución volvería silenciosamente a precisión completa. bf16 también realiza autocasting en CPU, por lo que no está restringido.

bf16 requiere soporte de hardware (Ampere o más reciente en GPUs de NVIDIA). En aceleradores más antiguos, como V100 o T4, use fp16.

fp16 y el escalador de pérdida

fp16 tiene un rango de exponente mucho menor que bf16, por lo que los gradientes se escalan hacia arriba antes de la pasada inversa y se escalan de nuevo hacia abajo antes del paso del optimizador. LeRobot delega esto en GradScaler de accelerate, que también detecta gradientes con desbordamiento y omite esas actualizaciones del optimizador. En el entrenamiento fragmentado, los gradientes son DTensor y PyTorch reduce la bandera de desbordamiento a través de la malla de dispositivos, de modo que un desbordamiento en un rango omite la actualización en todos ellos.

Las actualizaciones omitidas son normales, especialmente en los primeros pasos mientras la escala se calibra hacia abajo desde su valor inicial. Son visibles en los registros como la métrica scale que disminuye. La escala se guarda con cada punto de control (training_state/scaler_state.json) y se restaura al reanudar, por lo que una ejecución reanudada no tiene que recalibrarse.

El escalador se puede ajustar si los valores predeterminados no se adaptan a una política:

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

TF32

En GPUs NVIDIA Ampere y más recientes, LeRobot hace que las multiplicaciones de matrices float32 y las convoluciones de cuDNN usen TF32 de forma predeterminada (--tf32=true). Use --tf32=false para mantenerlas en precisión completa float32, y --tf32=null deja las configuraciones predeterminadas de PyTorch.