張量平行處理

張量平行處理是模型平行處理類型，其中特定模型權重、漸層與最佳化工具狀態會跨裝置分割。與管線平行處理原則相反，它可以保持個別權重不變，但會跨裝置分割權重、漸層或最佳化程式集，張量平行度會分割個別權重。這通常涉及特定作業、模組或模型層的分散式運算。

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

SMP v2 與變壓器引擎集成以實現張量並行性，並在 FSDP API 之 PyTorch 上運行。您可以同時啟用 PyTorch FSDP 和 SMP 張量平行度，並確定最佳模型平行度以獲得最佳效能。

在實踐中，張量平行性在以下情況下特別有用。

• 當使用較長的上下文長度進行訓練時，僅使用 FSDP 就會導致高激活內存。

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

與 SMP 張量平行度相容的 Hugging Face 變壓器型號

SMP v2 目前為下列 Hugging Face 變壓器型號提供張量平行度支援。

• GPT-NeoX

• 美洲駝 2 號

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

設定張量平行處理

對於tensor_parallel_degree，您可以選取張量平行程度的值。此值必須平均除以叢集中的 GPU 數目。例如，若要在使用具有 8 個 GPU 的執行個體時分片模型，請選擇 2、4 或 8。我們建議您從少數字開始，並逐漸增加它，直到模型適合 GPU 內存。

下列程式碼片段說明如何在遵循中介紹的兩個步驟程序時，將 SMP 初始化模組新增torch.sagemaker.init()至訓練指令碼，並以 JSON 格式設定 SMP 組態字典，以供訓練工作啟動器使用。開始使用 SageMaker 模型平行程式庫 v2您不需要對 PyTorch 模型或 PyTorch FSDP 組態進行任何變更。如需 tensor_parallel_degree 和 random_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 變壓器檢查點功能不兼容。為了處理這個問題，SMP 程式庫提供 API，以便在 Hugging Face 部變壓器表示中儲存轉換模型的檢查點，以及用來載入 Hugging Face 變壓器模型檢查點以進行微調的 torch.sagemaker.transform API。

如需使用 SMP v2 的張量平行處理原則功能時儲存檢查點的詳細資訊，請參閱。使用 SMP 時儲存並載入檢查點

如需有關套用 SMP v2 張量平行處理原則功能微調模型的詳細資訊，請參閱。微調