Ускорители PyTorch

LeRobot поддерживает несколько вариантов аппаратного ускорения как для обучения, так и для инференса.

К ним относятся:

• CPU: процессор выполняет все вычисления, специализированный ускоритель не используется
• CUDA: ускорение с помощью GPU NVIDIA и GPU AMD через ROCm (см. GPU AMD)
• MPS: ускорение с помощью графических процессоров Apple Silicon
• XPU: ускорение с помощью интегрированных и дискретных GPU от Intel

Начало работы

Для использования конкретного ускорителя должна быть установлена соответствующая версия PyTorch.

Для бэкендов CPU, CUDA и MPS следуйте инструкциям на странице установки PyTorch. Для бэкенда XPU следуйте инструкциям из документации PyTorch.

Видеокарты AMD (ROCm)

Сборки PyTorch для ROCm предоставляют доступ к рантайму HIP через пространство имён torch.cuda, поэтому --policy.device=cuda выбирает GPU AMD, и никаких изменений со стороны LeRobot не требуется. Для выбора конкретных карт предпочтительнее использовать HIP_VISIBLE_DEVICES, а не CUDA_VISIBLE_DEVICES.

Установка LeRobot поверх существующего ROCm PyTorch требует осторожности: LeRobot фиксирует версии torch<2.12, torchvision<0.27 и numpy<2.3, все они ниже тех, что поставляются в актуальных ROCm-образах, а [tool.uv.sources] разрешает torch из CUDA-индекса cu128, поэтому установка без ограничений заменяет ROCm-сборку и делает GPU непригодными для работы. Зафиксируйте эти пакеты с помощью uv pip install --override, чтобы остальное дерево зависимостей разрешалось штатно. docker/Dockerfile.rocm выполняет это и является самым быстрым способом получить рабочую среду; флаги запуска, необходимые для ROCm, описаны в docker/README.md.

Проверка установки

После установки наличие ускорителей можно проверить, запустив

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

Как запустить обучение или оценку

Чтобы выбрать нужный ускоритель, используйте флаг --policy.device при запуске lerobot-train или lerobot-eval. Например, чтобы использовать MPS на Apple Silicon, выполните:

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

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

Однако в большинстве случаев наличие ускорителя определяется автоматически, и параметр policy.device можно опустить в командах CLI.

Смешанная точность

Точность обучения определяется параметром --accelerator.mixed_precision, который принимает значения no (по умолчанию), bf16 и fp16:

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

fp16 требует ускорителя: при --policy.device=cpu он отклоняется, поскольку accelerate не создаёт масштабатор потерь, и выполнение будет незаметно переключаться на полную точность. bf16 автоматически преобразует типы и на CPU, поэтому на него это ограничение не распространяется.

bf16 требует аппаратной поддержки (Ampere или новее на GPU NVIDIA). На более старых ускорителях, таких как V100 или T4, используйте fp16.

fp16 и масштабировщик потерь

fp16 имеет значительно меньший диапазон экспоненты, чем bf16, поэтому градиенты масштабируются перед обратным проходом и возвращаются к исходному масштабу перед шагом оптимизатора. LeRobot делегирует эту задачу GradScaler из accelerate, который также обнаруживает переполнение градиентов и пропускает соответствующие обновления оптимизатора. При обучении с шардированием градиенты представляют собой DTensor, и PyTorch сводит флаг переполнения по сетке устройств, так что переполнение на одном ранке приводит к пропуску обновления на всех ранках.

Пропуск обновлений — это нормальное явление, особенно на первых шагах, пока масштаб снижается от начального значения. В логах это видно по падению метрики scale. Масштаб сохраняется вместе с каждой контрольной точкой (training_state/scaler_state.json) и восстанавливается при возобновлении, поэтому при перезапуске не требуется повторная калибровка.

Масштабировщик можно настроить, если значения по умолчанию не подходят для политики:

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

TF32

На GPU NVIDIA Ampere и новее LeRobot по умолчанию использует TF32 для матричных умножений float32 и свёрток cuDNN (--tf32=true). Используйте --tf32=false, чтобы сохранить полную точность float32, а --tf32=null оставляет настройки PyTorch по умолчанию.