張量平行處理 - Amazon SageMaker AI

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

張量平行處理

張量平行處理是模型平行處理類型,其中特定模型權重、漸層與最佳化工具狀態會跨裝置分割。與保持個別權重完好無損但將一組權重、漸層或最佳化工具分割到不同裝置相比,張量平行處理會碎片化個別權重。這通常涉及特定作業、模組或模型層的分散式運算。

如果單一參數使用多數 GPU 記憶體 (例如字彙量較大的大型內嵌資料表或具大量類別的大型 softmax 層),則需要張量平行處理。在這種情況,將此大型張量或作業視為原子單位不具效率,且會阻礙記憶體負載的平衡。

SMP v2 與 Transformer Engine 整合,以實作張量平行處理,並在 PyTorch FSDP APIs 上執行。您可以同時啟用 PyTorch FSDP 和 SMP 張量平行處理,並判斷最佳模型平行處理,以獲得最佳效能。

實際上,張量平行處理在下列案例中特別有用。

  • 當使用長內容長度進行訓練時,這會導致單獨使用 FSDP 的高啟用記憶體。

  • 使用超大型叢集進行訓練時,全域批次大小超過所需的限制。

與 SMP 張量平行處理相容的 Hugging Face Transformer 模型

SMP v2 目前為下列 Hugging Face 轉換器模型提供張量平行處理支援。

如需在這些模型上套用張量平行處理的參考組態,請參閱 組態提示

設定張量平行處理

對於 tensor_parallel_degree,您可以選取張量平行處理程度的值。值必須平均分配叢集中的 GPUs 數量。例如,若要在使用具有 8 個 GPUs 的執行個體時碎片化模型,請選擇 2、4 或 8。我們建議您從小數字開始,並逐漸增加,直到模型符合 GPU 記憶體。

下列程式碼片段示範如何將 SMP 初始化模組新增至torch.sagemaker.init()訓練指令碼,並設定訓練任務啟動器的 JSON 格式 SMP 組態字典,同時遵循 中介紹的兩個步驟程序使用 SageMaker 模型平行處理程式庫 v2。您不需要對 PyTorch 模型或 PyTorch FSDP 組態進行任何變更。如需 tensor_parallel_degreerandom_seed 參數的詳細資訊,請參閱 SMP v2 核心功能組態參數

SMP 組態

{ "tensor_parallel_degree": 8, "random_seed": 0 }

在您的訓練指令碼中

使用 初始化torch.sagemaker.init()以啟用 SMP v2,並使用 torch.sagemaker.transform API 包裝模型。

import torch.sagemaker as tsm tsm.init() from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_config(..) model = tsm.transform(model)

儲存和載入 Hugging Face 轉換器檢查點

SMP 程式庫轉換模型後,會變更模型的狀態字典 (state_dict)。這表示模型與原始 Hugging Face Transformer 檢查點功能不相容。為了處理此問題,SMP 程式庫提供 APIs,以從 Hugging Face Transformer 表示法中的轉換模型儲存檢查點,以及 torch.sagemaker.transform API 載入 Hugging Face Transformer 模型檢查點進行微調。

如需在使用 SMP v2 的張量平行處理功能時儲存檢查點的詳細資訊,請參閱 使用 SMP 進行檢查點

如需使用 SMP v2 張量平行處理功能微調模型的詳細資訊,請參閱 微調