Akselerator PyTorch

LeRobot mendukung berbagai opsi akselerasi perangkat keras untuk pelatihan dan inferensi.

Opsi-opsi ini meliputi:

• CPU: CPU menjalankan semua komputasi, tidak ada akselerator khusus yang digunakan
• CUDA: akselerasi dengan GPU NVIDIA, dan dengan GPU AMD melalui ROCm (lihat GPU AMD)
• MPS: akselerasi dengan GPU Apple Silicon
• XPU: akselerasi dengan GPU terintegrasi dan diskrit Intel

Memulai

Untuk menggunakan akselerator tertentu, versi PyTorch yang sesuai harus diinstal.

Untuk backend CPU, CUDA, dan MPS, ikuti instruksi yang disediakan di halaman instalasi PyTorch. Untuk backend XPU, ikuti instruksi dari dokumentasi PyTorch.

GPU AMD (ROCm)

Build ROCm dari PyTorch mengekspos runtime HIP melalui namespace torch.cuda, sehingga --policy.device=cuda memilih GPU AMD dan tidak diperlukan perubahan di sisi LeRobot. HIP_VISIBLE_DEVICES lebih disarankan daripada CUDA_VISIBLE_DEVICES saat memilih kartu tertentu.

Menginstal LeRobot di atas PyTorch ROCm yang sudah ada memerlukan perhatian khusus: batas torch, torchvision, dan numpy LeRobot mungkin berbeda dari yang disertakan dalam image ROCm, dan instalasi normal mengambil torch dari PyPI, di mana wheel Linux-nya adalah build CUDA, sehingga instalasi tanpa batasan akan menggantikan build ROCm dan membuat GPU tidak dapat digunakan. Kunci paket-paket tersebut dengan uv pip install --override agar sisa pohon dependensi teresolusi secara normal. docker/Dockerfile.rocm melakukan hal ini dan merupakan cara tercepat untuk mendapatkan lingkungan yang berfungsi; lihat docker/README.md untuk flag eksekusi yang diperlukan ROCm.

Memverifikasi instalasi

Setelah instalasi, ketersediaan akselerator dapat diverifikasi dengan menjalankan

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

Cara menjalankan pelatihan atau evaluasi

Untuk memilih akselerator yang diinginkan, gunakan flag --policy.device saat menjalankan lerobot-train atau lerobot-eval. Misalnya, untuk menggunakan MPS pada Apple Silicon, jalankan:

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

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

Namun, dalam kebanyakan kasus, keberadaan akselerator terdeteksi secara otomatis dan parameter policy.device dapat diabaikan dari perintah CLI.

Presisi campuran

Presisi pelatihan dimiliki oleh --accelerator.mixed_precision, yang menerima no (default), bf16, dan fp16:

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

fp16 memerlukan akselerator: opsi ini ditolak pada --policy.device=cpu, di mana accelerate tidak membangun loss scaler dan eksekusi akan secara diam-diam kembali ke presisi penuh. bf16 juga melakukan autocast pada CPU, sehingga tidak dibatasi.

bf16 memerlukan dukungan perangkat keras (Ampere atau lebih baru pada GPU NVIDIA). Pada akselerator yang lebih lama seperti V100 atau T4, gunakan fp16.

fp16 dan loss scaler

fp16 memiliki rentang eksponen yang jauh lebih kecil daripada bf16, sehingga gradien diskalakan naik sebelum langkah mundur dan diskalakan turun kembali sebelum langkah optimizer. LeRobot mendelegasikan hal ini ke GradScaler dari accelerate, yang juga mendeteksi gradien yang meluap dan melewati pembaruan optimizer tersebut. Dalam pelatihan yang dipartisi, gradien berupa DTensor dan PyTorch mereduksi flag overflow di seluruh mesh perangkat, sehingga overflow pada satu rank akan melewati pembaruan di semua rank.

Pembaruan yang dilewati adalah hal yang normal, terutama pada beberapa langkah pertama saat skala dikalibrasi turun dari nilai awalnya. Hal ini terlihat di log sebagai penurunan metrik scale. Skala disimpan bersama setiap checkpoint (training_state/scaler_state.json) dan dipulihkan saat resume, sehingga run yang diresume tidak perlu melakukan kalibrasi ulang.

Scaler dapat disetel jika nilai default tidak sesuai dengan kebijakan:

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

TF32

Pada GPU NVIDIA Ampere dan yang lebih baru, LeRobot membuat matmul float32 dan konvolusi cuDNN menggunakan TF32 secara default (--tf32=true). Gunakan --tf32=false untuk mempertahankan presisi float32 penuh, dan --tf32=null mempertahankan pengaturan default PyTorch.