翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

エキスパート並列処理

Mixture of Experts (MoE) モデルは、疎アプローチを採用したトランスフォーマーモデルの一種であり、従来の密結合モデルと比べてトレーニングが軽量化されます。この MoE ニューラルネットワークアーキテクチャでは、各入力に対して、モデルのコンポーネント (エキスパートと呼ばれる) の一部のみが利用されます。このアプローチにはいくつか利点があり、例えば、モデルのサイズが大きい場合でも、トレーニングを効率化し、推論を高速化できます。つまり、密結合のモデルをトレーニングする場合と同じコンピューティング予算で、MoE を使えばもっと大きなモデルやデータセットを扱うことができます。

MoE モデルは複数のエキスパート (専門家) で構成され、それぞれがニューラルネットワークを形成していますが、通常これはフィードフォワードネットワーク (FFN) です。ルーターと呼ばれるゲートネットワークが、どのトークンをどのエキスパートに送信するかを決定します。エキスパートは、それぞれが入力データの特定の側面の処理に特化しているので、モデルのトレーニングを高速化し、コンピューティングコストを削減しながらも、対応する密なモデルと同じパフォーマンス品質を達成できます。Mixture of Experts 全般の詳細については、NVIDIA Developer ウェブサイトの「Applying Mixture of Experts in LLM Architectures」を参照してください。

エキスパート並列処理は、並列処理の一種であり、MoE モデルのエキスパートを複数の GPU デバイス間に分散させます。

SMP v2 は、NVIDIA Megatron との統合によってエキスパート並列処理を実装し、MoE モデルのトレーニングに対応しています。また、PyTorch FSDP の API 上で動作します。PyTorch FSDP のトレーニングコードをそのまま使いつつ、SMP のエキスパート並列処理を有効にして MoE モデルをトレーニングできます。

SMP エキスパート並列処理と互換性のある Hugging Face Transformer モデル

SMP v2 は、現時点では、次の Hugging Face トランスフォーマーモデルでエキスパート並列処理をサポートしています。

エキスパート並列処理を設定する

expert_parallel_degree では、エキスパート並列処理の並列度を示す値を選択します。クラスター内の GPU の数を割り切れる値にする必要があります。例えば、GPU 数が 8 のインスタンスを使用している場合にモデルをシャーディング (分割) するには、2、4、または 8 を選択します。小さい数から始めて、モデルが GPU メモリに収まるまで徐々に数を増やすことをお勧めします。

次のコードスニペットでは、「SageMaker モデル並列処理ライブラリ v2 を使用する」で紹介した 2 ステップのプロセスに従って、SMP 初期化モジュール torch.sagemaker.init() をトレーニングスクリプトに追加し、トレーニングジョブランチャーの SMP 設定ディクショナリを JSON 形式で設定する方法を示しています。PyTorch モデルや PyTorch FSDP 設定については、一切変更する必要はありません。expert_parallel_degree パラメータの詳細については、「SMP v2 の主要機能の設定パラメータ」をご参照ください。

トレーニングスクリプト内

ステップ 1 の一環として、torch.sagemaker.init() でスクリプトを初期化して SMP v2 を有効にし、モデルを torch.sagemaker.transform API でラップします。その API に config パラメータを追加して MoE を有効にします。次のコードスニペットでは、汎用モデルクラスの AutoModelForCausalLM で SMP MoE を有効にする方法を示しています。MoE トランスフォーマーモデルの設定を、最初からトレーニングを行う場合は from_config メソッド、ファインチューニングする場合は from_pretrained メソッドを使用して取得します。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
}