多 GPU 訓練

本指南將說明如何使用 Hugging Face Accelerate 在多個 GPU 上訓練策略。

安裝

首先，請確保您已安裝 accelerate

pip install accelerate

使用多個 GPU 進行訓練

您可以透過兩種方式啟動訓練

選項 1：不使用配置檔（直接指定參數）

您無需執行 accelerate config，即可在指令中直接指定所有參數：

accelerate launch \
  --multi_gpu \
  --num_processes=2 \
  $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --wandb.enable=true

Accelerate 關鍵參數

• --multi_gpu：啟用多 GPU 訓練
• --num_processes=2：要使用的 GPU 數量
• --mixed_precision=fp16：使用 fp16 混合精度（若支援，亦可使用 bf16）

選項 2：使用 accelerate config

如果您偏好儲存配置，可以選擇執行以下指令，為您的硬體環境配置 accelerate：

accelerate config

此互動式設定將詢問有關您訓練環境的問題（GPU 數量、混合精度設定等），並將配置儲存以備日後使用。對於單機多 GPU 的簡單設定，您可以使用以下推薦設定：

• 運算環境：本機 (This machine)
• 機器數量：1
• 程序數量：(您想要使用的 GPU 數量)
• 要使用的 GPU ID：(留空以使用所有 GPU)
• 混合精度：fp16 或 bf16 (推薦用於更快的訓練)

然後透過以下指令啟動訓練：

accelerate launch $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --wandb.enable=true

運作原理

當您使用 accelerate 啟動訓練時：

1. 自動偵測：LeRobot 會自動偵測是否正在 accelerate 下執行。
2. 資料分佈：您的 Batch 會自動拆分到各個 GPU 上。
3. 梯度同步：在反向傳播期間，梯度會在各個 GPU 間進行同步。
4. 單一程序日誌：僅由主程序將日誌記錄至 wandb 並儲存檢查點 (checkpoint)。

學習率與訓練步數的縮放

重要：LeRobot 不會根據 GPU 數量自動縮放學習率或訓練步數。這讓您能完全掌控您的訓練超參數。

為什麼不自動縮放？

許多分散式訓練框架會自動按 GPU 數量縮放學習率（例如 lr = base_lr × num_gpus）。然而，LeRobot 始終保持您指定的學習率不變。

何時以及如何縮放

如果您在使用多個 GPU 時希望縮放超參數，則應手動執行。

學習率縮放

# Example: 2 GPUs with linear LR scaling
# Base LR: 1e-4, with 2 GPUs -> 2e-4
accelerate launch --num_processes=2 $(which lerobot-train) \
  --optimizer.lr=2e-4 \
  --dataset.repo_id=lerobot/pusht \
  --policy=act

訓練步數縮放

由於有效 Batch Size bs 會隨 GPU 數量增加（batch_size × num_gpus），您可能需要按比例減少訓練步數。

# Example: 2 GPUs with effective batch size 2x larger
# Original: batch_size=8, steps=100000
# With 2 GPUs: batch_size=8 (16 in total), steps=50000
accelerate launch --num_processes=2 $(which lerobot-train) \
  --batch_size=8 \
  --steps=50000 \
  --dataset.repo_id=lerobot/pusht \
  --policy=act

注意事項

• lerobot-train 中的 --policy.use_amp 旗標僅在沒有使用 accelerate 時才會生效。當使用 accelerate 時，混合精度由 accelerate 的配置控制。
• 訓練日誌、檢查點和上傳至 Hub 的作業僅由主程序執行，以避免衝突。非主程序會停用主控台輸出，以防止重複顯示日誌。
• 有效 Batch Size 為 batch_size × num_gpus。如果您使用 4 個 GPU 且設定 --batch_size=8，您的有效 Batch Size 即為 32。
• 學習率排程在多個程序中被正確處理——LeRobot 設定 step_scheduler_with_optimizer=False，以防止 accelerate 根據程序數量調整排程步數。
• 在儲存或推送模型時，LeRobot 會自動將模型從 accelerate 的分散式封裝中解包，以確保相容性。
• WandB 整合功能會自動僅在主程序中初始化，從而防止建立多個執行紀錄。

如需進階配置與疑難排解，請參閱 Accelerate 文件。若想深入了解如何在大量 GPU 上進行訓練，請查看這份精彩指南：Ultrascale Playbook。