サブロウ丸

Sabrou-mal サブロウ丸

主にプログラミングと数学

サーベイ: Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning

論文サーベイ 深層学習 分散深層学習
@article{zheng2022alpa, title={Alpa: Automating Inter-and Intra-Operator Parallelism for Distributed Deep Learning}, author={Zheng, Lianmin and Li, Zhuohan and Zhang, Hao and Zhuang, Yonghao and Chen, Zhifeng and Huang, Yanping and Wang, Yida and Xu, Yuanzhong and Zhuo, Danyang and Gonzalez, Joseph E and others}, journal={arXiv preprint arXiv:2201.12023}, year={2022} }

paper: https://arxiv.org/abs/2201.12023

目次

前半で背景や手法の概要を、後半で手法の詳細を紹介します。

背景

近年、深層機械学習分野では性能向上のためにモデルのサイズやデータセットは巨大化する傾向にあります*1。例えば自然言語処理モデルのGPT-3では100billion(1千億)以上のパラメタが使用されています。例えば1千億個のパラメタを半精度で管理するだけでも、単純計算で1490Gbyteのメモリを必要とします。

このような巨大なモデルの学習は、まずメモリ制約の点で単一のマシンだけで行うのは厳しいです。かといって複数のマシンを用いて効率的に学習させるのも簡単ではありません。効果的な並列手法はモデル構造や計算機構成に依存するため、ケースバイケースのチューニングを必要とします。これは新しいモデル開発の大きな足枷になっている、とのこと。

参考: モデルの巨大化

代表的な自然言語処理モデルのパラメタ数の推移を示した図です (billion = 10億)。 この巨大化の傾向は今も続いています。図は"Efficient large-scale language model training on gpu clusters using megatron-lm."のFigure1を引用。

事前知識

事前知識: 計算グラフ

計算グラフとは、モデルにおける計算順序をグラフ表現したものです。図ではResnet風のモデルの計算グラフを図示。青色の四角が演算子ノード、緑色のノードでモデルのパラメタを表現しています。例えばこの計算グラフではFの演算子でP1とXのテンソル積が行われます。

事前知識: 分散深層学習

分散深層学習とは複数台のマシンを用いて分散*2に機械学習を行うこと全般を指します。研究としては大きく二つの方向性に分けることができて、

  • ひとつの機械学習モデルの学習ワークロードを分散化することで、高速化の実現やハードウェアのメモリ要求を緩和するもの*3。
  • それぞれのマシンで異なるモデルを学習する。計算の最中ではモデル情報を交換する。マシンごとに使用できるデータが異なる場合に有用*4。

以降では前者の文脈における分散深層学習の説明を行います。

事前知識: 分散手法の分類

次に基本的な分散手法について紹介。そのために機械学習のワークフローから確認しましょう。まず入力データx_iに対するモデルfの出力値を算出します。いくつかの出力を求めたら、入力データに紐づく正解値を元に損失値を算出し、損失値が小さくなるようにモデルパラメタを更新します。

これに対し基本的な分散手法として次のようなものがあります。

  • (a) データ並列
    • 入力から出力を求める計算（順伝播）を分割
  • (b) モデル並列
    • モデルの計算要素を分割
  • (c) テンソル並列
    • モデル内のテンソル計算を入力、もしくはパラメタ（あるいはその両方）を分割

事前知識: モデル並列とpipeline

モデル並列について、もう少し詳細を見てみましょう。モデル並列ではモデルパラメタを分割してそれぞれのマシンで管理できるため要求されるメモリ量が減少します。下記のように4Stageにモデルを分割する例では管理するべきメモリ量は1台で計算する場合と比べて単純計算で1/4になります。

分割されたモデルでの順伝播/逆伝播のワークロードにはpipeline処理が用いられます。pipeline処理ではデータをモデルに入力する際に、そのデータをさらに細かくしたものを順にStageに流す、ということを行います。これにより後方のStageの待機時間が減ることで処理時間の短縮が見込まれるわけです。図では入力をa~hの8つのmicrobatchに分割したときのpipeline処理を示しています。図は"Alpa: Automating Inter-and Intra-Operator Parallelism for Distributed Deep Learning"のFigure5より引用。

論文概要

Alpaとはどんなもの?

モデル分割+並列計算ワークロード作成手法(Alpa)をこの論文で提案しています。

  • 入力
    • モデルの計算グラフ
    • 2次元meshで表現した計算リソース(N行M列で同性能のXPUを並べたもの)
  • 出力
    • 計算グラフの分割とそれぞれに割り当てる計算リソース
    • テンソル演算ごとの並列手法
  • 目的
    • 1単位のパイプライン処理*5の高速化
出力イメージ

計算グラフの分割と計算リソース割り当て例。図はResnet(風)計算グラフで、P、Q、Rはモデルパラメタ。例えば左上の演算子(F)ではP1とXのテンソル積(P1 X)と非線形変換が行われます。計算リソース割り当ては長方形型*6で行われます。

さらにStage内の演算（基本的にはテンソル積）において、どのテンソルのどの軸を基準に並列処理をするかを決めます。これは単一の演算子の並列処理自体の効率だけではなく、隣り合う演算子が採用する並列手法同士の相性として、発生する通信コストも考慮して決定されます。

先行研究と比べてどこがすごい?

並列戦略をチューニングする研究がいくつか提案されていて、また、モデルの自動分割に関する研究もありますが、それらはモデル並列単体（すなわち、データ並列との組み合わせは考えない）であったり分散計算機やモデルに強く依存しているため、その適用範囲が限られています。

その上で、この論文の提案の利点・新規性を挙げてみました。

  • データ並列など、さまざまな並列方法とモデル並列を組み合わせたワークロードを作成できる
  • 計算リソースの割り当ての自由度が高い、Stageごとに異なる数や形のXPU群を割り当てられる
  • 計算機割り当て時に、メモリ容量の制約を反映できる*7。すなわち、Stage計算時に計算機のメモリが不足しないことが（理論上）確約できる
  • 並列ワークロード策定時間が小さい

技術や手法のキモはどこ?

  1. 計算グラフの演算子をどのように並列化するか
  2. 計算グラフをどこで分割してどれほどの計算リソースを割り当てるか

の最適化を2段階で行います

これらを(1)を整数線形計画法(ILP)、(2)を動的計画法(DP)*8で定式化して最適化。個人的にはDPの部分で、2次元mesh状の計算リソースの割り当てについて、割り当てる形*9を考慮しなくても良いような（ある程度妥当そうな）仮定を設けている部分が非自明でなるほどなぁという感じ。

次に読むべき論文は?

同じ著者らの論文。通信に特化して議論しているらしい。

  • On Optimizing the Communication of Model Parallelism." arXiv e-prints (2022): arXiv-2211.

自動モデル分割手法として

  • [17] Dapple: A pipelined data parallel approach for training large models, 2021.
  • [38] Pipedream: generalized pipeline parallelism for dnn training, 2019.
  • [55] Supporting very large models using automatic dataflow graph partitioning, 2019.

数値実験

実装
  • Python(16K LoC)とC++(6K Loc)で実装
  • フロントエンドとしてJax, バックエンドにXLA*10を使用
  • 分散実行時にはdevice mesh workerとしてRayを使用
  • 実行時、計算にはXLA、通信にはNCCLを使用

https://raw.githubusercontent.com/google/jax/main/images/jax_logo_250px.png https://github.com/tensorflow/tensorflow/raw/master/tensorflow/compiler/xla/g3doc/images/xlalogo.png https://github.com/ray-project/ray/raw/master/doc/source/images/ray_header_logo.png

モデル

最大700億パラメタのモデルを用いて検証。Throughput(FLOPS)で評価。 比較対象はIntra(データ、テンソル)並列のみと、Inter(モデル)並列のみ、に加えて表に記載している代表的な並列手法と比較しています。

モデル 比較手法
GPT-3 Megatron-LM v2
Gshared MoE DeepSpeed
Wide-ResNet PP-DP(独自手法)

Wide-ResNetは目立った比較手法がないため、pipelineとDPのみを行う手法と比較しています。

計算環境

計算リソースは8node, 64GPUs (8 GPU/node)でAmazon EC2 p3.16xlarge instanceを用いています。

  • 8 NVIDIA V100 16GB GPUs
  • 64 vCPUs
  • 488 GB memory
Case Study

Wide-ResNet on 4(a) and 8(b) GPUsでの計算結果例。batch分割→batch分割では通信の発生なし。それ以外の分割が行われているのはメモリ制限の影響と思われます。

Throughput結果

  • GPTモデルでは、Transformer-basedモデルに特化して開発されているMegatron-LMのモデル分割モデルと比べて競合する性能を達成
  • Gshard MoEモデルでは手動でチューニングされたDeepSpeedモデルを比較して2ノードだと3.5×、4ノードだと9.7×のスピードアップを達成
  • Wide-ResNetでも高い線型スケーリング性を発揮
並列ワークロード策定時間

図は390億パラメタのGPT-3モデルで計算リソース64GPUの条件における、コンパイル = 並列ワークロード策定時間を示しています。1時間かかっていなくてすごい。



並列計算ワークロード策定手法の詳細

Alpaモデル分割の概要

  1. 定義したモデルの計算グラフを構築*11
  2. その計算グラフを分割してStageと呼ぶ、まとまりを作る。このときに、各々のStageにどれほどの計算リソースを割り当てるのか、も同時に最適化
  3. 2の内部では、"Stage内の演算を、割り当てられた計算リソースを用いたときの最適な並列処理"を求める問題を求解

モデル並列とpipeline

分割されたモデルでの順伝播/逆伝播のワークロードにはpipeline処理が行われます。pipeline処理では、データをモデルに入力する際に、そのデータをさらに細かくしたものを順にStageに流す、ということを行います。これにより後方のStageの待機時間を減らすことで処理時間の短縮が見込まれるわけです。図は入力をa~hの8つのmicrobatchに分割したときのpipeline処理を示しています。

Stageごとの処理時間は自身の計算量と計算リソースにより差が生じていることに注目してください。このときデータがすべてのStageを流れ終わるまでの時間は

と数式表現できます。ここでt_iはi番目のStageの一つのmicrobatchの処理時間、Bはmicrobatchの個数*12です。この値が小さいほど優れた分割とリソース分配であると言えます。

計算グラフをどこで分割してどれほどの計算リソースを割り当てるか

具体的な分割には動的計画法を利用。例えば与えられた計算グラフの一部をえいやとStage Aとして分割し、また与えられた計算リソースからえいやと2つを割り当てるとします。すると、残りの計算グラフと残りの計算リソースを与えられた場合の最適なモデル分割というよりサイズの小さい問題の解が求められれば、このStage Aを含む"モデル分割の良さ"の評価を行える、ということに気づきます。

なので、Stage Aとして考えられる計算グラフの分割を全列挙して、その全てに対し、残りの計算グラフと残りの計算リソースを与えられた場合の最適なモデル分割を求めてモデル分割の評価を行えば、最適なStage Aは何か?を知ることができるのです。

あえて数式化するならこんな感じ。ここでDは計算グラフGの可能な分割方法の集合で、sはStage Aに割り当てるリソース数、Cost(A, s)はStage Aにs個の計算リソースを割り当てたときの最適な処理時間です。この動的計画法を解くことで最適な計算グラフ分割と計算リソースの割り当てを近似的に求めることができます*13。

この動的計画法(DP)では計算リソースを"数"で管理していますが、本当は二次元メッシュ状の計算リソースにどのような形で配置するか、を考えなければならなりません。が、そうしなくても良いようにStageに割り当てる計算リソースは$(1, 2k), (k, M)$の形状の長方形のみを許容するようにします。図はNGな割り当て例。この条件さえ満たせば、それぞれのStageに割り当てた長方形リソースのピースを必ず元々のメッシュに重なりなく敷き詰められることが証明できます。

演算ノードの並列化

上で述べた動的計画法を実行するにはCost(A, s), すなわちStage Aを計算リソースsを用いて計算するときの最適な計算時間を求める必要があります。

もう少し噛み砕くと、Stage Aが持つ計算グラフの演算子を計算リソースsを用いてどのように並列計算すれば良いか？を考えることになります。 が、テンソル演算Y = PXの計算ひとつをとっても、複数の並列計算が考えられるのです。

  • [Y_1, Y_2] = P[X_1, X_2]
  • [Y_1; Y_2] = [P_1; P_2]X
  • [Y_{11}, Y_{12}; Y_{21}, Y_{22}] = [P_1; P_2][X_1, X_2]

などなど... この演算単体の計算時間は計算リソースに依存します。

この部分を少し詳細に見てみましょう。まず Y = PXという行列演算がモデルにおける基本的な演算になります。ここでXがモデルへの入力で、Pがモデルパラメタです。

ここで入力の次元を一つ増やして3次元とみなす。一般的に多次元の行列はテンソルと呼ばれます*14。行われる演算は行列の場合と同じ。

機械学習における複数の入力をモデルに入力するというワークフローを思いだしましょう。すなわちモデルパラメタPは固定で入力X = [X_1, ..., X_B]についてY = [PX_1, ..., PX_B]の出力が得られるように計算が行われます。このとき入力XをバッチB方向に次元を足したテンソルとみなすことができます。例えばPytorch機械学習フレームワークでは入力Xは[B, H_1, w]のテンソルとして実装します。

図はデータ並列の様子を表したものです。入力がバッチ方向に3分割され、異なるXPUが保有します。そしてそれぞれのXPUでテンソル演算を行い結果を保有する。このとき演算は完全に独立しているためXPU間で通信は発生しません。モデルパラメタPは全てのXPUが保有する必要があります。

図はテンソル並列の様子を表しています。この場合はモデルパラメタPが分割されてそれぞれのXPUで管理されます。このとき2回テンソル並列が行われた結果、それぞれのXPUが[B, H_2, w]のテンソルを出力として保有しています。さらに3回目のテンソル並列（P = [H_4, H_3]）を図のように行うことにすると、入力にあたるテンソルを全てのXPUで同期する必要があります。その同期のためにはAll-Reduce通信を用いてXPUが持つテンソルの要素を足し合わせ、それを全てのXPUに配置する、ということを行う必要があります。

理解のためにXPU1で行われる演算を記載します。[A, B]でA行B列の行列、A dot Bで行列A, Bの行列積を表現する。すると[H_2/3, H_1] dot [H_1, w] = [H_2/3, w] → [H_3, H_2/3] dot [H_2/3, w] = [H_3, w]という演算がバッチごとに行われます。

all-reduce通信では次のように情報の共有が行われます。

#1         #2         #3
[1a]       [2a]       [3a]
[1b]       [2b]       [3b]
[1c]       [2c]       [3c]
--------- All Reduce ----------
[1a+2a+3a] [1a+2a+3a] [1a+2a+3a]
[1b+2b+3b] [1b+2b+3b] [1b+2b+3b]
[1c+2c+3c] [1c+2c+3c] [1c+2c+3c]

仮に3番目のテンソル演算を上図のように行う場合は、全てのXPUで図のように入力テンソルを保有する必要があります。それを行うにはReduce-Scatterによる通信を使えば効率的に実現できます。

#1         #2         #3
[1a]       [2a]       [3a]
[1b]       [2b]       [3b]
[1c]       [2c]       [3c]
------- Reduce Scatter --------
[1a+2a+3a] []         []
[]         [1b+2b+3b] []
[]         []         [1c+2c+3c]

ここでの説明をまとめると、

モデルパラメタの管理 演算間の通信
データ並列 全てのXPUで同じものを保有 なし
テンソル並列 分割して管理できる あり

とそれぞれ、メモリ使用面と通信コストについてトレードオフを持つことがわかります。また、並列方法によって計算効率も異なることに注意してください。このことから、計算グラフ(Stage) 内における演算子の並列手法は、演算子ごとにみるだけではダメで、隣り合う演算子で採用する並列手法をみないと通信コストを考慮できないし、演算子全体をみないとメモリ使用量を考慮できない、ということが分かります。

ここで、計算グラフ全体をみたときにそれぞれの演算子がどの並列手法を採用すれば効率が良くなるのか？の問題を解くために次で紹介する整数線形計画法が用いられています。

整数線形計画法による定式化

Stageに含まれる演算ノード v \in Vについて、考えられる並列方針を s_v = [s_{v1}, s_{v2}, ..., s_{vk_v} のk_v次元ベクトルで表します。このとき、それぞれの要素は0か1のみの値をとり、さらに全ての要素の中でただ一つだけが1の値を取るようにします。このときs_{vi} = 1のときは演算ノードvではi番目の並列方針が採用されることを意味します。 またパラメータとしてi番目の並列方針を採用したときの計算コストをd_{vi}、通信コストをc_{vi}とします。最後に、演算ノードがvの出力が演算ノードuに入力されるとき、演算ノードvでs_{vi}、uでs_{uj}を採用した場合に演算ノード間で必要になる通信（再配置）コストを(i, j)成分に持つ行列をR_{vu}とします。

  • s_v: 演算vの並列方針
  • c_v: 演算vの通信コスト
  • d_v: 演算vの計算コスト
  • R_{vu}$: (i, j)成分は演算vが並列方針s_{vi}、演算uが並列方針s_{uj}を採用したときの、演算間で発生する通信（再配置）コスト

前述する記号を置くと、Stageにおけるコストは、計算グラフG = (V, E)について次の式で表せられます。

制約は

この問題は整数線形計画ソルバを用いて解くことができます。このとき、c, d, Rといったデータは事前にプロファイルを行うことで求めておく必要があります。また並列方針も与えておく必要があり、ここが大変だったろうなぁと推察します。

実際のコードはここ。

その他の工夫

  • 計算グラフの簡素化
    • モデル分割を行う動的計画法の前処理として実施
    • あらかじめ隣接する演算子のいくつかをまとめておくことで、計算量を減らす
  • Cross-mesh resharding
    • ノード間の通信をなるべくノード内の通信に置き換える

議論はあるか？

  • Stage間通信コスト
    • モデル分割の際に、Stage間の通信コストは考慮されていません。ただ、Stage間の通信量は小さいので、微々たる問題ではあるかも
  • XPU（Accelerator)性能、通信スピードの均一性の仮定
  • どの局所的メッシュを取ってきても、XPUの性能や通信スピードは同一的であるとみなす仮定は妥当でしょうか？また、動的計画法では、あるStageに割り当てる計算リソースは数字で管理しているが、それを実際にmesh状にマッピングするときにノード間を跨ぐような解も出てくるのでは？
  • 計算と通信のオーバーラップの精密化
    • いくつかの実験でメモリ不足が発生しています。おそらく計算と通信を別スレッドで実行するようにしているが、そことの連携は取っていないので、計算が終わった側から計算結果をどんどん次のStageに送って、そこでメモリ不足を引き起こしているのでは？

*1:これは古典的な機械学習の常識とは反していて、モデルの巨大化=パラメタ数が大きいほど過学習しやすく、学習済みモデルの汎用性が下がる、というデメリットがニューラルネットワークには当てはまらない、とされています。またパラメタ数を増やすと一旦性能は悪化するのですが、さらに増やし続けていくと性能向上がさらにみられるという二重降下と呼ばれる減少が知られています。そのためモデルの巨大化は現段階でも有効に働いています。 https://qiita.com/Uchiiita/items/0e2f1f1653e0c186de6d

*2:一般的に複数台のマシンを用いて並列に処理を行うことを分散処理、という。

*3:モデルの学習結果は、"メモリが無限の理想的なマシンがあるとして、そのマシン1台で行う学習の結果と一致する"、のが理想であり、実際それを満足する手法もあるが、効率化のためにそれを犠牲にするものもある。

*4:例えば医療データなどプライバシー性が高く秘匿性のあるデータは他のマシンに送信できない、という状況を想定。

*5:明記されていないがGpipe ("GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism"

*6:具体的には$(1, 2k)$か$(k, M)$、$M$は二次元メッシュの列数

*7:Alpaの出力結果を見ると計算機のメモリ容量に余裕がある場合はデータ並列が多用され（そのStageを計算するすべてのXPUでモデルパラメタを複製）、そうでない場合はテンソル並列（XPUでモデルパラメタを分割管理）が多用されるようなワークフローになっている

*8:1: モデル分割を動的計画法で最適化するのは常套手段ではある。 "PipeDream: generalized pipeline parallelism for DNN training.", "Automatic graph partitioning for very large-scale deep learning."

*9:例えば4個のXPUをあるStageに割り当てるとして、その配置方法はテトリスのピースの形だけ候補がある。が、本論文では$(1, 2k)$と$(k, M)$の長方形の形に限定して割り当てを行っている。ここでMは2次元meshの列数。

*10:XLA（Accelerated Linear Algebra）は、線形代数のためのドメイン固有のコンパイラ

*11:モデル中に分岐がある場合は入力データによって計算グラフが変わることがあるが、今回はそのようなケースは考えず計算グラフは常に一定とする。

*12:この図だとB=8

*13:例えばStage間の通信コストは無視しているため、近似的にしかこの分割問題を解くことはできない。

*14:n次元よりもn階のテンソルという言い方をする。0階のテンソルはスカラー（数値）で、1階のテンソルはベクトル、2階のテンソルが行列と、特別に名前があります