技術 約16分で読めます

MegaTrainはシングルGPUで120BパラメータLLMをフル精度訓練する

LLM 機械学習 GPU DeepSpeed メモリ最適化
いけさん目次

LLMの訓練にはGPUクラスタが必要、という常識を正面から覆す論文がarXivに投稿された。
Notre Dame大学とLehigh大学の研究チームが発表したMegaTrainは、シングルGPUで120BパラメータのLLMをフル精度（BF16/FP32混合、量子化なし）で訓練するシステムだ。
GitHubでコードも公開されている。

このブログでもLoRA学習環境の構築や非同期RL訓練の設計パターンを扱ってきたが、どちらもGPU VRAMの制約が大前提にあった。
MegaTrainはその前提自体を取り払う。
VRAMに収まらないなら量子化かLoRAで妥協する。その二択にフル精度訓練という第三の選択肢を追加した。

LLM訓練のメモリを分解する

なぜGPU VRAMがボトルネックになるのか。
訓練時のメモリ消費を構成要素に分解すると、パラメータ本体よりオプティマイザ状態のほうがはるかに大きいことがわかる。

BF16混合精度訓練 + Adamオプティマイザの場合、1パラメータあたりのメモリ消費は以下の通り。

構成要素データ型バイト/パラメータ内容
パラメータ（重み）BF162モデルの重み本体
勾配BF162バックワードパスで計算される勾配
Adamの一次モーメント（m）FP324勾配の指数移動平均
Adamの二次モーメント（v）FP324勾配の二乗の指数移動平均
マスターウェイトFP324更新時に必要なFP32コピー
合計16

Adamは各パラメータに対して勾配の移動平均（一次モーメント）と勾配の二乗の移動平均（二次モーメント）を保持し、学習率を適応的に調整するオプティマイザだ。
この2つのFP32テンソルだけで8バイト/パラメータを消費する。
さらにBF16で計算した重みをFP32のマスターウェイトに書き戻す必要があるため、合計16バイト/パラメータになる。

これをモデルサイズごとに算出すると以下の通り（アクティベーション除く）。

モデルパラメータ数必要メモリ（16B/param）H200 VRAM (141GB)A100 VRAM (40/80GB)
Qwen2.5-7B7.6B約122 GBぎりぎりOOM
Qwen2.5-14B14.7B約235 GBOOMOOM
Qwen2.5-32B32.5B約520 GBOOMOOM
Llama-70B70B約1,120 GBOOMOOM
GPT-OSS-120B120B約1,920 GBOOMOOM

7Bですら、アクティベーションを加えるとH200の141GBでは厳しい。
アクティベーションはフォワードパスの中間計算結果で、バックワードパスの勾配計算に必要になる。
14B以上はどのGPU単体でもVRAMに収まらない。

一方、CPUメモリなら1サーバーで1.5TB〜2TBを搭載できる。
120Bの1,920GBですら、2TBのサーバーなら収容可能だ。

ここにアクティベーションの問題もある。
アクティベーションはバッチサイズとシーケンス長に比例して増大し、数十GBから数百GBに達する。
Activation Checkpointing（中間結果を保存せず必要時に再計算する手法）を使えば、メモリ消費を計算時間とトレードオフで削減できる。
MegaTrainはレイヤー単位の実行により、アクティベーションの常駐を1レイヤー分に自然に抑えている。

既存アプローチの限界: ZeROからFSDPまで

このメモリ問題に対して、これまで複数のアプローチが提案されてきた。
代表的なのがMicrosoftのDeepSpeed ZeROシリーズとPyTorchのFSDP（Fully Sharded Data Parallelism）だ。

ZeROの3段階

ZeRO（Zero Redundancy Optimizer）は「複数GPUで同じデータを冗長に持つ無駄を排除する」という設計思想で、3つのステージに分かれる。

ステージ分割対象GPU間で冗長排除メモリ削減率
ZeRO-1オプティマイザ状態Adamのm, v最大4倍
ZeRO-2+ 勾配勾配テンソル最大8倍
ZeRO-3+ パラメータ重み本体も分割最大N倍（GPU数依存）

ZeRO-1はAdamの2つのモーメントだけをGPU間で分割する。
各GPUは全パラメータの重みと勾配を持つが、オプティマイザ状態は自分の担当分だけ。
ZeRO-2は勾配の分割を加え、ZeRO-3はパラメータ本体も分割する。

ZeRO-3なら理論上N台のGPUで1/Nのメモリで済むが、計算のたびに他のGPUからパラメータをall-gather（全GPU間で分散データを集約する通信操作）で集める必要があり、通信コストが増大する。

ZeRO-OffloadとZeRO-Infinity

ZeRO-3でもGPUの台数が足りなければメモリは不足する。
そこでCPUメモリやNVMeストレージを「あふれ先」として使うのがZeRO-OffloadとZeRO-Infinityだ。

graph TD
    subgraph "ZeRO-Offload / ZeRO-Infinity"
        GPU["GPU VRAM<br/>一次記憶（主）"]
        CPU["CPU メモリ<br/>スピル先（従）"]
        NVMe["NVMe SSD<br/>さらなるスピル先"]
        GPU -->|あふれ| CPU
        CPU -->|あふれ| NVMe
    end
    subgraph "MegaTrain"
        CPU2["CPU メモリ<br/>主記憶<br/>全パラメータ常駐"]
        GPU2["GPU VRAM<br/>計算装置<br/>一時バッファのみ"]
        CPU2 -->|"レイヤー単位ストリーミング"| GPU2
        GPU2 -->|"勾配退避"| CPU2
    end

ZeRO-Offloadはオプティマイザ更新をCPUで行い、フォワード/バックワードはGPUで実行する。
ZeRO-InfinityはZeRO-3にNVMeオフロードを加えた最大構成で、理論上は無限のモデルサイズに対応する。

しかし両方とも「GPUメモリが一次記憶」という設計は変わらない。
GPUに入りきらないものをCPUに追い出す設計では、PCIeバスを経由するデータ転送がランダムアクセス的になり、帯域を効率的に使えない。
結果として、モデルサイズが大きくなるほどスループットが急激に落ちる。

後述のベンチマークで、14BモデルでMegaTrainがZeRO-3の12.2倍のスループットを出した背景にはこの構造的な問題がある。

FSDP

PyTorch公式のFSDP（Fully Sharded Data Parallelism）はZeRO-3と同等の思想をPyTorchネイティブに実装したもの。
パラメータ、勾配、オプティマイザ状態をすべてGPU間でシャーディング（分割配置）する。
ZeRO-3と同じく、計算時に必要なパラメータをall-gatherで集めるため、GPU間通信がボトルネックになる構造は同じだ。

MegaTrainの設計: CPUメモリを主記憶にする

MegaTrainはこの主従関係をひっくり返す。
パラメータとオプティマイザ状態のすべてをCPUメモリに置き、GPUは「一時的な計算装置」としてのみ使う。

設計の中核は3つの要素で構成される。

  1. パイプライン化されたダブルバッファリング（CPUとGPU間のデータ転送を計算に隠蔽）
  2. Stateless Layer Templates（PyTorchの計算グラフを排除してメモリオーバーヘッドを削減）
  3. レイヤー単位のストリーミング実行（アクティベーションの常駐を1レイヤーに抑制）
graph TD
    subgraph "CPU メモリ（1.5〜2TB）"
        P["全パラメータ<br/>BF16"]
        O["オプティマイザ状態<br/>FP32 m, v"]
        M["マスターウェイト<br/>FP32"]
    end
    subgraph "GPU VRAM（バッファのみ）"
        B0["バッファ0"]
        B1["バッファ1"]
        CS["ComputeStream"]
    end
    P -->|"WeightStream<br/>レイヤーi+1転送"| B1
    B0 -->|"計算実行"| CS
    CS -->|"GradientStream<br/>勾配退避"| O
    O -->|"CPU側で<br/>パラメータ更新"| P

GPUのVRAM上に永続的に残るデータはほぼゼロ。
ダブルバッファ2面分のメモリだけで、どんなサイズのモデルでも訓練できる。
後述のスケーリング実験では、GPU側のメモリ割り当てをわずか3.83GBに固定したまま、7.6Bから43Bまでスケールしている。

パイプライン化されたダブルバッファリング

CPUからGPUへのデータ転送がボトルネックにならないのか。
この疑問に対するMegaTrainの回答が、3本のCUDAストリームを使ったパイプライン化されたダブルバッファリングだ。

CUDAストリームとは、GPU上で非同期に実行されるコマンドキューのこと。
異なるストリームに投入された処理は並行に実行できる。
MegaTrainはこれを利用して、計算・転送・退避を同時に走らせる。

CUDAストリーム役割方向
ComputeStreamフォワード/バックワード計算GPU内
WeightStreamパラメータ転送CPU → GPU
GradientStream勾配退避GPU → CPU

GPU上にはバッファが2面あり、ピンポン方式で交互に使う。

graph LR
    subgraph "時刻 t"
        A["バッファ0<br/>レイヤーi 計算中"] --> B["バッファ1<br/>レイヤーi+1 転送中"]
    end
    subgraph "時刻 t+1"
        C["バッファ0<br/>勾配退避 + i+2転送"] --> D["バッファ1<br/>レイヤーi+1 計算中"]
    end
    B --> D
    A --> C

ComputeStreamがバッファ0でレイヤーiの計算をしている間に、WeightStreamがバッファ1にレイヤーi+1のパラメータを転送する。
計算が終わったらGradientStreamがバッファ0の勾配をCPUへ退避し、次の計算はバッファ1で始まる。
この繰り返しで、GPUが「計算待ちのアイドル状態」になる時間を最小化する。

帯域は足りるのか: ハードウェア別の転送性能

このパイプラインが成立する条件は「パラメータの転送時間が、直前レイヤーの計算時間に完全に隠れること」だ。
Transformer 1レイヤーの計算時間が数十ms〜数百msであるのに対して、実際のハードウェアでどの程度の転送時間になるかを見てみる。

ハードウェアCPU-GPU接続帯域（片方向）1レイヤーのパラメータ（7B、約600MB）転送時間
GH200NVLink-C2C450 GB/s約600 MB約1.3 ms
H200PCIe Gen5 x1664 GB/s約600 MB約9.4 ms
A100 PCIePCIe Gen4 x1632 GB/s約600 MB約18.8 ms
RTX 3090PCIe Gen4 x1632 GB/s約600 MB約18.8 ms

GH200のNVLink-C2Cは450GB/s（片方向）と桁違いに速く、転送はほぼ瞬時に隠蔽される。
PCIe Gen4のA100やRTX 3090でも、バッチサイズが大きければ1レイヤーの計算時間は数百msに達するため、18.8msの転送は十分に隠蔽可能だ。

ただしバッチサイズが極端に小さい場合や、モデルが小さくて計算が軽い場合は転送時間が計算時間を上回る可能性がある。
MegaTrainのアブレーション実験によると、ダブルバッファリングを無効化するとスループットが31.3%低下する（266.3 → 182.91 TFLOPS）。
帯域が狭い環境ほどこの差は広がるだろう。

GH200は通常のGPUサーバーとは異なる統合アーキテクチャを採用している。
Grace CPU（Arm）とHopper GPUが同一パッケージに搭載され、CPUとGPUの間がNVLink-C2C（Chip-to-Chip）で直結されている。

通常のPCIe接続では、CPU → PCIeスイッチ → GPUという経路を経由し、各段でレイテンシが加算される。
NVLink-C2CはCPUとGPUの物理的な距離を最短にし、900GB/s（双方向合計）の帯域を実現する。
PCIe Gen5の128GB/s（双方向合計）と比較して約7倍だ。

MegaTrainの「CPUが主、GPUが従」設計は、このGH200アーキテクチャと特に相性がいい。
CPU-GPU間の転送がボトルネックにならないので、GPUの計算能力をほぼフルに使い切れる。

Stateless Layer Templates

PyTorchの通常のautograd（自動微分）では、フォワードパスで計算グラフが構築され、バックワードパスまでメタデータが保持される。
計算グラフにはテンソル間の依存関係、各オペレーションの入出力のshape、バックワード関数へのポインタなどが含まれる。
モデルが大きくなるほど、このグラフ自体のメモリ消費が無視できなくなる。

従来のPyTorchモデルとの違い

通常のPyTorchモデルでは、nn.Moduleが重みテンソルをnn.Parameterとして直接保持する。

# 通常のPyTorch（概念図）
Layer1.weight → Tensor(GPU)   # 永続的にGPUに常駐
Layer1.bias   → Tensor(GPU)
Layer2.weight → Tensor(GPU)
...
Layer80.weight → Tensor(GPU)  # 全レイヤーが同時にGPUメモリを消費

80レイヤーあれば80レイヤー分のパラメータが同時にGPU上のメモリを占有する。

MegaTrainのStateless Layer Templateは、AttentionブロックとMLPブロックの計算ロジック（CUDAカーネル）をテンプレートとして保持するが、重みへのポインタは持たない。

# MegaTrainのStateless Template（概念図）
TemplateA.compute_logic → CUDAカーネル  # 計算ロジックのみ
TemplateA.weight_slot   → None          # パラメータは未バインド

TemplateB.compute_logic → CUDAカーネル
TemplateB.weight_slot   → None

# 実行時に動的バインド
TemplateA.Bind(Buffer0) → レイヤー1を実行
TemplateB.Bind(Buffer1) → レイヤー2を実行（並行で転送済み）
TemplateA.Bind(Buffer0) → レイヤー3を実行（バッファ0を再利用）

パラメータがストリーミングで到着したタイミングでBindプリミティブを呼び、テンプレートの入力スロットにバッファのビューを動的にマッピングする。
テンプレートAでレイヤー1が実行されている間に、テンプレートBにレイヤー2のパラメータがバインドされる。

この設計には2つのメリットがある。

  • 永続的な計算グラフが不要になり、メタデータのオーバーヘッドが消える
  • GPU上のパラメータがバッファ2面分で固定されるため、モデルサイズに依存しないメモリ消費になる

MoEモデルの訓練: GPT-OSS-120Bの場合

ベンチマークに使われたGPT-OSS-120BはMoE（Mixture of Experts）モデルで、128個のエキスパートを持つ。
MoEは入力トークンごとにルーターネットワークが最適なエキスパートを選択し、全エキスパートではなく一部（通常2〜8個）だけを活性化する設計だ。
パラメータ数が120Bでも、1トークンあたりの計算に使われるのはごく一部のエキスパートだけなので、実効的な計算量はDenseモデルより少ない。

MoEモデルの訓練における特有の課題は、エキスパートの数だけパラメータが膨れ上がることだ。
120Bパラメータの大半は128個のエキスパートの重みが占め、共有レイヤー（Attention、ルーターなど）は比較的小さい。
128エキスパートの全パラメータを同時にGPU VRAMに載せるのは、現行のどのGPU単体でも不可能に近い。

MegaTrainのレイヤー単位ストリーミングはMoEの各エキスパートに対しても同じように機能する。
ルーターがエキスパートを選択した後、選択されたエキスパートのパラメータだけをCPUからGPUに転送して計算し、勾配をCPUに戻す。
未選択のエキスパートはCPUメモリに眠ったままで、GPU VRAMを消費しない。

推論側で同じ発想を実装したのがHypuraのNVMeエキスパートストリーミングだ。
Hypuraはエキスパートの重みをNVMe SSDからオンデマンドでGPU（Metal）にストリーミングし、LRUキャッシュで99.5%のヒット率を実現する。
MegaTrainは訓練側でCPUメモリからストリーミングする。
どちらもGPU VRAMにすべてを載せず、必要なパラメータだけをオンデマンドで送る設計だ。

ベンチマーク結果

実験にはQwen2.5シリーズ（7B/14B/32B/72B）とGPT-OSS-120B（MoE、128エキスパート）が使われた。
評価データセットはMetaMathQA（約39.5万件の英語数学問題）。

GH200での性能

モデルTFLOPS備考
Qwen2.5-7B284
Qwen2.5-14B264ZeRO-3 Offloadの1.84倍
Qwen2.5-32B250+

デバイスメモリの割り当てをわずか3.83GBに固定したまま、レイヤー数を28（7.6Bパラメータ）から180（43Bパラメータ）まで増やすスケーリング実験も行われた。

56レイヤーの時点でFSDPは43 TFLOPSまで落ちるが、MegaTrainは264 TFLOPSを維持し、6.14倍の差がつく。
FSDPはレイヤー数が増えるほどall-gatherの通信オーバーヘッドが蓄積して性能が劣化するのに対し、MegaTrainはレイヤー数に関係なくストリーミングのパイプラインが一定のスループットを維持する。

A100 PCIe（40GB）での性能

A100はGH200やH200と比べてPCIe Gen4接続で帯域が狭いが、それでもMegaTrainの優位は圧倒的だった。

モデルMegaTrainGeminiZeRO-3
7B128 TFLOPS52.8（2.42倍）36.0（3.56倍）
14B122 TFLOPS15.0（8.13倍）10.0（12.20倍）
32B114 TFLOPSOOMOOM

14BモデルでDeepSpeed ZeRO-3の12.2倍のスループット。
32Bではベースラインがメモリ不足で動かない中、MegaTrainは114 TFLOPSで訓練を続行した。

Geminiは「GPU-CPU間でCPU側のオプティマイザ更新をオーバーラップさせる」アプローチで、ZeRO-Offloadと同系統の設計。
MegaTrainとの差が大きい理由は、Geminiがオプティマイザ更新のオーバーラップにとどまるのに対し、MegaTrainはフォワード/バックワード計算自体をパイプライン化している点にある。

コンシューマGPUでの実験

GPUモデルバッチサイズTFLOPS
RTX A6000 (48GB)14B956.82
RTX 3090 (24GB)14B330.19
RTX A60007B1255.73
RTX 30907B535.09

RTX 3090のVRAM 24GBで14Bモデルのフル精度訓練が動く。
従来ならLoRAかQLoRAで妥協するしかなかったサイズだ。
LoRAは重み行列を低ランク近似で効率的に更新するパラメータ効率化手法で、QLoRAはさらにベースモデルを4bit量子化して省メモリ化する。

精度の検証

フル精度訓練を謳うなら、精度が落ちていないことを示す必要がある。
MetaMathQAでの結果は以下の通り。

モデルZeRO-3 OffloadZeRO-InfinityMegaTrain
7B88.93%88.97%88.99%
14B92.41%92.36%92.52%

MegaTrainのほうがわずかに高い。
パラメータのストリーミングやバッファリングが数値精度に影響しないことが確認された。
これは当然の結果で、MegaTrainは量子化も近似も行わず、同じBF16/FP32のデータをCPU-GPU間でコピーしているだけだからだ。

超長コンテキスト訓練（GH200）

Qwen2.5-7Bでコンテキスト長を1Kから512Kまで伸ばす実験もある。

コンテキストバッチサイズステップ時間TFLOPSデバイスメモリ
1K15827.05秒284.774.2 GB
8K2027.3秒283.274.5 GB
64K255.3秒331.377.1 GB
512K1871.4秒407.481.9 GB

Attention計算はコンテキスト長の2乗に比例して増えるが、レイヤー単位の実行によりアクティベーションの常駐を1レイヤー分に抑えている。
512Kでデバイスメモリが81.9GBに収まっているのは、GH200の統合メモリアーキテクチャを活用した結果だ。

TFLOPSが1Kの284.7から512Kの407.4に上がっているのは意外に見えるかもしれない。
これはコンテキスト長が伸びるとAttention計算（Q×K^Tの行列積）の比率が増え、Tensor Coreの稼働率が上がるため。
短いコンテキストでは通信やバッファ操作の比率が相対的に高い。

フル精度訓練 vs パラメータ効率化手法

MegaTrainの登場で「フル精度訓練」という選択肢が現実的になったが、LoRAやQLoRAが不要になるわけではない。
両者は目的が異なる。

観点フル精度訓練（MegaTrain）LoRA / QLoRA
更新対象全パラメータ低ランクアダプタのみ（全体の0.1〜1%）
精度最高（量子化なし）ベースモデルの量子化でわずかに劣化あり
スループット全パラメータの転送が必要更新パラメータが少なく高速
ハードウェア要件大量のCPUメモリ（1〜2TB）少ないGPU VRAMで動作
向いている用途プリトレーニング、ドメイン適応タスク特化ファインチューニング

LoRA学習環境の構築で扱ったように、LoRAはMac mini M4 Proの24GBユニファイドメモリでも動く手軽さがある。
一方、LoRAは全パラメータの一部しか更新しないため、ベースモデルの知識構造を大きく変えるような訓練（新しいドメインへの適応やプリトレーニングの継続）には向かない。

MegaTrainが埋めるのはこのギャップだ。
「フル精度で全パラメータを更新したいが、GPUクラスタは持っていない」という需要に、コンシューマGPU + 大量CPUメモリという構成で応える。

「シングルGPU」の実態と今後の展開

Hacker Newsのコメントでは「H200＋1.5TBホストメモリをシングルGPUと呼ぶのは」という声もあった。
たしかにGH200＋大容量メモリという構成は一般的ではないし、14Bモデルのスループットが341 tokens/秒という数字も、推論特化の最適化と比べれば遅い。

ただ、この研究の価値は速度競争にはない。

1つは、既存フレームワークの構造的な非効率を可視化したこと。
ZeRO-3 Offloadが14Bモデルで10 TFLOPSしか出せない一方、同じA100でMegaTrainが122 TFLOPS出せるのは、12倍の性能差が「ハードウェアの限界」ではなく「ソフトウェア設計の問題」だったことを意味する。
今後DeepSpeedやFSDPがこの設計を取り込んでもおかしくない。

もう1つは、メモリ階層の設計思想そのものの転換だ。
HypuraのNVMeストリーミングは推論側でSSDをメモリ階層に組み込んだ。
非同期RL訓練は生成と訓練を分離してGPU利用率を95%以上に引き上げた。
MegaTrainは訓練側でCPUメモリを主記憶にする。

論文では2つの拡張方向に言及している。
1つはマルチGPUへの展開で、テンソル並列やエキスパート並列との組み合わせ。
もう1つはSSDを記憶階層に追加するティアードストレージ（階層化ストレージ）で、1兆パラメータの訓練を視野に入れている。
シングルGPUで120Bが動くなら、4〜8GPU構成で1Tに届く可能性は十分にある。