サブロウ丸

サブロウ丸

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

サーベイ: Mesh-tensorflow:Deep learning for supercomputers

@article{shazeer2018mesh,
  title={Mesh-tensorflow: Deep learning for supercomputers},
  author={Shazeer, Noam and Cheng, Youlong and Parmar, Niki and Tran, Dustin and Vaswani, Ashish and Koanantakool, Penporn and Hawkins, Peter and Lee, HyoukJoong and Hong, Mingsheng and Young, Cliff and others},
  journal={Advances in neural information processing systems},
  volume={31},
  year={2018}
}

paper: https://arxiv.org/pdf/1811.02084.pdf

論文は下記のmeshライブラリの設計思想の説明のようなものになっている。 meshライブラリはtensorflowのthird-party的な立ち位置で開発されている。

背景

  • モデルの巨大化が進行中(自然言語処理モデルGPT-3は1750億のパラメタを持つ)
  • 単一のメモリにモデルが載らないケースも

どんなもの?

  • データ並列、モデル並列を簡易に記述できる機能をtensorflowのthird-partyとして開発

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

技術や手法のキモはどこ

  • データ並列をデータテンソルのbatch方向の分割とみなし、その一般化としてデータをbatch以外の任意の方向に分割可能と捉える
  • クラスタコンピュータもmesh shape (512coreを[16, 16, 2]や[32, 2]と表現する)によるテンソルとみなし、
    • データテンソルの次元とクラスタコンピュータテンソルの次元の単射写像layoutをユーザーに定義させる(!!斬新で柔軟な発想だと思う)
    • そうすることでデータ並列もモデル並列も同じ記法で表現可能に
  • また行列演算として、分割された行列が立方的であるほど1計算あたりの通信量を減らすことができることを説明 (Appendix)

実装としては

  • tensorflowの計算グラフ内のそれぞれの演算をどのdeviceで実行するのかをユーザーが与えたmesh, layout情報をもとに自動で付与してくれる機能を提供
  • 実際の並列実行部分はtensorflow内部で行われる
  • たとえばdevice間の通信は, peer-to-peerのsend/recvで行われる
  • mesh-tensorflowとしてはより効率的な集団通信を行いたいので、内部でsend/recvによる集団通信を再実装する必要があり、実際実装されている

どうやって有効だと検証した?

  • Transformer modelでの実装例を示した
  • また10epochの実験的な実行を行なって実際に学習が実行可能なことを示した
    • 512 core TPUv2 cluster serverで実験
    • BERTを学習(バッチサイズ256, sentenceの長さ256)
      • バッチ方向16分割、モデル方向(語彙方向、feed forwardの中間層、attention head、experts)32分割
      • データは languagemodel_wiki_noref_v128k_11k (wikipediaからの50億tokenを含むテキストコーパス)
      • 最大48億パラメタモデルで13時間ほどだった(10epoch実行に)

https://github.com/tensorflow/mesh/tree/master/mesh_tensorflow/bertにあるコードからmeshとlayoutの実験設定を取り出し。

mesh_shape = [("batch", 8)]
layout = [("batch", "batch"), ("vocab", "model"), ("intermediate", "model"), ("num_heads", "model"), ("experts", "batch")]

議論はある?

  • 層ごとのlayoutのカスタマイズは無理 (なので、単調のアーキテクチャであれば問題ないが、複数のアーキテクチャを含む場合は異なるlayoutの方が効率が良くなることがありそう)
  • かなり頻繁にall-reduceやgatherが必要になるため、通信がボトルネックになりやすく効率を追求するのは難いのでは

次に読むべき論文は?

批評

Although thisapproach works well for models of sizes up to 20 billion parameterson NVIDIA DGX A100 servers (with 8 80GB-A100 GPUs), it breaksdown for larger models. Larger models need to be split acrossmultiple multi-GPU servers, which leads to two problems: (a) theall-reduce communication required for tensor parallelism needsto go through inter-server links, which are slower than the high-bandwidth NVLink [9] available within a multi-GPU server, and(b) a high degree of model parallelism can create small matrixmultiplications (GEMMs), potentially decreasing GPU utilization. (Narayanan, Deepak, et al. "Efficient large-scale language model training on GPU clusters using megatron-LM." Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 2021.)