專家平行處理

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

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

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

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

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

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

• 混合

設定專家平行處理

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

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

注意

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

注意

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

• SMP 2.3.0 版及更新版本

• SageMaker Python SDK 2.214.4 版及更新版本

• PyTorch 2.2.0 版及更新版本

在您的訓練指令碼中

在步驟 1 中，使用 初始化您的指令碼torch.sagemaker.init()以啟用 SMP v2，並使用 torch.sagemaker.transform API 包裝模型，將 config 參數新增至 API 以啟用 MoE。下列程式碼片段示範如何使用從頭開始訓練的 from_config方法或微調from_pretrained的 方法，為AutoModelForCausalLM提取 MoE 轉換器模型組態的一般模型類別啟用 SMP MoE。 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 組態字典。 PyTorch

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