SageMaker 分散式模型平行化最佳實務 - Amazon SageMaker AI

本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。

SageMaker 分散式模型平行化最佳實務

當您使用 SageMaker 模型平行化程式庫 v2 (SMP v2) 執行分散式訓練任務時,請遵循下列指南。

為分散式訓練設定正確的組態

若要預估並尋找套用 SMP v2 提供之分散式訓練技術的最佳起點,請檢閱下列清單。每個清單項目討論使用 SageMaker 模型平行化程式庫第 2 版的核心功能 的優點以及潛在的權衡。

組態提示

本節提供如何根據全域批次大小需求來決定最佳模型組態,以實現最佳輸送量的指導方針。

首先,無論您的模型大小為何,我們都建議下列設定。

  1. 使用您可以使用的最強大執行個體類型。

  2. 永遠開啟混合精確度,因為它可為改善效能和減少記憶體提供顯著的優勢。我們建議您使用 bfloat16,因為它比 float16 更加精確。

  3. 若適用,開啟 SageMaker 分散式資料平行化程式庫 (而非使用 NCCL),如與針對AWS基礎設施最佳化的 SMDDP 程式庫的相容性所示。一個例外情形是僅限張量平行化使用案例 (hybrid_shard_degree = 1tensor_paralle_degree > 1)。

  4. 如果您的模型有超過 600 億個參數,建議使用延遲參數初始化。您也可以使用延遲參數初始化來加速任何模型的初始化。

  5. 我們建議您啟用啟用檢查點

根據您的模型大小,我們建議您從下列指引開始。

  1. 使用碎片資料平行化。

    1. 根據您要放入 GPU 記憶體的批次大小,選擇適當的碎片資料平行程度。一般而言,您應該從模型適合 GPU 記憶體的最低程度開始,同時可將網路通訊的額外負荷降至最低。如果您看到正在發生快取排清的警告,建議您提高碎片程度。

    2. 根據最大本機批次大小和所需的全域批次大小 (如果有的話) 來決定 world_size

    3. 您可以實驗啟用卸載。視案例而定,它可以解決您的記憶體需求,而不必增加碎片程度,這表示通訊較少。

  2. 同時使用 PyTorch FSDP 的碎片資料平行化和 SMP v2 的張量平行化,如張量平行化中所述。

    1. 在大型叢集上訓練時,僅使用 FSDP 會導致全域批次大小變得太大,導致模型的收斂問題。一般而言,大多數研究工作會將批次大小保持在 400 萬個字符以下。在這種情況下,您可以透過將 PyTorch FSDP 與 SMP v2 的張量平行化組合在一起來減少批次大小,以解決這個問題。

      例如,如果您有 256 個節點且序列長度為 4096,即使每個 GPU 的批次大小為 1,也會導致全域批次大小為 800 萬個字符。不過,當您使用程度為 2 的張量平行化,以及每個張量平行群組的批次大小為 1 時,這會成為每個 GPU 的 1/2 批次大小,也就是 400 萬個字符。

    2. 使用 8k 等長內容進行訓練時,16k 啟用記憶體可能會變得非常高。FSDP 不會對啟用進行分片,而啟用可能會導致 GPU 耗盡記憶體。在這種情況下,您可以透過將 PyTorch FSDP 與 SMP v2 的張量平行化結合來進行高效率的訓練。

參考組態

SageMaker 模型平行化訓練團隊根據使用 torch.sagemaker.transform 轉換為 SMP 變壓器模型的 Llama 2 模型實驗提供了以下參考點,並在序列長度為 4096 和混合精確度 (FP16 或 BF16) 的 ml.p4d.24xlarge 執行個體上進行訓練。

模型 模型大小 (模型參數數量) 執行個體數量 碎片資料平行程度 張量平行程度 啟用檢查點 啟用卸載 批次大小
Llama 2 7B 1 8 1 TRUE FALSE 4
70B 32 256 1 TRUE FALSE 2
175B 64 128 4 TRUE TRUE 6

您可以從上述組態中推斷,以估算模型組態的 GPU 記憶體使用量。例如,如果您增加 100 億個參數模型的序列長度,或將模型的大小增加到 200 億,則可能需要先降低批次大小。如果模型仍然不符合,請嘗試增加張量平行處理程度。

使用 SageMaker AI 主控台和 Amazon CloudWatch 監控和記錄訓練任務

若要監控系統層級指標 (例如 CPU 記憶體使用率、GPU 記憶體使用率和 GPU 使用率),請使用透過 SageMaker AI 主控台提供的視覺效果。

  1. 在左側導覽窗格中,選擇訓練

  2. 選擇 Training jobs (訓練任務)

  3. 在主窗格中,選擇您要查看其更多詳細資訊的訓練任務名稱。

  4. 瀏覽主窗格,並找到監視器區段以查看自動化視覺效果。

  5. 若要查看訓練任務日誌,請選擇監視器區段中的檢視日誌。您可以在 CloudWatch 中存取訓練任務的分散式訓練任務日誌。如果您已啟動多節點分散式訓練,您應該會看到多個日誌串流,其標記格式為 algo-n-1234567890algo-1 日誌串流會追蹤主節點 (第 0 個) 的訓練日誌。

如需更多詳細資訊,請參閱 用於監控和分析訓練任務的 Amazon CloudWatch 指標

許可

若要執行具有模型平行化的 SageMaker 訓練任務,請確認您有 IAM 角色的正確許可,例如: