@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]と表現する)によるテンソルとみなし、
- また行列演算として、分割された行列が立方的であるほど1計算あたりの通信量を減らすことができることを説明 (Appendix)
実装としては
- tensorflowの計算グラフ内のそれぞれの演算をどのdeviceで実行するのかをユーザーが与えたmesh, layout情報をもとに自動で付与してくれる機能を提供
- 実際の並列実行部分はtensorflow内部で行われる
- たとえばdevice間の通信は, peer-to-peerのsend/recvで行われる
- mesh-tensorflowとしてはより効率的な集団通信を行いたいので、内部でsend/recvによる集団通信を再実装する必要があり、実際実装されている
どうやって有効だと検証した？
- Transformer modelでの実装例を示した
- また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.)
