PyTorch 가속기

LeRobot은 학습과 추론 모두에 대해 여러 하드웨어 가속 옵션을 지원합니다.

이 옵션에는 다음이 포함됩니다:

• CPU: CPU가 모든 계산을 수행하며, 전용 가속기는 사용되지 않습니다
• CUDA: NVIDIA GPU를 통한 가속, 그리고 ROCm을 통한 AMD GPU 가속(자세한 내용은 AMD GPU 참고)
• MPS: Apple Silicon GPU를 활용한 가속
• XPU: Intel 통합 및 독립형 GPU를 활용한 가속

시작하기

특정 가속기를 사용하려면 적합한 버전의 PyTorch가 설치되어 있어야 합니다.

CPU, CUDA, MPS 백엔드의 경우 PyTorch 설치 페이지에 제공된 지침을 따르세요. XPU 백엔드의 경우 PyTorch 문서의 지침을 따르세요.

AMD GPU(ROCm)

ROCm 버전의 PyTorch는 HIP 런타임을 torch.cuda 네임스페이스를 통해 노출하므로, --policy.device=cuda를 사용하면 AMD GPU가 선택되며 LeRobot 측의 변경은 필요하지 않습니다. 특정 카드를 선택할 때는 CUDA_VISIBLE_DEVICES보다 HIP_VISIBLE_DEVICES를 사용하는 것이 권장됩니다.

기존 ROCm PyTorch 위에 LeRobot을 설치할 때는 주의가 필요합니다. LeRobot의 torch, torchvision, numpy 버전 범위가 ROCm 이미지에 포함된 버전과 다를 수 있으며, 일반적인 설치 방식은 PyPI에서 torch를 가져오는데 PyPI의 Linux 휠은 CUDA 빌드이므로 제약 조건 없이 설치하면 ROCm 빌드가 대체되어 GPU를 사용할 수 없게 됩니다. 나머지 의존성 트리가 정상적으로 해석되도록 해당 패키지를 uv pip install --override로 고정하세요. docker/Dockerfile.rocm가 이를 수행하며, 작동하는 환경을 가장 빠르게 구축하는 방법입니다. ROCm이 요구하는 실행 플래그는 docker/README.md를 참고하세요.

설치 확인

설치 후, accelerate의 가용성을 확인하려면 다음을 실행합니다.

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

학습 또는 평가 실행 방법

원하는 가속기를 선택하려면 lerobot-train 또는 lerobot-eval 실행 시 --policy.device 플래그를 사용하세요. 예를 들어, Apple Silicon에서 MPS를 사용하려면 다음을 실행하세요:

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는 loss scaler를 생성하지 않고 실행이 조용히 전체 정밀도로 되돌아갑니다. bf16는 CPU에서도 자동 캐스팅되므로 제한되지 않습니다.

bf16는 하드웨어 지원이 필요합니다(NVIDIA GPU의 경우 Ampere 이상). V100 또는 T4와 같은 이전 세대 가속기에서는 fp16을 사용하세요.

fp16 및 손실 스케일러

fp16은 bf16보다 지수 범위가 훨씬 좁기 때문에, backward pass 전에 gradient를 확대하고 optimizer step 전에 다시 축소합니다. LeRobot은 이 작업을 accelerate의 GradScaler에 위임하며, GradScaler는 overflow된 gradient를 감지하고 해당 optimizer 업데이트를 건너뜁니다. sharded training 환경에서는 gradient가 DTensor이며, PyTorch가 device mesh 전체에 걸쳐 overflow 플래기를 reduce합니다. 따라서 한 rank에서 overflow가 발생하면 모든 rank의 업데이트가 건너뜁니다.

업데이트가 건너뛰어지는 것은 정상적인 현상이며, 특히 초기 단계에서 scale이 초기 값에서 하향 보정되는 동안에는 흔히 발생합니다. 로그에서는 scale 지표가 감소하는 형태로 확인할 수 있습니다. 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

Ampere 및 그 이후의 NVIDIA GPU에서 LeRobot은 기본적으로 float32 matmul과 cuDNN 컨볼루션에 TF32를 사용합니다(--tf32=true). --tf32=false를 사용하면 float32의 전체 정밀도를 유지하며, --tf32=null은 PyTorch의 기본 설정을 그대로 유지합니다.