オープンソースの中国コミュニティチームは、共有の名のもとに、オープンソースの中国コミュニティの背後にあるストーリーを伝える初のライブブロードキャストを行いました。」

この記事は、Huawei クラウドコミュニティ「大規模モデル LLM の分散トレーニング」、著者: Hua Shanghua_Lancer から共有されたものです。

言語モデルのパラメーターと必要なトレーニングデータの量が急速に増加しているため、単一マシン上の限られたリソースでは大規模な言語モデルのトレーニングの要件を満たすことができなくなりました。分散トレーニング (分散トレーニング) システムは、大量のコンピューティングおよびメモリリソース要件の問題を解決するように設計する必要があります。

分散トレーニングシステム環境では、リソースのボトルネックを解決するために、モデルトレーニングタスクを複数のサブタスクに分割し、それらのサブタスクを複数のコンピューティングデバイスに分散する必要があります。しかし、数万のコンピューティングアクセラレーションチップを含むクラスターを使用して、数千億、さらには数兆のモデルパラメーターを持つ大規模な言語モデルをトレーニングするにはどうすればよいでしょうか?これには、クラスターアーキテクチャ、並列戦略、モデルアーキテクチャ、メモリの最適化、コンピューティングの最適化などの一連のテクノロジが含まれます。

分散機械学習システム、分散トレーニングクラスターアーキテクチャ、分散トレーニング並列戦略の基本概念を詳細に紹介し、DeepSpeed を例としてクラスター上で大規模な言語モデルをトレーニングする方法を紹介します。

1. 分散型研修の概要

分散トレーニングとは、機械学習または深層学習モデルのトレーニングタスクを複数のサブタスクに分解し、それらを複数のコンピューティングデバイスで並行してトレーニングすることを指します。図 1 は、単一のコンピューティングデバイスと複数のコンピューティングデバイスの例を示しています。ここでのコンピューティングデバイスは、中央処理装置 (CPU)、グラフィックスプロセッシングユニット (GPU)、またはテンソルプロセッシングユニット (Tensor Processing Unit) です。 TPU) は、ニューラルネットワークプロセッサ (ニューラルネットワークプロセッシングユニット、NPU) である場合もあります。

メモリは同じサーバー内の複数のコンピューティングデバイス間で共有できないため、これらのコンピューティングデバイスが 1 つのサーバーにあるか複数のサーバーにあるかに関係なく、そのシステムアーキテクチャは分散システムのカテゴリに分類されます。モデルトレーニングタスクには多くの場合、入力として多数のトレーニングサンプルがあり、1 台のコンピューティングデバイスを使用して完了することも、モデルトレーニングタスク全体をサブタスクに分割して異なるコンピューティングデバイスに分散して並列コンピューティングを実現することもできます。

その後、各計算機の出力を組み合わせて、最終的に単一の計算機と同等の計算結果を得る必要があります。各コンピューティングデバイスはサブタスクのみを担当する必要があり、複数のコンピューティングデバイスが並行して実行できるため、全体の計算をより速く完了し、最終的にコンピューティングプロセス全体を高速化できます。

図 1 単一のコンピューティングデバイスのコンピューティングと複数のコンピューティングデバイスの例

分散トレーニングシステムの設計を人々に促す最も重要な理由の 1 つは、単一のコンピューティングデバイスのコンピューティング能力ではもはやモデルのトレーニングをサポートするのに十分ではないということです。図 2 は、機械学習モデルの計算能力要件と、同じ期間に単一のコンピューティングデバイスが提供できる計算能力を示しています。図に示されているように、機械学習モデルは 2013 年の AlexNet から始まり、2022 年の 5,400 億パラメータを備えた Palm モデルまで、18 か月ごとに 56 回のペースで開発されています。モデルパラメーターのスケールが増加するにつれて、トレーニングデータの量の要件も指数関数的に増加し、コンピューティング能力の需要が高まります。

ただし、近年の CPU の計算能力の増加はムーアの法則よりもはるかに低く、計算加速デバイス (GPU、TPU など) が機械学習モデルに大量の計算能力を提供していますが、その成長率は依然として低いです。 18 か月ごとに 2 倍になるというムーアの法則を超えていません。機械学習モデルの開発に対応するには、増大するモデルのコンピューティング能力要件を満たすことができるのは分散トレーニングシステムのみです。

図 2 機械学習モデルのパラメータの増加とコンピューティングハードウェアの計算能力の増加の比較

分散トレーニングの全体的な目標は、全体的なトレーニング速度を向上させ、モデルトレーニングの全体的な時間を短縮することです。合計のトレーニング速度は、次の式を使用して簡単に見積もることができます。

総トレーニング速度 ∝ 単一デバイスの計算速度 × 計算デバイスの総数 × マルチデバイスの加速率

その中で、単一デバイスの計算速度は、主に単一計算加速チップの計算速度とデータ I/O 能力によって決まります。単一デバイスのトレーニング効率を最適化するための主な技術的手段には、混合精度トレーニング、オペレーターが含まれます。融合、勾配累積など; 分散トレーニングシステム内のコンピューティングデバイスの数が増えるほど、理論上のピークコンピューティング速度は高くなりますが、コンピューティングデバイスの数が増加すると、速度が急激に上昇します。加速率の低下、マルチデバイスの加速率は計算によって決定され、通信効率は最適化のためのアルゴリズムとネットワークトポロジの組み合わせによって決定されます。分散トレーニング並列戦略の主な目的は、マルチデバイスの加速率を向上させることです。分散型トレーニングシステム。

大規模な言語モデルのパラメーターの量と使用されるデータの量は非常に膨大であるため、トレーニングを完了するために分散トレーニングアーキテクチャが使用されます。文書 [5] では、NVIDIA V100 GPU を使用した GPT-3 のトレーニングプロセスのみが紹介されています。文書 [31] では、OPT が 992 個の NVIDIA A100 80G GPU を使用し、完全共有データ並列 [129] と Megatron-LM Tensor Parallelism [130] を採用していることが紹介されています。 , 全体のトレーニング時間はほぼ 2 か月です。

BLOOM[33] モデルの研究者は、使用されているハードウェアとシステムアーキテクチャの詳細を明らかにしました。モデルのトレーニングには合計 3.5 か月かかり、48 個のコンピューティングノードが使用されました。各ノードには 8 つの NVIDIA A100 80G GPU (合計 384 GPU) が含まれており、ノード内の GPU 間の通信に 4*NVLink を使用します。ノードは、4 枚の Omni-Path 100 Gbps ネットワークカードで構築された拡張 8 次元ハイパーキューブグローバルトポロジネットワークを使用して相互に通信します。

文献 [37] では、LLaMA モデルのトレーニングで使用されるクラスターの具体的な構成とネットワークトポロジは示されていませんが、さまざまなパラメータースケールの合計 GPU 時間が示されています。 LLaMA モデルのトレーニングには A100-80GB GPU が使用され、LLaMA-7B モデルのトレーニングには 82432 GPU 時間が必要となり、LLaMA-13B モデルのトレーニングには 135168 GPU 時間が必要となり、LLaMA-33B モデルのトレーニングには 530432 GPU 時間がかかり、LLaMA-65B モデルのトレーニングには最大 1022362 GPU のコストがかかります時間。 LLaMA で使用される学習データの量は OPT モデルや BLOOM モデルをはるかに上回っているため、モデルパラメーターの数は上記 2 つのモデルよりはるかに少ないにもかかわらず、必要な計算量は依然として膨大です。

分散トレーニングシステムを使用すると、大規模な言語モデルのトレーニングサイクルを、単一のコンピューティングデバイスでの数十年から、数千台のコンピューティングデバイスを使用した場合の数十日に短縮できます。ただし、分散トレーニングシステムは、クラスター内のすべてのリソースが完全に活用されるようにするために、コンピューティングの壁、ビデオメモリの壁、通信の壁などのさまざまな課題を克服する必要があり、それによってトレーニングプロセスを加速し、トレーニングサイクルを短縮します。

• 計算の壁: 単一のコンピューティングデバイスが提供できる計算能力と、大規模な言語モデルに必要な総計算量との間には、大きな差があります。 2022 年 3 月にリリースされた NVIDIA H100 SXM のシングルカード FP16 の計算能力はわずか 2000 TFLOP ですが、GPT-3 は
合計 314 ZFLOP の計算能力を必要とします。この 2 つの違いは 8 桁あります。

• ビデオメモリの壁: 単一のコンピューティングデバイスでは、大規模な言語モデルのパラメータを完全に保存することはできません。 GPT-3 には 1,750 億個のパラメータが含まれており、FP16 形式で保存すると、700 GB のコンピューティングデバイスメモリ領域が必要になりますが、NVIDIA H100 GPU には 80 GB のビデオメモリしかありません。

• コミュニケーションウォール: 分散トレーニングシステムでは、コンピューティングデバイス間で頻繁なパラメータの送信と同期が必要です。通信遅延と帯域幅の制限により、これがトレーニングプロセスのボトルネックになる可能性があります。 GPT-3 トレーニングプロセス中に、分散システム内に 128 のモデルコピーがある場合、各反復中に少なくとも 89.6 TB の勾配データを送信する必要があります。 2023 年 8 月の時点で、単一の InfiniBand リンクは 800Gb/s を超える帯域幅しか提供できません。コンピューティングの壁とビデオメモリの壁は、単一のコンピューティングデバイスの限られたコンピューティングおよびストレージ能力と、モデルの膨大なコンピューティングおよびストレージ要件との間の矛盾から生じます。この問題は分散トレーニング手法を使用することで解決できますが、分散トレーニングではコミュニケーションの壁という課題に直面します。複数のマシンとカードをトレーニングする中で、これらの問題が徐々に明らかになりました。大規模なモデルのパラメータが増加すると、対応するクラスターサイズも増加し、これらの問題がより顕著になります。同時に、大規模なクラスターが長時間トレーニングされると、機器の故障がトレーニングプロセスに影響を与えたり、中断したりする可能性があり、分散システムに高い要求も課せられます。

2. 分散トレーニングの並行戦略

分散トレーニングシステムの目標は、単一ノードモデルのトレーニングを同等の分散並列モデルトレーニングに変換することです。大規模な言語モデルの場合、トレーニングプロセスは、データと損失関数に基づく最適化アルゴリズムを使用してニューラルネットワークモデルのパラメーターを更新するプロセスです。シングルノードモデルトレーニングシステムの構造は図 3 に示されており、主にデータとモデルの 2 つの部分で構成されます。トレーニングプロセスは、複数のデータミニバッチ (ミニバッチ) によって完了します。

図内のデータは、小さなデータのバッチを表しています。トレーニングシステムは、小さなバッチのデータを使用して、損失関数と最適化アルゴリズムに基づいて勾配を生成し、モデルパラメーターを修正します。大規模言語モデルの多層ニューラルネットワークの実行プロセスは、計算グラフ (Computational Graph) で表現できます。このグラフには相互接続された複数の演算子 (演算子) があり、各演算子はニューラルネットワーク層 (ニューラルネットワーク層) を実装し、パラメーターはトレーニング中にこの層によって更新される重みを表します。

図 3 単一デバイスのモデルトレーニングシステム

計算グラフの実行プロセスは、順計算と逆計算の 2 段階に分けることができます。前方計算のプロセスでは、データを最初の演算子に読み込み、対応する出力構造を計算し、最後の演算子が終了するまで前方計算プロセスを繰り返します。逆計算プロセスは、最適化関数と損失に基づいて、各オペレーターが順番に勾配を計算し、その勾配を使用してローカルパラメーターを更新します。逆計算が完了し、データミニバッチの計算が完了すると、システムは次のデータミニバッチを読み取り、次のラウンドのモデルパラメーターの更新を続行します。

単一デバイスのモデル学習システムのプロセスによれば、並列加速を実行すると、データとモデルの 2 つの次元から考慮できることがわかります。まず、データを分割し (パーティション)、同じモデルを複数のデバイスにコピーし、異なるデータシャードを並列実行できます。この方法は通常、データ並列処理 (DP) と呼ばれます。モデルを分割し、モデル内の演算子を複数のデバイスに分散してそれぞれを完了することもできます。この方法は通常、モデル並列処理 (MP) と呼ばれます。非常に大規模な言語モデルをトレーニングする場合、より高度な並列処理を実現するために、データとモデルを同時に分割する必要があることがよくあります。この方法は、ハイブリッド並列処理 (HP) と呼ばれます。

2.1. データの並列性

データ並列システムでは、各コンピューティングデバイスにはニューラルネットワークモデル全体の完全なコピー (モデルレプリカ) があり、反復処理の際、各コンピューティングデバイスにはデータサンプルのバッチのサブセットのみが割り当てられ、サンプルのバッチに基づいて割り当てられます。データのサブセットは、ネットワークモデルの前方計算に使用されます。バッチ内のトレーニングサンプルの数が N で、並列計算に M 台のコンピューティングデバイスが使用され、各コンピューティングデバイスに N/M サンプルが割り当てられると仮定します。順方向計算が完了すると、各コンピューティングデバイスはローカルサンプルに基づいて損失誤差を計算して勾配 Gi (i はアクセラレータカード番号) を取得し、ローカル勾配 Gi をブロードキャストします。すべてのコンピューティングデバイスは、他のアクセラレーションカードによって与えられた勾配値を集計し、平均勾配 (ΣNi=1Gi)/N を使用してモデルを更新してバッチトレーニングを完了する必要があります。図 4 は、2 台のコンピューティングデバイスで構成されるデータ並列トレーニングシステムの例を示しています。

図4 2ノードデータ並列学習システムの例

データ並列トレーニングシステムは、コンピューティング機器を追加することで、全体のトレーニングスループットと 1 秒あたりのグローバルバッチサイズ (1 秒あたりのグローバルバッチサイズ) を効果的に向上させることができます。単一のコンピューティングデバイスのトレーニングと比較した主な違いは、各コンピューティングデバイスの最終結果がすべてのプロセスの勾配の平均になるように、逆計算の勾配をすべてのコンピューティングデバイスで同期する必要があることです。

一般的なニューラルネットワークフレームワークには、TensorFlow DistributedStrategy、PyTorch Distributed、Horovod DistributedOptimizer などのデータ並列処理の特定の実装があります。 Transformer アーキテクチャに基づく大規模な言語モデルの各演算子はバッチデータではなく単一データに依存するため、データの並列処理はその計算ロジックに影響を与えません。一般に、各トレーニングデバイスの順方向計算は独立しており、計算ロジックには影響しません。同期の問題が含まれます。データ並列トレーニングは加速率が最も高くなりますが、モデルのコピーを各デバイスにバックアップする必要があり、比較的大量のビデオメモリを消費します。

PyTorch DistributedDataParallel を使用して単一サーバー上で複数のアクセラレータカードトレーニングを実装するコードは次のとおりです。まず、データセットのサンプルをランダムに中断し、それらをさまざまなコンピューティングデバイスに分散する DistributedSampler クラスを構築します。

クラス DistributedSampler(サンプラー):
  def __init__(self、dataset、num_replicas=None、rank=None、shuffle=True、seed=0):
    num_replicas が None の場合:
        dist.is_available() でない場合:
            raise RuntimeError("配布パッケージが利用可能である必要があります")
        num_replicas = dist.get_world_size()
    ランクがなしの場合:
        dist.is_available() でない場合:
            raise RuntimeError("配布パッケージが利用可能である必要があります")
        ランク = dist.get_rank()
    self.dataset = データセット #データセット
    self.num_replicas = num_replicas #プロセス数のデフォルトは world_size (GPU の数)
    self.rank = ランク # 現在どのプロセス/GPU に属しているか
    self.epoch = 0
    self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
    #プロセスあたりのサンプル数
    self.total_size = self.num_samples * self.num_replicas #データセット内の合計サンプル数
    self.shuffle = shuffle # データセットをシャッフルするかどうか
    self.seed = シード

def __iter__(self):
# 1. シャッフル処理: データセットの順序を乱す
    self.shuffle の場合:
        # エポックとシードに基づいて難読化する
        g = トーチ.ジェネレーター()
        # ここで self.seed は固定値です。 set_epoch によって self.epoch を変更すると、初期化シードを変更できます。
        # これにより、各エポックでのデータセットのシャッフル順序を異なるものにすることができるため、各エポックで、
        # 各 GPU は異なるデータを取得するため、より適切なトレーニングが促進されます。
        g.manual_seed(self.seed + self.epoch)
        インデックス = torch.randperm(len(self.dataset),generator=g).tolist()
    それ以外：
        インデックス = リスト(範囲(len(self.dataset)))
    # データ補足
    インデックス += インデックス[:(self.total_size - len(インデックス))]
    len(インデックス) == self.total_size をアサート
    # データを割り当てる
    インデックス = インデックス[self.rank:self.total_size:self.num_replicas]
    len(インデックス) == self.num_samples をアサートします
    iter(インデックス)を返す
def __len__(自分自身):
    self.num_samples を返す
def set_epoch(self, epoch):

    self.epoch = エポック

DistributedSampler を使用して、次のように完全なトレーニングプログラムサンプル main.py を構築します。

引数解析をインポートする
私たちを輸入してください
インポートシュティル
インポート時間
インポート警告
numpyをnpとしてインポート
warnings.filterwarnings('無視')
輸入トーチ
torch.nn を nn としてインポート
torch.nn.Parallel をインポート
torch.backends.cudnn を cudnn としてインポート
import torch.dist として配布
torch.optim をインポート
torch.utils.data をインポートする
torch.utils.data.distributed をインポートします
torch.utils.data.distributed からインポート DistributedSampler
モデルから DeepLab をインポート
データセットからのインポート都市景観
parser = argparse.ArgumentParser(description='DeepLab')
parser.add_argument('-j', '--workers'、default=4、type=int、metavar='N',
help='データ読み込みワーカーの数 (デフォルト: 4)')
parser.add_argument('--epochs'、デフォルト=100、type=int、metavar='N'、
help='実行する合計エポック数')
parser.add_argument('--start-epoch'、default=0、type=int、metavar='N'、
help='手動エポック番号 (再起動時に役立ちます)')
parser.add_argument('-b', '--batch-size'、default=3、type=int、
メタ変数='N')
parser.add_argument('--local_rank'、default=0、type=int、help='分散トレーニングのノードランク')
args = parser.parse_args()
torch.distributed.init_process_group(backend="nccl") #初期化
print("トレーニングに GPU を使用: {}".format(args.local_rank))
# モデルを作成する
モデル = DeepLab()
torch.cuda.set_device(args.local_rank) #現在のグラフィックカード
model = model.cuda() # モデルはグラフィックス カードに配置されます
モデル = torch.nn.Parallel.DistributedDataParallel(モデル, device_ids=[args.local_rank],
    Output_device=args.local_rank, find_unused_parameters=True) # データ並列処理
基準 = nn.CrossEntropyLoss().cuda()
オプティマイザー = torch.optim.SGD(model.parameters(), args.lr,
    勢い=args.momentum、weight_decay=args.weight_decay)
train_dataset = 都市景観()
train_sampler = DistributedSampler(train_dataset) # データを配布する
train_loader = torch.utils.data.DataLoader(train_dataset、batch_size=args.batch_size、
    shuffle=False、num_workers=args.workers、pin_memory=True、sampler=train_sampler)

次のコマンドラインから上記のプログラムを起動します。

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 main.py

2.2 モデルの並列性

モデル並列処理は、単一ノードのメモリ不足の問題を解決するためによく使用されます。 1,750 億個のパラメータを含む GPT-3 モデルを例に挙げます。モデル内の各パラメータが 32 ビット浮動小数点数で表される場合、モデルは 700GB (つまり 175G × 4 バイト) のメモリを占有する必要があります。 16 ビット浮動小数点数で表され、各モデルのコピーには 350GB のメモリが必要です。 NVIDIA が 2022 年 3 月にリリースした H100 アクセラレータカードは、80GB のビデオメモリのみをサポートしており、モデル全体を完全に搭載することはできません。モデルの並列性は、計算グラフの観点から次の 2 つの形式に分類できます。

(1) モデルのレイヤーに従って異なるデバイスに分割します。つまり、レイヤー間並列処理またはオペレーター間並列処理 (オペレーター間並列処理)、パイプライン並列処理 (パイプライン並列処理、PP) とも呼ばれます。

(2) 計算層のパラメータを異なるデバイス、つまり層内並列処理または演算子内並列処理 (Intra-operator Parallelism) に分割します。これは、テンソル並列処理 (Tensor Parallelism、TP) とも呼ばれます。

2 ノードモデルの並列トレーニングシステムのサンプルを図 4.9 に示します。左側はパイプライン並列処理であり、モデルの異なる層は異なるデバイスに分割されており、右側はテンソル並列処理であり、同じ層内の異なるパラメータです。デバイス内のさまざまな計算に分割されます。

パイプラインの並列処理

パイプライン並列処理 (PP) は、モデルの各層をセグメントで処理し、各セグメントを異なるコンピューティングデバイスに分散する並列コンピューティング戦略です。これにより、前後のステージがパイプラインおよびバッチで動作できるようになります。パイプライン並列処理は通常、単一のコンピューティングデバイス上のメモリ不足の問題を効果的に解決するために、大規模モデルの並列システムに適用されます。図 4.6 は、順方向計算と逆方向計算を含む 4 つの計算デバイスで構成されるパイプライン並列システムを示しています。このうち、F1、F2、F3、および F4 はそれぞれ異なるデバイス上にある 4 つの順方向パスを表し、B4、B3、B2、および B1 は同じく 4 つの異なるデバイス上にある逆方向逆方向パスを表します。ただし、図からわかるように、計算グラフの下流デバイス (Downstream Device) は、自身の計算を開始する前に、上流デバイス (Upstream Device) の計算が完了するまで待機するために、長時間アイドル状態を維持する必要があります。タスク。

図5 2ノードモデル並列学習システムの例

この状況により、デバイスの平均使用量が大幅に減少し、パイプラインバブルとも呼ばれるモデル並列処理バブルが形成されました。

図 6 パイプラインの並列例

単純なパイプライン戦略によって生成される並列バブルにより、システムはコンピューティングリソースを完全に活用できなくなり、システム全体のコンピューティング効率が低下します。並列バブルを減らすために、文献 [131] では、ミニバッチをさらに小さなマイクロバッチに分割し、パイプライン並列スキームを使用して一度に 1 つのマイクロバッチ×データを処理する GPipe メソッドが提案されています。

現段階の計算が完了して結果が得られると、マイクロバッチの結果が下流のデバイスに送信され、同時に次のマイクロバッチのデータの処理が開始されるため、コストを削減できます。ある程度平行な泡。図 7GPipe ポリシーパイプラインの並列例。図に示すように、順方向 F1 の計算は F11、F12、F13、F14 に分割されます。F11 の計算が計算機 1 で完了した後、F21 の計算が計算機 2 で開始されます。同時に、F12 がコンピューティングデバイス 1 の計算で並行して開始されます。元のパイプライン並列方式と比較して、GPipe パイプライン方式は並列バブルを効果的に減らすことができます。

図 7 GPipe ポリシーパイプラインの並列例

GPipe 戦略は特定の並列バブルを減らすことができますが、逆方向計算はミニバッチ内のすべての前方計算が完了した後にのみ開始できます。したがって、依然として多くの並列バブルが生成され、システムの並列効率が低下します。 Megatron-LM[132] は、1 つの順方向チャネルと 1 つの逆方向チャネルで構成される 1F1B パイプライン戦略を提案しました。 1F1B パイプライン戦略では、タスクスケジューリングメカニズムが導入され、ダウンストリームデバイスがアップストリームの計算を待機している間に他の並列タスクを実行できるようになり、デバイスの使用率が向上します。 1F1B は、図 8 に示すように、ノンインターリーブとインターリーブの 2 つのスケジューリング方法を提供します。

1F1B ノンインターリーブスケジューリングモードは 3 つの段階に分割できます。 1 つ目はウォームアップフェーズで、コンピューティングデバイスでさまざまな数の前方計算が実行されます。次のフェーズは順方向-逆方向フェーズであり、コンピューティングデバイスは順方向計算を実行し、次に逆方向計算を順番に実行します。最後のフェーズは逆方向フェーズであり、コンピューティングデバイスは最後の逆方向計算を完了します。 GPipe 戦略と比較して、非インターリーブスケジューリングモードはメモリの節約において優れたパフォーマンスを発揮します。ただし、1 ラウンドの計算を完了するには GPipe 戦略と同じ時間が必要です。

1F1B インターリーブスケジューリングモードでは、マイクロバッチの数がパイプラインステージの整数倍である必要があります。各デバイスは、複数の連続するレイヤーの計算を単独で担当するのではなく、モデルナゲットと呼ばれる複数のレイヤーのサブセットを処理できます。具体的には、以前のモデルでは、デバイス 1 がレイヤー 1 ～ 4 を担当し、デバイス 2 がレイヤー 5 ～ 8 を担当する、というようになります。ただし、新しいモードでは、デバイス 1 はレイヤー 1、2、9、10 を処理でき、デバイス 2 はレイヤー 3、4、11、12 などを処理できます。このモードでは、各デバイスがパイプラインの複数のステージに割り当てられます。たとえば、デバイス 1 は、ウォームアップフェーズ、順方向計算フェーズ、および逆方向計算フェーズのタスクの一部に関与する場合があります。各デバイスは、異なるステージでコンピューティングタスクを並行して実行できるため、パイプラインの並列性をより有効に活用できます。このモードは、メモリ消費量の点で優れたパフォーマンスを発揮するだけでなく、計算効率も向上し、大規模モデルの並列システムがより効率的に計算タスクを完了できるようになります。

図 8 1F1Bパイプライン並列戦略の例

PyTorch には、パイプラインを実装するための API 関数 Pipe も含まれています。具体的な実装については、「torch.distributed.pipeline.sync.Pipe」クラスを参照してください。次のように、この API を使用して、2 つの異なるコンピューティングデバイスに配置される 2 つの線形レイヤーを含むサンプルを構築できます。

{#
ステップ 0. 最初に RPC フレームワークを初期化する必要があります。
os.environ['MASTER_ADDR'] = 'ローカルホスト'
os.environ['MASTER_PORT'] = '29500'
torch.distributed.rpc.init_rpc('ワーカー'、ランク=0、ワールドサイズ=1)
# ステップ 1: 2 つの線形層を含むモデルを構築する
fc1 = nn.Linear(16, 8).cuda(0)
fc2 = nn.Linear(8, 4).cuda(1)
# ステップ 2: 2 つのレイヤーを nn.Sequential でラップします。
モデル = nn.Sequential(fc1, fc2)
# ステップ 3: パイプを構築する (torch.distributed.pipeline.sync.Pipe)
モデル = パイプ(モデル、チャンク = 8)
# トレーニング/推論を行う
入力 = torch.rand(16, 16).cuda(0)
Output_rref = モデル(入力)
}

テンソル並列性

Tensor Parallelism (TP) は 2 つの問題を解決する必要があります。1 つはモデルの特定の構造と演算子の種類に応じてパラメーターを異なるデバイスに分割する方法、もう 1 つは分割後の数学的一貫性を確保する方法です。大規模な言語モデルは、Transformer 構造に基づいています。Transformer 構造は、主に、埋め込み表現 (Embedding)、行列乗算 (MatMul)、およびクロスエントロピー損失 (Cross Entropy Loss) の計算という 3 つの演算子で構成されます。

これら 3 種類の演算子はまったく異なるため、パラメータを異なるデバイスに分割するために、対応するテンソル並列戦略 [130] を設計する必要があります。エンベディングオペレーターの場合、語彙の総数が非常に多い場合、単一のコンピューティングデバイスのビデオメモリではエンベディングレイヤーのパラメーターを収容できなくなります。たとえば、語彙の数が 64000、埋め込み表現の次元が 5120、タイプが 32 ビット精度の浮動小数点数を使用する場合、パラメータのレイヤー全体に必要なビデオメモリは約 64000 × 5120 × 4/1024 になります。 /1024 = 1250MB、逆勾配も同じです。ストレージだけで 1250MB、ほぼ 2.5GB が必要です。

プレゼンテーション層に埋め込まれたパラメータは単語の次元に応じて分割でき、各コンピューティングデバイスは単語ベクトルの一部のみを保存し、各デバイス上で部分的な単語ベクトルを要約することで完全な単語ベクトルが取得されます。図 4.9 は、単一ノードの埋め込みと 2 ノードのテンソル並列処理の概略図を示しています。

単一ノードで Embedding 操作を実行します。bz はバッチサイズ (バッチサイズ)、Embedding のパラメーターサイズは [word_size, hidden_size] で、[bz, hidden_size] テンソルが計算されます。図 4.9 の Embedding テンソル並列例では、Embedding パラメーターが word_size 次元に沿って 2 つのブロックに分割されており、各ブロックのサイズは [word_size/2, hidden_size] で、それぞれ 2 つのデバイスに保存されます。各ノードが独自の単語リストをクエリするときに、単語リストが見つからない場合、その単語の表現は 0 になります。各デバイスをクエリした後、最後に、AllReduce_Sum 通信を介してデバイス間で合計された [bz, hidden_size] 結果テンソルが取得されます。完全な完全な結果から、ここでの出力結果は単一のコンピューティングデバイスによって実行された結果と一致していることがわかります。

図 9 2 ノードの埋め込み演算子のテンソル並列例

行列乗算のテンソル並列処理 (MatMul) は、行列ブロック乗算原理を最大限に活用する必要があります。たとえば、次の行列乗算 Y = X ×A を実装するには、X は次元 M × N の入力行列、A は次元 N ×K のパラメータ行列、Y は次元 M ×K の結果行列です。パラメータ行列 A が非常に大きい場合、または 1 枚のカードのビデオメモリ容量を超える場合は、パラメータ行列 A を複数のカードに分割し、集団通信を通じて結果を収集して、最終結果が数学的に正確であることを確認できます。 1 台のコンピューティングデバイスに相当する計算結果。パラメータ行列 A をセグメント化するには 2 つの方法があります。

(1) パラメータ行列 A が列に分割され、行列 A が列に分割されます: A = [A1,A2]

(2) パラメータ行列 A が行に分割され、行列 A が行に分割されます。

図 10 は、パラメータ行列を列ごとに分割する例を示しています。パラメータ行列 A は、A1 と A2 を 2 つのコンピューティングデバイスにそれぞれ配置します。 2 つのコンピューティングデバイスは、それぞれ Y1 = X ×A1 と Y2 = X ×A2 を計算します。計算が完了すると、複数のコンピューティングデバイスが相互に通信して他のコンピューティングデバイスの計算結果を取得し、それらを結合して最終結果行列 Y を取得します。この結果は、単一のコンピューティングデバイスの計算結果と数学的に等価です。

図 10 2 ノードの行列乗算演算子テンソルを列ごとに並列分割する例

図 11 は、パラメータ行列を列と行で分割する例を示しています。行列の乗算規則を満たすには、入力行列 X を列 X = [X1|X2] で分割する必要があります。同時に、行列はブロックに分割され、2 つの計算デバイスに配置され、それぞれ Y1 = X1 × A1 と Y2 = X2 × A2 が計算されます。計算が完了した後、複数のコンピューティングデバイスが通信して他のカード上の計算結果を取得および削減し、最終的な結果行列 Y を取得できます。同様に、この分割方法は、数学的計算の等価性を保証するだけでなく、単一のコンピューティングデバイスがビデオメモリに対応できないという問題を解決し、分割によって単一のコンピューティングデバイスがパラメータ A に対応できることを保証することもできます。

Transformer の FFN 構造には 2 つの全結合 (FC) 層が含まれています。つまり、2 つの行列乗算があり、図 4.12 に示すように、これら 2 つの行列乗算では上記 2 つのセグメンテーション方法が採用されています。第１のＦＣ層のパラメータ行列は列ごとにブロックに分割され、第２のＦＣ層のパラメータ行列は行ごとにブロックに分割される。このように、最初の FC 層の出力は 2 番目の FC 層のデータ入力要件 (列ごとに分割) を正確に満たすため、最初の FC 層の後の要約通信動作を省略できます。マルチヘッドセルフアテンションメカニズムのテンソル並列性は FFN に似ています。複数の独立したヘッドを備えているため、行列分割方法は FFN よりも容易です。詳細については[130]を参照してください。

分類ネットワークの最後の層は通常、Softmax 演算子と Cross_entropy 演算子を使用して、クロスエントロピーロス (クロスエントロピーロス) を計算します。カテゴリの数が非常に多い場合、単一のコンピューティングデバイスのメモリではロジット行列を保存および計算できなくなります。このタイプのオペレータの場合、カテゴリの次元に従って分割することができ、同時に、中間結果の通信を通じて最終的なグローバルなクロスエントロピー損失を取得できます。

図 11 2 ノードの行列乗算演算子テンソルを行ごとに並列分割する例

図 12 FNN 構造テンソル並列図

最初に計算するのはソフトマックス値です。式は次のとおりです。

このうち、p はテンソル並列度のデバイス番号を表します。 Softmax の計算結果を取得した後、同時にラベル Target をカテゴリごとに分割し、各デバイスが損失の一部を取得します。最後に、すべてのカテゴリの損失を取得するために別の通信が実行されます。プロセス全体では、クロスエントロピー損失の計算を完了するために 3 回の少量の通信のみが必要です。 PyTorch は、きめ細かいテンソルレベルの並列 API、DistributedTensor を提供します。また、「nn.Module」でテンソル並列処理を実行するための粗粒度のモデルレベル API も提供します。次のコード行を使用して、大きなテンソルをシャーディングできます。

輸入トーチ
torch.distributed._tensor からインポート DTensor、DeviceMesh、Shard、distribute_tensor
# 利用可能なデバイス (マルチホストまたは単一ホスト) でデバイス メッシュを構築します
device_mesh = DeviceMesh("cuda", [0, 1, 2, 3])
# 行単位のシャーディングを行う場合
rowwise_placement=[シャード(0)]
# 列ごとのシャーディングを行う場合
Colwise_placement=[シャード(1)]
big_tensor = torch.randn(888, 12)
# 返された分散テンソルは、プレースメントで指定されたディメンション全体でシャーディングされます
rowwise_tensor = distribution_tensor(big_tensor, device_mesh=device_mesh,placements=rowwise_placement)

パラメータとしてすでに「torch.Tensor」を持っている「nn.Linear」のようなモジュールの場合、モデルレベルでテンソル並列処理を実行するためのモジュールレベル API「distribute_module」も提供されています。リファレンスコードは次のとおりです。

輸入トーチ
torch.distributed._tensor から DeviceMesh、Shard、distribute_tensor、distribute_module をインポート
クラスMyModule(nn.Module):
    def __init__(自分自身):
        super().__init__()
        self.fc1 = nn.Linear(8, 8)
        self.fc2 = nn.Linear(8, 8)
        self.relu = nn.ReLU()
        def forward(self, input):
            return self.relu(self.fc1(入力) + self.fc2(入力))
    メッシュ = DeviceMesh(device_type="cuda", メッシュ=[[0, 1], [2, 3]])
    def shard_params(mod_name, mod, メッシュ):
        rowwise_placement = [シャード(0)]
        def to_dist_tensor(t): 戻り distribute_tensor(t, メッシュ, rowwise_placement)
        mod._apply(to_dist_tensor)
    sharded_module = distribution_module(MyModule()、メッシュ、partition_fn=shard_params)
    def shard_fc(mod_name, mod, メッシュ):
        rowwise_placement = [シャード(0)]
        mod_name == "fc1"の場合:
            mod.weight = torch.nn.Parameter(distribute_tensor(mod.weight, メッシュ, rowwise_placement))
    sharded_module = distribution_module(MyModule()、メッシュ、partition_fn=shard_fc)

2.3 ハイブリッド並列処理

ハイブリッド並列処理 (HP) は、データ並列処理、パイプライン並列処理、テンソル並列処理などの複数の並列戦略を組み合わせたものです。ハイブリッド並列処理では、さまざまな並列戦略を組み合わせることで、さまざまな並列戦略を最大限に活用して、コンピューティングのパフォーマンスと効率を最大化できます。

数千億規模の大規模な言語モデルの場合、通常、各サーバー内でテンソル並列戦略が使用されます。この戦略には大量のネットワーク通信が含まれるため、サーバー内の異なるコンピューティングデバイス間で高速通信帯域幅を利用する必要があります。サーバ。パイプライン並列処理により、モデルのさまざまなレイヤーが複数のステージに分割され、各ステージは異なるマシンによって計算されます。これにより、複数のマシンの計算能力を最大限に活用し、計算結果や中間データをマシン間の高速通信で転送できるため、全体の計算速度と効率が向上します。

最後に、データ並列戦略が外側の層に重ね合わされて、同時実行数が増加し、全体的なトレーニング速度が向上します。データの並列処理により、トレーニングデータは並列処理のために複数のサーバーグループに分散され、サーバーの各グループは異なるデータバッチを処理します。これにより、複数のサーバーのコンピューティングリソースを最大限に活用し、トレーニングの同時実行性が向上し、全体のトレーニング速度が向上します。

BLOOM は、トレーニングに Megatron-DeepSpeed[104] フレームワークを使用します。これは主に 2 つの部分で構成されます。Megatron-LM は、テンソル並列機能とデータ読み込みプリミティブを提供します。DeepSpeed は、ZeRO オプティマイザー、モデルパイプライン、および従来の分散トレーニングコンポーネントを提供します。このようにして、データ、テンソル、パイプラインの 3 次元並列処理を実現できます。BLOOM モデルのトレーニングで使用される並列コンピューティング構造を図 14 に示します。

BLOOM モデルのトレーニングでは、48 台の NVIDIA DGX-A100 サーバーのクラスターを使用します。各 DGX-A100 サーバーには、8 個の NVIDIA A100 80GB GPU が含まれており、合計 384 個になります。 BLOOM トレーニングで採用された戦略は、データの並列化のためにまずクラスターを 48 個のグループに分割することです。

次に、モデル全体をパイプライン並列化のために 12 ステージに分割します。各ステージのモデルは、テンソル並列処理のために 4 つの GPU に分割されています。同時に、BLOOM は、モデルのビデオメモリ占有をさらに削減するために、ZeRO (Zero Redundancy Optimizer) [134] も使用します。上記の 4 つのステップを通じて、数百の GPU の効率的な並列コンピューティングを実現できます。

図 14 BLOOM モデルのトレーニングで使用される並列計算の構造

2.4 コンピューティングデバイスのメモリの最適化

現在の大規模言語モデルのトレーニングでは通常、Adam 最適化アルゴリズムが使用されます。このアルゴリズムには、各パラメーターの勾配に加えて、1 次運動量 (Momentum) と 2 次運動量 (Variance) が必要です。一般に、Adam 最適化アルゴリズムは SGD アルゴリズムよりも優れており安定していますが、コンピューティングデバイス上のメモリを大幅に消費します。

メモリ使用量を削減するために、ほとんどのシステムは混合精度トレーニング方法を採用しています。つまり、FP16 (16 ビット浮動小数点) または BF16 (Bfloat16) と FP32 (32 ビット浮動小数点) 形式の両方の値があります。。 FP32、FP16、BF16 は図 4.15 のように表されます。 FP32 では、ビット 31 は符号ビット、ビット 30 ～ 23 は指数を表すために使用され、ビット 22 ～ 0 は仮数を表すために使用されます。 FP16 では、ビット 15 は符号ビット、ビット 14 ～ 10 は指数を表すために使用され、ビット 9 ～ 9 は仮数を表すために使用されます。 BF16 では、ビット 15 は符号ビット、ビット 14 ～ 7 は指数を表すために使用され、ビット 6 ～ 0 は仮数を表すために使用されます。 FP16 は FP32 に比べて値の範囲が非常に狭いため、計算処理中にオーバーフローやアンダーフローが発生しやすくなります。 FP16 と比較すると、BF16 は精度と引き換えに、より広い値範囲を実現します。ただし、FP16 および BF16 は FP32 と比較して精度が低いため、トレーニングプロセス中に勾配の消失やモデルの不安定性が発生する可能性があります。

したがって、これらの問題を解決するには、動的損失スケーリング (Dynamic Loss Scaling) や混合精度オプティマイザー (Mixed Precision Optimizer) などのいくつかのテクノロジーを使用する必要があります。混合精度の最適化のプロセスを図 4.16 に示します。 Adam オプティマイザーの状態には、FP32 に保存されたモデルパラメーターのバックアップが含まれており、1 次運動量と 2 次運動量も FP32 形式で保存されます。モデルパラメータの数が Φ で、モデルパラメータと勾配が FP16 形式で保存されると仮定すると、合計 2Φ + 2Φ + (4Φ + 4Φ + 4Φ) = 16Φ バイトのストレージが必要です。

そのうち、アダムステータスが75％を占めます。動的損失スケーリングバックプロパゲーションの前に、損失変化 (dLoss) は手動で 2K 倍増加するため、バックプロパゲーション中に取得された活性化関数の勾配はオーバーフローしません。バックプロパゲーションの後、重み勾配は 2K 倍減少し、通常の値に戻ります。たとえば、75 億のパラメーターを含むモデルの場合、FP16 形式が使用される場合、必要なコンピューティングデバイスメモリは 15 GB だけですが、モデル状態はトレーニングフェーズ中に実際に 120 GB を消費します。

モデルの状態に加えて、コンピューティングカードが占有するメモリには、アクティベーション値 (Activation)、さまざまな一時バッファ (Buffers)、使用できないビデオメモリのフラグメント (Fragmentation) などを含む残留状態 (Residual States) もあります。アクティベーション値はチェックポインティング (Activation Checkpointing) を使用してアクティベーション値のメモリフットプリントを大幅に削減できるため、モデルの状態、特に Adam オプティマイザの状態をどのように削減するかがメモリフットプリントの問題を解決する鍵となります。

図 16 混合精度の最適化プロセス

上記は、分散機械学習システム、分散トレーニングクラスターアーキテクチャ、および分散トレーニング並列戦略の基本的な概念についての簡単な紹介です。DeepSpeed は、クラスター上で大規模な言語モデルをトレーニングする方法の例です。次の記事でそれをお知らせします。注目してサポートしてください。あなたのサポートが私の創作の原動力です。

参考内容：

(1) 30 個の大規模な言語モデルのトレーニング関連データセットの収集丨共有 - Zhihu https://zhuanlan.zhihu.com/p/612243919。

(2) 大規模言語モデルのトレーニングデータの 4 つの一般的な処理方法 - Zhihu https://zhuanlan.zhihu.com/p/673045395。

(3) 「大規模言語モデル: 理論から実践へ」Zhang Qi 他 —北京: 電子産業新聞

(4) 大規模言語モデルのレビュー - 中国人民大学 http://ai.ruc.edu.cn/research/science/20230605100.html。

クリックしてフォローし、できるだけ早くHuawei Cloudの新しいテクノロジーについて学びましょう~

理論と実践により、分散トレーニングを理解できるようになります