GENIAC: 172B 事前学習知見, 2024

[AkihikoWatanabe]

https://zenn.dev/tokyotech_lm/articles/deb8012251bb68

[AkihikoWatanabe]

LLMの事前学習における知見がまとまっている記事とのこと

[AkihikoWatanabe]

・Megatron LMで学習 　→ 3D Parallelismなどの分散学習手法によりHF Trainerより高速 　→ Data Parallelim、Tensor Parallelism、 Pipeline Parallelismを組み合わせたもの ・GPUメンテナンス、不良で学習が継続できなかった場合はcheckpointをロードして学習 ・学習曲線が安定しているように見えるがSpikeは発生している。発生時はgradient normが急激に上昇する ・LlamaなどのLLMからの継続的事前学習ではなくfrom scratchから学習しているので透明性が高い ・Transformer engineを利用 ・AdamWを利用 ・attention dropout, hidden dropoutは0.0

この際、 通信を多く必要とする分散手法のワーカー（Tensor Parallelワーカー）はノード内に配置するようにMegatron-LMのデフォルトではなっているため、今回もそれを利用しました。このようにする理由は、ノード内の通信はNVLinkにより、ノード間通信よりも高速であるためです。また、Data Parallelの勾配平均化のための通信を考慮して、Data Parallelワーカーも可能な限りノード内に配置するMegatron-LMデフォルトの挙動を利用しました。 Pipeline Parallelismは他の並列化手法と比較して通信量が少ないP2P(Point-to-Point)通信であるため、パイプラインステージはノード間で配置するようにしました。これも、Megatron-LMデフォルトの挙動です。

勉強になる

・通常のデータ並列はoptimizer stateをworker間で複製するので遅い。Deep Speed Zero 1のように分散して保有することで高速化 ・Tensor Parallelでself attention, MLPの計算を並列化できる ・LayerNormalization, Dropoutの演算もメモリ効率の観点から並列化 ・学習を安定させるためにz-lossを利用 ・batch skippingとは、gradient clippingを行っていてもなおspikeが生じる場合に、100 step前に戻り、spikeが生じた付近のデータを数百iteration程度スキップすること