專家平行化 - Amazon SageMaker AI

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

專家平行化

專家混合 (MoE) 模型是一種採用稀疏方法的轉換器模型,與訓練傳統的密集模型相比,訓練更輕量。在此 MoE 神經網路架構中,每個輸入只會使用稱為專家的模型元件子集。此方法提供數種優點,包括更有效率的訓練和更快速的推論,即使模型大小較大也一樣。換句話說,使用相同的運算預算來訓練完整密集模型,您可以在使用 MoE 時容納較大的模型或資料集。

MoE 模型由多個專家組成,每個專家都由神經網路組成,通常是向前饋送網路 (FFN)。稱為路由器的閘道網路會決定要傳送哪些權杖給哪個專家。這些專家專門處理輸入資料的特定層面,使模型能夠更快地訓練、降低運算成本,同時實現與其對等密集模型相同的效能品質。若要進一步了解一般的專家混合,請參閱 NVIDIA 開發人員網站上的部落格文章在 LLM 架構中套用專家混合

專家平行化是一種平行化,可處理跨 GPU 裝置分割 MoE 模型專家。

SMP v2 與 NVIDIA Megatron 整合,實作專家平行化以支援訓練 MoE 模型,並在 PyTorch FSDP API 上執行。您可以繼續使用 PyTorch FSDP 訓練程式碼,並為訓練 MoE 模型啟用 SMP 專家平行化。

Hugging Face Transformer 模型與 SMP 專家平行化相容

SMP v2 目前為下列 Hugging Face Transformer 模型提供專家平行化支援。

設定專家平行化

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

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

注意

您可以搭配混合碎片資料平行化使用專家平行化。請注意,專家平行化目前與張量平行化不相容。

注意

此專家平行化訓練功能可在 SageMaker 程式庫和 PyTorch 程式庫的下列組合中使用:

  • SMP v2.3.0 及更新版本

  • SageMaker Python SDK 2.214.4 版及更新版本

  • PyTorch v2.2.0 及較新版本

在您的訓練指令碼中

步驟 1 中,使用初始化指令碼 torch.sagemaker.init() 啟用 SMP v2,並使用 torch.sagemaker.transform API 包裝模型,將 config 參數新增至 API 以啟用 MoE。下列程式碼片段示範如何為一般模型類別 AutoModelForCausalLM 啟用 SMP MoE,使用從頭開始訓練的 from_config 方法或微調的 from_pretrained 方法來提取 MoE 轉換器模型組態。若要進一步了解 SMP MoEConfig 類別,請參閱torch.sagemaker.moe.moe_config.MoEConfig

# Import the torch.sagemaker.transform API and initialize. import torch.sagemaker as tsm tsm.init() # Import transformers AutoModelForCausalLM class. from transformers import AutoModelForCausalLM # Import the SMP-implementation of MoE configuration class. from torch.sagemaker.moe.moe_config import MoEConfig # Define a transformer model with an MoE model configuration model = AutoModelForCausalLM.from_config(MoEModelConfig) # Wrap it by torch.sagemaker.transform with the SMP MoE configuration. model = tsm.transform( model, config=MoEConfig( smp_moe=True, random_seed=12345, moe_load_balancing="sinkhorn", global_token_shuffle=False, moe_all_to_all_dispatcher=True, moe_aux_loss_coeff=0.001, moe_z_loss_coeff=0.001 ) )

SMP 組態

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

{ ..., # other SMP config parameters "expert_parallel_degree": 8 }