# 1兆 (1T) パラメータ規模のLLMの事前学習検証 ## 概要 [[Preferred Elements]](PFE)が2024年7月から約1か月、[[Preferred Networks]]のPLaMo系ソフトウェアを基盤に1兆(1T)パラメータ級のLLM事前学習を検証した技術記事である。[[GENIAC]]の助成事業として、[[NEDO]]の計算資源提供支援を受けて実施された。 目的は、[[Mixture-of-Experts]](MoE)によって計算量を大きく増やさずモデル容量を1T級へ拡張できるかを、実際の学習で確かめることだった。評価軸は学習速度とモデル性能である。 > [!key-insight] 実証されたこと > 400基のH100 GPU上で1Tパラメータ級MoEを致命的なtrain loss発散なしに学習できた。一方、学習速度はPLaMo-100Bの半分以下で、通常の学習型ルーティングではエキスパート崩壊が起きた。 ## 事前学習設定 - PLaMo-100Bを基盤に、1Tパラメータでモデル並列しやすいようレイヤ数などを調整 - MoEによりMLP部分のパラメータを12倍へ拡張 - [[DeepSeek-AI]]のDeepSeek MoEを参考にshared expertを導入 - PLaMo-100Bの学習後半とほぼ同じデータセットを使い、学習量は100Bトークン - 400基の[[NVIDIA H100]] GPUを使用 100Bトークンは近年のLLM事前学習としては短い。著者は、1T級の実現可能性を限られた期間で検証するため、学習量を抑えて試行錯誤に時間を割いた。 ## メモリ消費量 400基のH100は合計32TBのGPUメモリを持つ。AdamWはパラメータ1個あたり12バイトの状態を必要とするため、1Tパラメータではオプティマイザ状態だけで約12TB、総メモリの37.5%を占める。 そこで[[FP8-LM]]と8-bit Optimizersを参考にオプティマイザ状態を低精度で保持し、使用量を約4〜5TBへ削減した。これにより活性化保存など、他の一時領域へメモリを配分できるようになった。 ## MoEの学習安定化 ### 学習型ルーティングの崩壊 当初は[[Switch Transformer-C]]の系譜にある一般的なMoEを採用した。小さなRouterが各トークンをTransformerのMLP/FFNに相当するエキスパートへ割り当て、Routerとエキスパートの重みを同時に最適化する構成である。 ![[_attachments/pretraining-1t/fig01-moe-routing.png]] （図1。Routerがトークンを複数のFFNエキスパートへ割り当てるMoEの構造。） 1TモデルではRouterが学習に失敗し、ほぼ常に同じエキスパートへトークンを割り当てた。load balancing lossは均等割り当てを促すが、この規模では崩壊を防げず、実際に使われるパラメータが減少した。容量を超えたトークンのtoken droppingも発生し、train lossは約7.8で下げ止まった。 ![[_attachments/pretraining-1t/fig02-training-loss-collapse.png]] （図2。通常の1T-MoEはtrain loss約7.8で停滞し、Hash Layersは約1.57まで低下した。） 小規模な約30Bパラメータ実験では問題が起きなかったため、原因候補として、Transformerの深さと、エキスパート単体で簡単な問題を解けるほどエキスパート当たりのパラメータが大きいことを挙げている。詳細な原因究明には小さな再現ケースが必要だとする。 ![[_attachments/pretraining-1t/fig03-balanced-routing.png]] （図3。理想的なルーティングではトークンが各エキスパートへ均等に割り当てられる。） ![[_attachments/pretraining-1t/fig04-routing-collapse-token-dropping.png]] （図4。1T-MoEで特定エキスパートへの集中とtoken droppingが起き、未処理トークンが学習を妨げる。） ### Hash Layersへの変更 問題を回避するため、[[Hash Layers]]へ移行した。Hash Layersは学習したRouterではなくハッシュ関数でエキスパートを決めるため、Routerの正のフィードバックによる崩壊を避けられる。 Hash Layersを使った学習では最終train lossが約1.57となり、通常のMoEより大幅に低かった。 ![[_attachments/pretraining-1t/fig05-hash-layer-learning-curve.png]] （図5。Hash Layersは通常の1T-MoEより低いtrain lossへ到達した。） 代替案として、load balancing lossの係数を上げる、2048エキスパートのようにより疎なモデルにする、[[Expert Choice Routing]]を使う、という選択肢も検討された。係数を上げる方法は効果がなく、残る2つは学習速度などの理由で採用されなかった。 ## 学習速度 1Tモデルの実効効率は約220 TFLOP/s/GPUで、PLaMo-100Bの約540 TFLOP/s/GPUの半分以下だった。MoEはFLOPsを増やさずパラメータを増やせるが、エキスパートの分散配置により通常のTransformerよりメモリアクセス量が多く、メモリ帯域が律速になりやすい。 メモリアクセス/FLOP比を下げる一般的な方法はバッチサイズを増やすことだが、巨大モデルでは活性化保存のためメモリ制約が厳しい。再計算とモデル並列の最適化が不十分で、今回の検証ではバッチサイズを十分に増やせなかった。 ## モデル性能 同じ100Bトークン条件で1B、3B、13Bモデルを学習し、train lossを比較した。小規模モデルの外挿では1Tモデルのlossはさらに低いと予測されるが、実測値は予測より悪かった。 ![[_attachments/pretraining-1t/fig06-train-loss-scaling.png]] （図6。1B・3B・13Bのtrain lossからの外挿に対して、1Tモデルのtrain lossは高い位置にある。） 原因候補は、1Tモデルに対する100Bトークンの不足と、MoEでパラメータを増やすためのアーキテクチャ設計不足である。前者については学習後半でloss低下が落ち着いているため、単純な学習不足とは断定できない。Expert Choice RoutingでMoEのスケーリング則を改善した先行研究との類似性から、Hash Layersでも大規模MoE特有の現象が起きた可能性を示唆する。 downstream taskなどの評価は、アーキテクチャ変更で状況が変わりうるため優先されなかった。 ## 横断的知見 - **条件付き計算は容量拡張と実効効率を自動的には両立しない**: [[Mixture-of-Experts]]は活性化パラメータを抑えながら総パラメータを増やすが、本記事の220 TFLOP/s/GPUはPLaMo-100Bの540 TFLOP/s/GPUを下回る。FLOPs削減だけでなく、メモリアクセス、バッチサイズ、モデル並列を同時に最適化する必要がある。(Source: [[@2024__Preferred Networks__1兆 (1T) パラメータ規模のLLMの事前学習検証]], [[LLM分散学習]]) - **MoEの負荷分散はモデル規模で相転移しうる**: 約30Bでは機能した学習型Routerが1Tで崩壊し、load balancing lossの係数変更でも改善しなかった。これは既存の[[負荷分散]]が補助損失・ゲーティング関数・システム配置という複数層へ展開してきた理由を、規模拡大の失敗事例から裏づける。(Source: [[@2024__Preferred Networks__1兆 (1T) パラメータ規模のLLMの事前学習検証]], [[負荷分散]]) - **Hash Layersは「学習しないルーティング」という独立した解法軸を示す**: 既存のSwitch TransformerやExpert Choice RoutingがRouterの学習を設計し直すのに対し、Hash Layersは割り当てをハッシュ関数へ固定する。精度・専門化・通信効率を犠牲にせず大規模化できるかは未検証だが、ルーティング崩壊を避けるために学習可能性そのものを外す設計が成立することを示す。(Source: [[@2024__Preferred Networks__1兆 (1T) パラメータ規模のLLMの事前学習検証]], [[Hash Layers]], [[Expert Choice Routing]]) - **スケーリング則の外挿にはアーキテクチャの同一性が必要である**: 本記事では小型の通常モデルから1TのMoEモデルへ外挿したtrain lossが外れた。データ量不足だけでなく、密モデルからMoEへの条件付き計算の変更が外挿の前提を壊した可能性がある。([[LLMスケーリング則]]) ## 未解決の問い - 1T級でRouterが崩壊する閾値は、モデルの深さ・エキスパート当たりのパラメータ数・エキスパート数のどの組み合わせで決まるか。 - Hash Layersは学習型Routerと比べて、エキスパートの専門化・未知データへの汎化・通信負荷でどのような差を持つか。 - Expert Choice Routing、Hash Layers、load balancing lossの係数調整を、同じ1T級モデル・同じ計算予算で比較した場合の精度と実効効率はどうなるか。 - 1Tモデルで100Bトークンが不足していたのか、MoEアーキテクチャの設計が不適切だったのかを、十分な小規模スケーリング実験で切り分けられるか。 - 低精度オプティマイザ状態が、メモリ削減と引き換えに学習安定性・最終性能へ与える影響はどの程度か。 - 1T級MoEで、バッチサイズ・再計算・モデル並列を共同最適化したとき、220 TFLOP/s/GPUをどこまで改善できるか。 ## 関連 - エンティティ: [[Preferred Networks]] / [[Preferred Elements]] / [[Hiroaki Mikami]] / [[PLaMo-100B]] / [[NVIDIA H100]] / [[GENIAC]] / [[NEDO]] / [[Switch Transformer-C]] - 概念: [[Mixture-of-Experts]] / [[Hash Layers]] / [[Expert Choice Routing]] / [[負荷分散]] / [[条件付き計算]] / [[LLM分散学習]] / [[LLMスケーリング則]] / [[言語モデル事前学習]] / [[混合精度訓練]] ## 出典 - [[.raw/articles/pretraining-1t-2024-08-20.md]] - [1兆 (1T) パラメータ規模のLLMの事前学習検証｜株式会社Preferred Networks](https://www.preferred.jp/ja/blog/tech/pretraining-1t)