近年、Transformerアーキテクチャに基づく大規模言語モデル（LLM）は卓越した性能により広く注目を集めています。産業規模のLLMは高品質なサービスを提供するために長いシーケンス入力を処理する必要があります。しかし、メモリ消費はシーケンス長に対して二次的に増加し、長いシーケンスの訓練の拡張性を制限しています。既存の並列手法は実行中に冗長なテンソルを生成し、メモリ最適化の余地があります。また、テンソル並列（TP）は計算と通信の有効な重畳を実現できません。これらの問題に対して、本稿ではTransformer訓練のコア計算単位（すなわち2つの連続した行列乗算およびその間に存在し得る関数演算 O=f（AB）C）に特化した汎用並列手法「メモリ効率的テンソル並列（METP）」を提案します。METPは計算Oのサブタスクを複数のデバイスに割り当て、集合通信の代わりにポイントツーポイント通信（send/recv）を用いてサブマトリクスを交換し計算を完了し、冗長なテンソルの生成を回避します。ダブルバッファリング技術により計算と通信の深い重畳を実現し、長いシーケンスTransformer訓練を導く完全重畳の理論条件を提案します。理論解析により、並列度pにおいてFlashAttentionを用いない場合の注意計算のメモリコストはO（1/p³）であり、FlashAttentionを用いて多頭自己注意を計算する場合、TPに比べ少なくとも41.7%のメモリ節約が可能であることが示されました。8台のA100 GPU環境を基にした実験により、METPは他の手法と比べてシーケンス長を2.38～2.99倍に伸ばせることが実証されました。

分散学習;大規模言語モデル;長いシーケンス;機械学習システム;メモリ効率;テンソル並列