Bộ tăng tốc PyTorch

LeRobot hỗ trợ nhiều tùy chọn tăng tốc phần cứng cho cả huấn luyện và suy luận.

Các tùy chọn này bao gồm:

• CPU: CPU thực hiện tất cả các phép tính, không sử dụng bộ tăng tốc chuyên dụng
• CUDA: tăng tốc bằng GPU NVIDIA và bằng GPU AMD thông qua ROCm (xem GPU AMD)
• MPS: tăng tốc với GPU Apple Silicon
• XPU: tăng tốc với GPU tích hợp và rời của Intel

Bắt đầu

Để sử dụng bộ tăng tốc cụ thể, cần cài đặt phiên bản PyTorch phù hợp.

Đối với các backend CPU, CUDA và MPS, hãy làm theo các hướng dẫn được cung cấp trên trang cài đặt PyTorch. Đối với backend XPU, hãy làm theo các hướng dẫn từ tài liệu PyTorch.

GPU AMD (ROCm)

Các bản build ROCm của PyTorch phơi bày runtime HIP thông qua namespace torch.cuda, do đó --policy.device=cuda sẽ chọn một GPU AMD và không cần thay đổi nào phía LeRobot. HIP_VISIBLE_DEVICES được ưu tiên hơn CUDA_VISIBLE_DEVICES khi chọn các card cụ thể.

Việc cài đặt LeRobot trên nền PyTorch ROCm hiện có cần thận trọng: LeRobot cố định torch<2.12, torchvision<0.27 và numpy<2.3, tất cả đều thấp hơn mức mà các image ROCm hiện tại cung cấp, và [tool.uv.sources] phân giải torch từ index CUDA cu128, do đó một lệnh cài đặt không ràng buộc sẽ thay thế bản build ROCm và khiến các GPU không thể sử dụng được. Hãy cố định các gói đó bằng uv pip install --override để phần còn lại của cây phụ thuộc được phân giải bình thường. docker/Dockerfile.rocm thực hiện điều này và là cách nhanh nhất để có một môi trường hoạt động; xem docker/README.md để biết các cờ chạy mà ROCm yêu cầu.

Xác minh cài đặt

Sau khi cài đặt, khả năng sử dụng bộ tăng tốc có thể được xác minh bằng cách chạy

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

Cách chạy huấn luyện hoặc đánh giá

Để chọn bộ tăng tốc mong muốn, hãy sử dụng cờ --policy.device khi chạy lerobot-train hoặc lerobot-eval. Ví dụ, để sử dụng MPS trên Apple Silicon, hãy chạy:

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

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

Tuy nhiên, trong hầu hết các trường hợp, sự hiện diện của bộ tăng tốc sẽ được phát hiện tự động và tham số policy.device có thể được bỏ qua trong các lệnh CLI.

Độ chính xác hỗn hợp

Độ chính xác của quá trình huấn luyện được quản lý bởi --accelerator.mixed_precision, chấp nhận no (mặc định), bf16 và fp16:

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

fp16 yêu cầu một bộ tăng tốc: nó bị từ chối khi dùng --policy.device=cpu, vì accelerate không xây dựng loss scaler và quá trình chạy sẽ âm thầm quay lại độ chính xác đầy đủ. bf16 tự động chuyển đổi (autocast) trên CPU, nên không bị hạn chế.

bf16 yêu cầu hỗ trợ phần cứng (Ampere hoặc mới hơn trên GPU NVIDIA). Trên các bộ tăng tốc cũ hơn như V100 hoặc T4, hãy sử dụng fp16.

fp16 và bộ điều chỉnh loss

fp16 có phạm vi số mũ nhỏ hơn nhiều so với bf16, do đó gradient được nhân lên trước bước backward và chia nhỏ lại trước bước optimizer. LeRobot giao việc này cho GradScaler của accelerate, thành phần này cũng phát hiện gradient tràn và bỏ qua các bước cập nhật optimizer tương ứng. Trong chế độ huấn luyện sharded, gradient là các DTensor và PyTorch giảm cờ tràn trên toàn bộ lưới thiết bị, nên nếu một rank bị tràn thì tất cả các rank đều bỏ qua bước cập nhật.

Việc bỏ qua các lần cập nhật là bình thường, đặc biệt là trong vài bước đầu tiên khi scale đang hiệu chỉnh giảm từ giá trị ban đầu. Chúng có thể được quan sát trong log qua việc chỉ số scale giảm xuống. Scale được lưu cùng với mỗi checkpoint (training_state/scaler_state.json) và được khôi phục khi tiếp tục chạy, do đó một lần chạy được tiếp tục không cần phải hiệu chỉnh lại.

Bộ điều chỉnh có thể được tinh chỉnh nếu các giá trị mặc định không phù hợp với một chính sách:

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

TF32

Trên các GPU NVIDIA Ampere và mới hơn, LeRobot mặc định sử dụng TF32 cho các phép nhân ma trận float32 và phép tích chập cuDNN (--tf32=true). Sử dụng --tf32=false để giữ chúng ở độ chính xác float32 đầy đủ, và --tf32=null sẽ giữ nguyên cài đặt mặc định của PyTorch.