Megatron-DeepSpeed を使用して大規模モデルをトレーニングしたいけれども、そのモデルが実装に含まれていない場合は、Megatron-DeepSpeed パッケージに移植することも可能です。Transformer ベースのモデルと想定すると、既存のコードをもとに独自の実装を簡単に追加することができます。

このページでは、Megatron-DeepSpeed を使用した LLaMA モデルの実装例を紹介します。
https://arxiv.org/pdf/2302.13971v1.pdf

【モデルを把握する】
実装プロセスをスムーズに進めるには、どの既存モデルが新しく実装するモデルと最も類似しているのかを把握してから、違いを洗い出しておくことが重要です。LLaMA の場合は GPT と類似しているので、GPT をベースにすることができます。また、要件やデータによっては、GPT 以外のトークナイザーを使用する必要があるかもしれない点にも注意しておきましょう。

LLaMA 系のモデルと GPT 系のモデルには、主に次のような違いがあります。
　・使用する活性化関数が異なる
　・位置埋め込みが異なる
　・レイヤー正規化が異なる
　・LLaMA モデルには、全 Transformer レイヤーの後に追加のレイヤー正規化も含まれる
　・LLaMA ではモデル開始 / 終了時の埋め込みレイヤーに異なる重み付けを使用する
　・LLaMA の MLP レイヤーでは GPT よりも小さい射影隠れ状態サイズを使用する
　・LLaMA は線形レイヤー内でバイアスを使用しない

【独自モデルの実装】
すべての Megatron-DeepSpeed モデルは、モデル・アーキテクチャーを定義する x_model.py (例: gpt_model.py) と、引数管理や補助関数など学習処理の詳細を処理する pretrain_x.py の、2 つの固有ファイルを持ちます。

新しくモデルを実装する際は、引数の追加や既存の引数への修正が必要です。新しい引数は通常、arguments.py ファイルに追加します。ただし、独自モデルに固有の新しい引数もあると思いますので、その場合は pretrain_x.py 内の ‘extra args provider’ を使用して新しい引数を追加すれば、すべてのモデルで使用される汎用 arguments.py への追加は不要です。
また、独自モデルの設定を既存の引数に合わせて、新しくデフォルト値を設定する必要もあります。そのため、Megatron-DeepSpeed では pretrain 関数を呼び出すときに ‘args defaults’ ディクショナリーを渡します。

LLaMA の場合、(レイヤー正規化を置き換える) RMSNorm と SwiGLU アクティベーション関数の追加が必要です。(ほかのモデルで使用されている既存の layernorm と gelu アクティベーション関数を含めた) すべてのオプションをサポートするために、arguments.py へ新しい引数 ‘layernorm-type’ と ‘activation-func-type’ を追加します。さらに、LLaMA はバイアスなしで線形レイヤーを使用するため、‘no-bias’ も追加します。LLaMA では回転位置埋め込み (RoPE) を使用しますが、コードはすでにさまざまなタイプの位置埋め込みをサポートしているので、既存 ‘position-embedding-type’ 引数への追加オプションとしては ‘rotary’ を追加するのみです。また、extra args provider を使用して、以下のとおり必要なデフォルトすべてを LLaMA に合わせて設定します。

新しい llama_model.py は gpt_model.py とよく似ています。適用するフローを検討する必要があり、Megatron-DeepSpeed はパイプと非パイプどちらのモジュールにも対応していますが、Habana では (パイプライン並列処理を行わない場合でも) パイプモジュールのみサポートしている点に注意してください。
https://github.com/HabanaAI/Model-References/blob/master/PyTorch/nlp/DeepSpeedExamples/Megatron-DeepSpeed/megatron/model/llama_model.py

モデルの forward 関数はこのファイル内の LLaMAModel クラス配下で定義します。ParallelTransformer の関連する引数に応じて LayerNorm レイヤーを RMSNorm へ置き換えるなど、適用する変更の一部は別のクラス内のものです。
https://github.com/HabanaAI/Model-References/blob/master/PyTorch/nlp/DeepSpeedExamples/Megatron-DeepSpeed/megatron/model/transformer.py

パイプモジュールは別のクラスとして定義し (LLaMAModelPipe 対 LLaMAModel)、このパイプモジュールは DeepSpeed のパイプエンジンを使用して (Tensor とパイプラインの) モジュール並列処理を可能にします。パイプモジュールの場合、結合レイヤーには TiedLayerSpec、結合なしレイヤーには LayerSpec を使い、全レイヤーのシーケンスを定義します。結合レイヤーは、例えば GPT の埋め込みレイヤーのように、レイヤー間の重み共有を処理します。LLaMA の埋め込みには結合レイヤーを使わず、‘vocab parallel projection’ を最終の埋め込みレイヤーとして定義します。

【新レイヤーの追加】
新しい独自モデルのレイヤーの一部はまだ実装されていません。例えば、LLaMA には RMSNorm を追加する必要があり、レイヤー正規化すべてを RMSnorm レイヤーに置き換えるほかに、megatron/model へ rmsnorm.py を追加します。

【既存実装の修正】
また、既存の実装に修正を加える必要があるかもしれません。例えば SwiGLU アクティベーション関数を使用するには、以下のとおり transformer.py 内で forward 関数そのものを編集します。