內容平行處理 - Amazon SageMaker AI

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

內容平行處理

內容平行處理是一種模型平行處理,可沿著序列維度分割模型啟用。與僅分割 LayerNorm和 的其他序列平行處理技術不同RMSNorm,內容平行處理會沿著序列維度分割網路輸入和所有中繼啟用。

SMP v2 與 Transformer Engine 整合,可實現內容平行處理,並可與 FSDP和 SMP 搭配使用 PyTorch張量平行處理。您可以同時啟用所有三個平行處理以進行模型訓練。內容平行處理對於具有大型啟用大小和長序列長度的訓練模型很有幫助。它可加速注意力分數和注意力輸出的運算,允許每個裝置僅沿著序列維度計算一部分分數和輸出。雖然張量平行處理也透過沿著隱藏維度的分割加速運算,但內容平行處理的優勢更重要,因為運算需求隨序列維度增加四分之一。

Hugging Face Transformer 模型與SMP內容平行處理相容

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

設定內容平行處理

將整數值設定為 參數,該context_parallel_degree參數會平均分配叢集GPUs中的 。例如,如果您有 8 個GPU執行個體,請針對 使用 2、4 或 8context_parallel_degree。我們建議從小context_parallel_degree值開始逐步增加,直到模型符合所需輸入序列長度的GPU記憶體。

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

在您的訓練指令碼中

步驟 1 中,使用 初始化指令碼torch.sagemaker.init()以啟用 SMP v2,並使用 torch.sagemaker.transform 包裝模型API。

從 SMP v2.6.0 開始,您可以使用 引數cp_comm_type來決定要使用的內容平行處理實作。SMP 程式庫目前支援兩種實作: p2pall_gatherp2p 實作會在注意力實作期間使用 peer-to-peer收發呼叫來累積金鑰值,並以非同步方式執行,允許與運算重疊。 all_gather實作會改用 AllGather 集合操作,並同步執行。

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

SMP 組態

步驟 2 SageMaker PyTorch 中,將下列參數新增至估算器的SMP組態字典。

{ ..., # other SMP config parameters "context_parallel_degree": 2 }