【大規模訓練】変圧器におけるテンソルモデルの並列性

名前を付けたくない小さなPクラスメートにテキスト@

0序文

モデルの規模と計算能力の向上に伴い、トレーニングは単一のカードに限定されることはなくなり、モデルトレーニングへのマルチマシンとマルチカードの並列アプローチになりました。同時に、多くの異なる並列トレーニング最適化スキームがあり、モデルの並列処理はその1つです。以下では、3つの側面を紹介します。

モデル並列処理の動機と現状
モデルの並列処理の原理
データの並列処理とモデルの並列処理の組み合わせ

1モデル並列処理の動機と現状

1.1動機

グラフィックカードのメモリサイズによって制限される大きなモデルをトレーニングする場合、計算メモリを削減するために一連の技術的手法を適用する必要があります。主に2つの側面から：

1枚のカードで最適化されたコンピューティングメモリ。
複数のカードを使用して共有します。

オプティマイザー状態シャーディング（ZeRO）では、実際に使用されるいくつかの方法について説明しました。テンソル並列処理およびレイヤー内並列処理とも呼ばれるモデル並列処理は、その1つであり、ブロックマトリックスの原理を使用し、単一レイヤーのパラメーターが大きすぎるという問題の解決に重点を置いています。

非常に大きなモデルは、奥行きが広がるだけでなく、幅も広がります。これらの割合はモデル設計者が把握する必要があり、このような問題について調査・分析を行った作品もあります。モデルの大幅な拡張によってもたらされる超大容量ビデオメモリの場合、パイプライン並列処理（レイヤー間並列処理）またはいくつかの計算シーケンス最適化方法（チェックポインティングなど）を使用して問題を解決できます。モデル幅の拡大、モデルの並列処理は間違いなく現在のより成熟したプログラムです。

変圧器で線形層を広範に使用すると、モデルの並列処理の利点を十分に活用できます。

1.2ステータスクォー

顔認識タスクで使用される初期の方法はInsightFaceです。FCファイナルレイヤーとレイヤーのlossパラメーター分割と並列計算を実行します。

transformer現在広く使用されている方法はMegatronで、データ並列処理+モデル並列処理+パイプライン並列処理を使用して3072カードで1兆の大規模モデルトレーニングをサポートします。モデル並列処理の設計では、などのモデル機能transformerを最大限に活用し、通信を減らします。AttentionMLPlinear

最近、トーチシャードPytorchに基づくツールキットがオープンソースになりました。これは、のインターフェイスをなどのテクノロジーの混合使用と互換性があり、モデルの並列処理を軽量に使用する方法をユーザーに提供します。MegatronAMPZeRO

さらに、Parallelformers、fairscaleなどのオープンソースライブラリがあります。原理と設計は類似しているため、ここでは詳しく説明しません。

2モデルの並列処理の原理

モデルの並列処理の主なアイデアは、ネットワークレイヤーの入力、パラメーター、および操作を異なるカードに分割することであり、焦点はパラメーターにあります。

2.1数学的原理

レイヤーの場合linear、最も簡単なアイデアは、ブロック行列計算ルールを使用して結果の一貫性を取得することです。

次の式(1)と、行列をとによってそれぞれブロック(2)に分割します。AB

\（\ left [\ begin {matrix} X \ end {matrix} \ right] \ times \ left [\ begin {matrix} A_1＆A_2 \ end {matrix} \ right] = \ left [\ begin {matrix} XA_1 ＆XA_2 \ end {matrix} \ right] \ tag {1} \）

\（\ left [\ begin {matrix} Y_1＆Y_2 \ end {matrix} \ right] \ times \ left [\ begin {matrix} B_1 \\\\ B_2 \ end {matrix} \ right] = \ left [\ begin {matrix} Y_1B_1 + Y_2B_2 \ end {matrix} \ right] \ tag {2} \）

非linearレイヤーの場合、追加の設計は必要ありません。受信した入力が並列の場合、データの通信を遅らせることを検討したり、通信量を減らす方法を見つけることができます。たとえば、softmaxレイヤーは数学的な推移性に依存できます。つまり、極大値の最大値はグローバル最大値です。結合法則、つまりローカル合計の合計はグローバル合計であり、削減の効果を実現します。トラフィック。

2.1`transformer`モデル

2.1.1`linear`並列処理

次の図では、、をネットワーク内のアクティベーション値として扱い、、をXネットワークY内のレイヤーのパラメーターとして扱います。ZABlinear

最初の行に示されている並列方法は、列に従ってパラメーターを分割することです。これにより、取得された出力を列に従ってスプライスでき、並列処理なしで同等の結果を取得できます。

2行目は、パラメーターが行に分割されていることを示しています。行列の乗算の形状に一致させるには、入力を列に分割して、出力Z1とZ2を追加して最終結果を取得できるようにする必要がありZます。

操作を簡単にするために、ここでのセグメンテーションは等分割を指します。

各分割により、2つの追加の通信が発生します（図に示されているのはforwardフェーズであり、それに応じて、backwardフェーズにはルールに対応する1つの通信もあります）。通信コストを削減するために、比較的簡単な考え方は、2つのセグメンテーションを組み合わせることができれば、上の行の最後all_gatherと下の行の最初をsplit省略できるというものです。つまり、図の中央の破線で示されているように、パラメータ行列の列セグメンテーションと行セグメンテーションの組み合わせです。

トランスの構造は当然この組み合わせをサポートします。次の2つの図は、メガトロンの論文からのものです。

最初にMLPモジュール。図のパラメータは列にA分割され、得られた結果はモジュール。計算は現在の要素にのみ関連し、満たされているため、次の行列まで逆方向に伝播し続けることができます。乗算が発生します。A1A2GeLUGeLUGeLU([XA1, XA2])=[GeLU(XA1, XA2)]

図のfとgは、通信スライシングに関連する操作を表しています。MLPでは、fの前方はスライス（split）、後方は合計（all_reduce）、gの前方は合計（all_reduce）、後方はスライス（分割）です。

上の図に示すように、アテンションには同様の組み合わせがあります。マルチヘッドアテンション、、、を使用することが多くQ、ヘッド内KでVのみ相互作用し、ヘッドを並列に分割するだけなので、2つのマトリックスブロックの乗算間で、次のように計算されます。同等であり、前後の2つの通信操作はMLPに似ています。

Multiheads Attentionの典型的なモデル並列バージョンは、次のコードに示されています。

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.world_size = get_tensor_model_parallel_world_size() # 模型并行路数
        self.dim_per_partition = divide(dim, self.world_size)
        self.dim_per_attention_head = divide(dim, num_heads)
        self.num_heads_per_partition = divide(num_heads, self.world_size)

        self.qkv = ColumnParallelLinear(dim, dim * 3, gather_output=False) # gather_output=False，延迟 gather 的时间，和后面的 RowParallelLinear 组合起来
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = RowParallelLinear(dim, dim, input_is_parallel=True) # input_is_parallel=True，输入已经是按列切分的了
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, self.num_heads_per_partition, 3, self.dim_per_attention_head).permute(3, 0, 2, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, self.dim_per_partition)
        x = self.proj_drop(self.proj(x))
        return x

2.2.2`cross_entropy_loss`並列処理

cross_entropy_loss前の出力が列並列linearである場合、即時の通信コストは比較的大きくなります。実際、通信を節約する方法があります。並列バージョンにアクセスしますcross_entropy_loss。最初に計算するのは、次の式に示すように、softmaxの値です。ここpで、はモデルの並列グループのカード番号を表します。

\（softmax（x_i）= \ frac {e ^ {x_i}} {\ sum_ {j} {e ^ {x_j}}} = \ frac {e ^ {x_i -x _ {\ max}}} {\ sum_ { j} {e ^ {x_j-x _ {\ max}}}} = \ frac {e ^ {x_i-x _ {\ max}}} {\ sum_p {\ sum_ {k} {e ^ {x_k-x _ {\ max}}}}} \ tag {3} \）

\（x _ {\ max} = \ max_p（\ max_k（x_k））\ tag {4} \）

浮動小数点のオーバーフローを防ぐために、一般に、\（x_ {i} \）の値は計算中に\（x _ {\ max} \）から減算されるため、このプロセスでは2つの通信が行われます。 1つは通信Get \（x _ {\ max} \）です。ここで、\（\ underset {k} {\ max}（x_k）\）xは、モデルの並列グループ内のカードの最大値\（\ underset {p} {\ max}（\ underset {k} {\ max}（x_k））\）は、通信を通じて取得されたモデル並列グループの全体的な最大値を指しますx。別のコミュニケーションは、部分を合計することによって全体の合計を取得することです。

softmaxの値を取得した後target、いくつかのカテゴリの損失を分割して取得し、最後にすべてのカテゴリの損失を取得するために合計を行う必要があります。これにより、非常に小さな通信で合計3回の通信cross_entropy_lossで。

2.2.3追加処理

複数の緯線を組み合わせる場合、途中でドロップアウトなどの操作がある場合は、可能な限りランダム性の一貫性を確保する必要があります。

さらに、グローバルパラメータ情報が必要な場所では、対応するパラメータ通信も行う必要がありますclip_grad_norm。

3データの並列処理とモデルの並列処理の組み合わせ

実際のアプリケーションでは、モデルの並列処理を複数の並列メソッドと組み合わせることがよくあります。データ並列と混合する場合の特定のプロセスについては、以下で詳しく説明します。パイプライン並列については、当面は説明しません。

3.1メガトロンウェイ

メガトロンでは、モデルの並列処理は通常、単一のマシン内で使用され、データの並列処理はマシン間で使用されます。

次の図を例にとると、合計8枚のカード、2ウェイデータパラレル、4ウェイモデルパラレルがあり、データパラレルグループのカードのシリアル番号は0〜3で、合計2つのグループです。モデル並列グループ（モデル並列グループ）グループ内）内蔵カードのシリアル番号は0〜1、合計4グループです。

通信を初期化するときは、これらのカード番号とグループを明確に定義する必要があります。トレーニング中、各反復での初期パラメーター調整と勾配調整の両方で、対応するグループを指定する必要があります。

同じデータ並列グループ内では、純粋なデータ並列アプローチと比較してほとんど変更がありません。PyTorchを例にとると、データ並列モデルは通常、DistributedDataParallelを使用します。1つの方法は、パラメーターprocess_groupをデータ並列グループに置き換えることです。

# model 的 ddp 要传入 data_parallel_group 作为它的 process_group：在初始化 broadcast 阶段和 average_gradient 阶段起作用

model = DistributedDataParallel(model, process_group=data_parallel_group)

以下に、同じモデルの並列グループの動作について説明します。

まず、分散トレーニングのさまざまなカードのパラメーターが初期化中に同期されたままになるようにするには、上記のDDPに暗黙的に含まれるデータ並列グループでブロードキャストされるパラメーターに加えて、追加の非並列レイヤーを作成する必要があります。モデル並列グループブロードキャストのパラメータ。
第二に、同じモデルの並列グループでは、処理されたデータは同じです。これには、データサンプラーでの保証が必要なだけでなく、ネットワークに入力されたデータが異なる変形を起こさないようにするために、通常はフォワードの前に、データはグループ内で分析されます。ブロードキャストを実行します。
このように、パッチ埋め込みなどの非並列層では、グループ内で同じ動作が実行されます。並列アテンションと並列mlpを含む並列トランスフォーマー層などの並列層では、同じデータが引き続き処理されます。グループですが、パラメーターがスライスされているため、結果は部分的な結果になり、通信によって結合されます。

Megatronは、上記の関数を組み合わせて、DistributedDataParallelと同様のクラスに書き込みます。

3.2InsightFaceの方法

モデルのほとんどのレイヤーをモデルによって並列に分割できる場合、上記の方法は実装が比較的簡単で、スケーラビリティが高くなります。ただし、モデル内のほとんどのレイヤーを分割できない場合、このアプローチはより大きな問題を引き起こします。冗長な計算です。

InsightFaceの顔認識タスクでは、ほとんどのフロントレイヤーは分割できないレイヤーであり、最後の完全に接続されたレイヤーのみが過剰なパラメーターの問題を抱えています。この場合、完全に接続されたレイヤーとそれに続くcross_entropyのみがモデル化されます。 -並列計算であり、データ並列アプローチを維持します。

この並列分割方法に適応するには、データがモデルに並列に転送されてall_gather完全な入力を取得するときに、アクティブ化値に対して1回実行する必要があります。

この方法はより柔軟に使用できますが、ユーザーの要件は高くなります。パラメータの通信グループの分割は、より繊細に設計する必要があります。同時に、通信を適切に挿入するために、層間の接続を正確に判断する必要があります。

参考文献

オプティマイザー状態シャーディング（ZeRO）：https://zhuanlan.zhihu.com/p/394064174
InsightFace：https://github.com/deepinsight/insightface/tree/master/recognition
メガトロン：https://github.com/NVIDIA/Megatron-LM
トーチシャード：https://github.com/KaiyuYue/torchshard
Parallelformershttps://github.com/tunib-ai/parallelformers
フェアスケール：https://github.com/facebookresearch/fairscale
メガトロンの論文：https://arxiv.org/abs/1909.08053

読んでいただきありがとうございます、コメント欄にメッセージを残して議論してください〜

PSこの記事が気に入ったら、もっといいねをして、もっと多くの人に見てもらいましょう：D

パブリックアカウント「SenseParrots」をフォローして、人工知能フレームワークに関する最新の業界トレンドと技術的思考を入手してください。