世界を動かす技術を、日本語で。

MegaTrain: 単一GPU上での100B以上のパラメータを持つLLMのフルプレシジョントレーニング

概要

  • MegaTrain は100B以上のパラメータを持つ大規模言語モデルを 単一GPU で効率的に学習可能
  • ホストメモリ(CPUメモリ) にパラメータ・オプティマイザ状態を保存し、GPUを一時的な計算エンジンとして活用
  • CPU-GPU間帯域幅 のボトルネックを2つの最適化で克服
  • DeepSpeed ZeRO-3 と比較して高いスループットを実現
  • 7Bモデルの512kトークンコンテキスト学習 も単一GH200で可能

MegaTrain: 単一GPUによる100B+パラメータLLM学習システム

  • MegaTrain は、100B以上のパラメータを持つ大規模言語モデルを フル精度単一GPU 上で学習できる新しいメモリ中心型システム
  • 従来のGPU中心型システムと異なり、 パラメータとオプティマイザ状態をCPUメモリに保存 し、GPUは一時的な計算エンジンとして使用
  • 各レイヤーごとに パラメータをストリーミング転送 し、計算後に勾配を出力する設計
  • GPU上の永続的なデバイス状態を最小化 し、メモリ効率を向上

最適化手法

  • CPU-GPU間の帯域幅ボトルネック を回避するため、2つの主要な最適化を導入
    • パイプライン化・ダブルバッファリング実行エンジン
      • パラメータのプリフェッチ、計算、勾配のオフロードを 複数のCUDAストリームで重複実行
      • GPU実行の連続性 を確保
    • ステートレスなレイヤーテンプレート
      • 永続的なautogradグラフを廃止し、 動的に重みをバインド
      • グラフメタデータの削減 と柔軟なスケジューリングを両立

パフォーマンスと応用

  • H200 GPU(1.5TBホストメモリ)最大120Bパラメータモデル の安定学習を実現
  • 14Bモデル学習時、DeepSpeed ZeRO-3(CPUオフロード時)比で1.84倍の学習スループット
  • GH200単体で7Bモデル・512kトークンコンテキスト学習 も可能
  • メモリ効率・帯域幅活用の最適化により、 大規模モデル学習のコスト削減と柔軟性向上

対象分野・参考情報

  • 対象分野 :自然言語処理、分散・並列・クラスタ計算、オペレーティングシステム
  • arXiv論文番号 :arXiv:2604.05091 [cs.CL]
  • 論文リンクhttps://doi.org/10.48550/arXiv.2604.05091
  • 著者 :Zhengqing Yuan

Hackerたちの意見

MegaTrainはパラメータやオプティマイザの状態をホストメモリ(CPUメモリ)に保存して、GPUを一時的な計算エンジンとして扱うんだ。各レイヤーごとにパラメータをストリーミングして、勾配を計算することで、持続的なデバイス状態を最小限に抑えてる。これ、めっちゃすごいよね。家にあるのはRTX 3080の10GB VRAMだけだから、40Mや50Mパラメータ以上の大きなモデルのトレーニングには苦労してる。OOMエラーが出るし、最適化もいっぱいしなきゃいけない。PCにはもっとCPU RAMがあるから、これを使えばローカルでトレーニングできるモデルのサイズが増えるかも。

これらのアーキテクチャを最大限に活用するには、基本的に「重み」から知識や能力をもっと移動させて、ハードウェアの能力に比例した形でシステムの補完的な部分に組み込むことが鍵だと思う。ここ数ヶ月で、小さなモデルがこのAIトランスコーディングのニッチを占めるようになって、爆発的に増えてる。私が期待しているのは、ツール適応の爆発が起こること。LLMがある程度固定されたツールや問題ケースと組み合わさることで、一般性を少し犠牲にして特化した役割を果たせるようになること。トランスコーディング関連のタスクは、一般的にデバイスのユーザーが実際にやっていることと同期しているから、ユーザーのハードウェアの能力ややりたいことに密接に関連してる。だから、今のところほとんどの人はこういうことを意識してないと思う。習慣的にそうなってるんじゃないかな。今や個人用コンピュータを「孤立したハードウェア」と考えるのが意味を持つようになったから。ローカルデバイスでのMoEに対して正しいアプローチを取れば(ローカルLLMが向かっている方向)、VRAMに重みを出し入れする高コストを、ユーザーがまだ価値を感じるような極端なバッチユースケースで相殺できるかもしれない。LoRaはこれにとってすごく役立つけど、時には数層の違い以上の専門知識や特化が必要なこともある。今、これを実験中なんだ。論文と同じ基本原則だけど、技術的な最適化よりも作業負荷の最適化に近い。あと、これは文字通り機械文化の始まりだから、なんかクールだよね。

これ、めっちゃすごいよね。家にあるのはRTX 3080の10GB VRAMだけだから、40Mや50Mパラメータ以上の大きなモデルのトレーニングには苦労してる。OOMエラーが出るし、最適化もいっぱいしなきゃいけない。私も同じGPUを使ってるけど、何かをトレーニングする気になるかどうかすら不安だよ。どんなトレーニングをそのGPUでやったのか教えてくれない? :)

記事の主張は、もっと多くの計算能力とVRAMが必要だって前提になってるね。トリックで往復が少なくなるけど、完全にはなくならないし。50Mって言ったんじゃないよね?50Bのことかな?試してみるしかないけど、大きなコンテキストに期待しすぎない方がいいよ。もし彼らの技術がうまくいくなら、8096kのコンテキスト制限はまだOOMになると思う。2048ならいけるかも。これはこの論文のトリックを使わずに自分の実験から推測してるんだけどね。

モデルを何にトレーニングしてるのか聞いてもいい?そのためのトレーニングデータはどうやって生成してるの?

これ、たぶん小規模なファインチューニングの仕事にしか使われないだろうね。事前トレーニングのスケールには遅すぎる。

事前トレーニングのスケールには遅すぎる。でも「遅すぎる」っていうのは客観的な事実じゃないんだよね。どれだけの忍耐と電気代をかけられるか次第だし。AI画像生成の界隈では、モデルが画像を生成するのに5秒以上かかると文句を言う人がいる一方で、限られたハードウェアで30分待つのを喜んでる人もいる。だから「遅すぎる」って何を意味するのか判断するのは難しい。かなり主観的だよね。

モデルのコピーをRAMに複数のマシンに分散させて、各マシンがモデルの異なる部分の重みを更新して、ネットワーク越しに更新を同期させるってことだね。

これがどれくらい上手くいくのか気になってた :) もっと進められると思うけど、問題は勾配と更新がどれだけ圧縮できるかだね。

https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_... にすごく似てるね。このPyTorchのプリミティブだけでどれだけ再現できるか気になる。

H200で120Bモデルをトレーニングするのに実際どれくらい時間かかるんだろう?8台あったらどうなる?

これって別に新しいことじゃないよ。自分も似たようなことをずっとやってたけど、論文を書くのは面倒だっただけ。(: たぶん「少ないVRAMで巨大モデルをどうやってトレーニングするか?」って真剣に考える人は似たようなアイデアを思いつくと思う。でも、実際にはその分野のほとんどの人はやらないんだよね。単一のGPUで巨大モデルをトレーニングする実用性はかなり低いから。(例えば、彼らは14Bモデルを単一の3090で341トークン/秒取ってるけど、自分の方法では単一の4090で約1kトークン/秒だった。それでも遅いけど)それに、トレーニングを早くしたりVRAM使用量を減らしたりするための他のトリックもあるのに、使ってないんだよね。例えば、勾配のオフロードは必要ないし(オプティマイザーを修正すれば、勾配を直接オプティマイザーの状態に蓄積できる)、Adamの代わりにMuonを使うこともできる(これならAdamの半分のVRAMで済む)、量子化も使える(パラメータとオプティマイザーの状態両方に;例えば、Muonを4ビットに量子化したら結構うまくいった)とか。

341はあなたの1トークン/秒よりも2桁も速いから、彼らのやってることがそんなに明白じゃないように見えるね。341トークン/秒が遅いかどうかの基準もないけど、3090にしては速い感じがする。

方法を書いてくれない?読んでみたいんだけど。

「POCかGTFO」って言うけど、俺は光速旅行を発明したんだ。明らかだったけど、まだ論文は書いてないけどね :)

GPUはもう脳じゃなくて手なんだよ。脳はRAMだね。突然、奥さんが疑問に思った256GBのDDR5構成が「研究インフラ」になっちゃった。

H200 GPUで1.5TBのホストメモリ、確かに1台のGPUだけど、スリムってわけじゃないよね。

比較が128 H100の時なら、これはすごくいいアップグレードだね。それにH100やH200は時間単位でそんなに高くなく借りられるし。

RTX 3060で泣く

この技術がAppleのMシリーズみたいな統一メモリアーキテクチャでどう機能するのか、ちょっと興味あるんだ。プロセスを重ねて処理を早くしてるみたいだけど、メインメモリが統一されてるからGPUにデータを行き来させる必要がないっていうのも、何か利点があると思うんだよね。もっと賢い人、教えてくれない?

なんで誰も大きな研究所以外は手に入れられないものについて話さないの?トレーニングセットのことなんだけど。オープンなセットもいくつかあるけど、OpenAIやGoogle、Anthropicが実際に使ってるものとは全然違うと思うんだ。間違ってたら嬉しいけど。