はじめに

大規模言語モデルの知識蒸留では、教師モデルのlogitsを事前計算して保存する手法が有効です。学習時にGPU上に教師モデルをロードする必要がなくなり、VRAM制約の厳しい環境でも大規模な生徒モデルを学習できます。

しかし、教師モデルが32B規模のMoE（Mixture of Experts）の場合、RTX 3090（24GB VRAM）では事前計算のスループットが約1,060 tok/sに留まり、200Mトークン分の処理に48時間を要していました。GPU利用率はわずか37%。GPUが6割以上遊んでいる。

本記事では、9つの仮説を順に検証し、最終的にスループットを**約2.65倍（2,809 tok/s）**まで高速化した過程を記録します。推測ではなく計測で原因を特定した、失敗込みの実践記録です。

環境

GPU: NVIDIA RTX 3090 (VRAM 24GB)
OS: WSL2 Ubuntu on Windows
PC RAM: 64GB DDR4
教師モデル: llm-jp-4-32b-a3b-base（32B MoE, アクティブ3B, 4bit量子化）
フレームワーク: PyTorch 2.6 + transformers + bitsandbytes
目的: Top-128 logitsの事前計算・保存

問題: GPU利用率37%、1,060 tok/s

教師モデルをbitsandbytesで4bit量子化してGPU上にロードし、MB=14（バッチサイズ14）で事前計算を実行したところ、以下の状態でした。

指標	値
スループット	1,060 tok/s
GPU利用率（SM）	37%
メモリコントローラ	16%
VRAM使用量	21.0 GB
CPU利用率	6%（ほぼアイドル）
PCIe RX	50 MB/s
RAM	余裕あり（スワップなし）

GPUもCPUもRAMも余裕がある。なのにスループットが上がらない。MBを16以上に上げるとCUDA driver error: device not readyで停止してしまう。何が原因なのか。

仮説と検証: 9つの失敗

ここから計測ベースで原因を絞り込みます。結論を先に言えば、仮説1〜9のうち最初の9つは全て的外れでした。

仮説1: ネットワークI/Oが律速

CulturaXデータセットをストリーミングモードで読み込んでいるため、毎回ダウンロードが発生しているのではないか。

# 5秒間のダウンロード速度を計測
RX1=$(cat /proc/net/dev | grep eth0 | awk '{print $2}')
sleep 5
RX2=$(cat /proc/net/dev | grep eth0 | awk '{print $2}')
echo "Download speed: $(( (RX2 - RX1) / 5 / 1024 / 1024 )) MB/s"

結果: 0 MB/s。データはローカルキャッシュ済みでした。❌ ネットワークは原因ではない。

※AIの回答： [MoEの律速です]

仮説2: tokenize速度が律速

CulturaXのテキストをtokenizeする処理がボトルネックではないか。

tokenize単独のベンチマークを実行したところ、64,797 tok/s。事前計算の1,060 tok/sの約60倍速い。❌ tokenizeは原因ではない。

※AIの回答： [MoEの律速です]

仮説3: RAM不足・スワップ

WSL2のRAMが不足してスワップが発生し、GPU待ちが起きているのではないか。

free -h
vmstat 1 5

結果: Swap使用量=0B、si/so=0/0、Available=十分。❌ RAMは原因ではない。

AIの回答： [MoEの律速です]

仮説4: DataLoaderのtensor生成が遅い

以前の学習スクリプトで「DataLoaderが88%のボトルネック」になった経験があります（torch.tensor()を毎サンプル呼ぶとPython→C++のブリッジコストが積み重なる問題）。同じパターンではないか。

data_streamをバルクバッファ方式に変更し、数千チャンクをNumPy配列で一括保持→torch.from_numpy()で一回だけtensor変換する方式に変更しました。

結果: 速度変化なし（1,065 tok/s）。❌ DataLoaderは原因ではない。

※AIの回答： [MoEの律速です]

仮説5: np.savez_compressedが遅い

shard保存時のgzip圧縮がCPUを占有し、GPUが待たされているのではないか。

np.savez_compressedをnp.savez（非圧縮）に変更。

結果: 速度変化なし。❌ 保存処理は原因ではない。

※AIの回答： [MoEの律速です]

仮説6: CPU→GPU転送の非効率

pin_memory()とnon_blocking=TrueでPCIe転送を最適化すれば改善するのではないか。

bx = torch.stack(batch_x).pin_memory().to(DEVICE, non_blocking=True)

結果: 速度変化なし。❌ CPU→GPU転送は原因ではない。

※AIの回答： [MoEの律速です]

仮説7: WSL2のメモリ割り当てが過大

.wslconfigでmemory=48GBに設定されており、64GBのうち48GBをWSL2が確保。Windows側のRAMが5.6GBしか残らず、GPUドライバに影響しているのではないか。

memory=24GBに修正してWSL再起動。

結果: 速度変化なし（1,060 tok/s）。❌ WSL2のメモリ設定は原因ではない。

※AIの回答： [MoEの律速です]

仮説8: offload層がfp16のまま（部分的成功）

offload_tmpディレクトリを確認したところ、16GBのfp16ファイルが存在。MoE Expert層の後半9層分が4bit量子化されずにfp16のままディスクにオフロードされていました。

GPU上: layer 0-22（4bit化済み）→ 17.4GB
ディスク: layer 23-31（fp16のまま）→ 16GB
合計: 33.4GB（4bitなら18GB程度のはず）

device_map="auto"は、fp16サイズ（64GB）で配置計画を立てるため、24GBに収まらない分をオフロードします。今回の環境では、オフロードされた層に4bit量子化が適用されていませんでした。

device_map={"": 0}で全層をGPUに強制配置したところ、全層が4bit化され、教師ロード後のVRAMは18.4GBに収まりました。

しかし、速度は変わらない。1,060 tok/sのまま。ディスクI/Oは解消されたはずなのに、なぜ？

⚠️ 部分的成功。全層4bit化はVRAMを改善したが、スループットは向上せず。

※AIの回答： [MoEの律速です]

仮説9: MBを上げればGPU利用率が上がる

全層4bit化でVRAMに余裕ができたので、MB=16, 18, 20, 22を試す。

結果: 全てCUDA driver error: device not readyで失敗。

warmup（1回目のforward）は成功するが、2回目以降で必ずlm_headのF.linearでクラッシュ。MBの値に関係なく、MB=14でも連続forwardで不安定。

File ".../modeling_qwen3_moe.py", line 1060, in forward
    logits = self.lm_head(hidden_states[:, slice_indices, :])
RuntimeError: CUDA driver error: device not ready

❌ 単純なMB増加では解決しない。

※AIの回答： [MoEの律速です]

転機: エラーの共通点

9つの仮説がすべて失敗した時点で、エラーのスタックトレースを改めて見直しました。

全てのクラッシュに共通しているのは、lm_headのF.linearで発生していること。これはMoEのforward（Attention + Expert routing）とは別の、最終段の線形変換です。

真の原因: lm_headの出力テンソルが巨大すぎる

このモデルの語彙数は196,608です。lm_headの出力テンソルは以下のサイズになります。

MB × SEQ_LEN × vocab × dtype
= 14 × 512 × 196,608 × 2 bytes (BF16)
≈ 2.8 GB

教師モデル本体が18.4GBでVRAMに載っている状態で、MoE forwardの中間テンソルと合わせると、lm_headの出力テンソルが加わった瞬間に24GBを超過していました。

MB	lm_head出力	合計推定	結果
14	2.8 GB	〜22 GB	ギリギリ動く
16	3.2 GB	〜23 GB	TDR/OOM
20	4.0 GB	〜24 GB	TDR/OOM
22	4.4 GB	〜25 GB	確実にOOM

つまりMBの上限を決めていたのはMoEのforward速度ではなく、lm_headの出力テンソルサイズでした。GPU利用率37%の原因は、MBが14に制限されていたことによるバッチ不足だったのです。

※AIの回答： [すごい発見です]

発想: なぜlm_headを分割しようと思ったか

エラーが常にlm_headで発生していることに気づいた時点で、lm_headの出力テンソルサイズを計算しました。MB=14で2.8GB。これがVRAMの残りをほぼ食い尽くしていると分かりました。

ここで考えたのは、「全語彙のlogitsをバッチ全体で一度に保持する必要があるのか？」という点です。

事前計算の目的はTop-128のlogitsを保存することであり、196,608語彙すべてのlogitsを同時にVRAM上に保持する必要はありません。1サンプル分のlogits（200MB）さえあればTop-128を抽出でき、抽出後は即座に解放できます。

つまり、問題は「lm_headが遅い」ことではなく、「lm_headの出力を溜めすぎている」ことでした。処理を分割すれば、計算量は変わらずにVRAMピークだけを削減できます。

この発想は、以前の学習スクリプトでgradient checkpointingを導入した経験と似ています。gradient checkpointingも「計算量は増えるがVRAMを削減する」トレードオフであり、VRAM制約下では計算速度よりもメモリ管理が優先されるという原則が共通しています。

解決: lm_headをチャンク分割

196,608語彙のlogitsをバッチ全体で一括生成する必要はありません。保存する必要があるのはTop-128ですが、計算自体は全語彙に対して行われます。今回の最適化は計算量の削減ではなく、VRAMピーク削減によって大きなMBを利用可能にした点が本質です。

lm_headを1サンプルずつ適用し、Top-128を即座に取得してから巨大テンソルを解放する方式に変更しました。

# 旧: 一括（ピークVRAM = MB × 512 × 196608 × 2B = 数GB）
logits = teacher(input_ids=bx).logits          # 巨大テンソル生成
top = torch.topk(logits, 128, dim=-1)

# 新: チャンク分割（ピークVRAM ≈ 200MB）
outputs = teacher.model(input_ids=bx)          # lm_head手前まで
hidden = outputs[0]                            # (MB, 512, hidden_dim)

for i in range(MB):
    logit_i = teacher.lm_head(hidden[i:i+1])   # 1サンプル分のみ
    top_i = torch.topk(logit_i, 128, dim=-1)   # Top-128即座取得
    # → logit_iを即座に解放、VRAMは200MBのまま

この変更により、lm_headのピークVRAMは1 × 512 × 196,608 × 2B ≈ 200MBに抑えられます。2.8GB → 200MB、約14分の1です。

結果: MB=56、2,809 tok/s

lm_headのVRAM制約が解消されたことで、MBを大幅に引き上げることが可能になりました。

MB	tok/s	VRAM_peak	倍率
14（旧・一括lm_head）	1,060	21.0 GB	1.00x
20	1,520	19.9 GB	1.43x
24	1,696	20.4 GB	1.60x
28	1,920	20.8 GB	1.81x
32	2,088	21.1 GB	1.97x
36	2,161	21.6 GB	2.04x
40	2,373	21.8 GB	2.24x
48	2,595	22.6 GB	2.45x
56	2,809	23.3 GB	2.65x

GPU利用率も37%から67%に向上しました。

200Mトークン分の事前計算時間は、48時間から約18時間に短縮されました。

なぜMB増加で速度が上がるのか

MoEモデルのforward処理には、Expert routingや4bit逆量子化など、バッチサイズに依存しない固定オーバーヘッドがあります。MBが小さいと、このオーバーヘッドの比率が大きくなり、GPUの演算ユニットが遊びます。

MBを上げることで、1回のforwardあたりの有効計算量が増え、GPU利用率が向上します。旧構成ではlm_headの出力テンソルがVRAMを圧迫してMB=14が上限でしたが、チャンク分割によりこの制約が解消されました。

失敗から学んだこと

AIの回答は「間違い」というより、過去の大量事例から見た確率の高い候補でした。
AIの回答を鵜呑みにしない事！

9つの仮説検証を振り返ると、失敗には共通のパターンがありました。

#	仮説	結果	反省
1	ネットワークI/O	❌	ストリーミング＝毎回DLとは限らない
2	tokenize速度	❌	64K tok/sは十分すぎた
3	RAM不足	❌	計測すれば一瞬で分かった
4	DataLoaderのtensor生成	❌	学習時の経験が転用できると思い込んだ
5	圧縮保存	❌	先入観で疑った
6	CPU→GPU転送	❌	PCIe帯域は十分だった
7	WSL2メモリ設定	❌	Windows側の問題を疑ったが無関係
8	offload層のfp16	⚠️	全層4bit化は成功したが速度に効かず
9	MB増加	❌	lm_headのVRAMスパイクで全滅

最大の教訓は、「GPUが遅い＝演算が遅い」ではないということです。今回の原因は出力テンソルのVRAM占有であり、演算性能ともメモリ帯域とも関係ありませんでした。エラーのスタックトレースが最初から答えを示していたのに、別の仮説を先に検証してしまったのは反省点です。

補足: 全層4bit化

仮説8で発見した問題ですが、独立した知見として記録します。

今回の環境では、device_map="auto"による配置では一部のMoE Expert層がGPU外（ディスク）に配置され、期待した4bit化によるVRAM削減効果が得られていませんでした。

GPU上: layer 0-22（4bit化済み）→ 17.4GB
ディスク: layer 23-31（fp16のまま）→ 16GB

device_map={"": 0}で全層をGPUに強制配置することで、全層が4bit化され、VRAM=18.4GBに収まりました。offloadなしの構成が可能になり、ディスクI/Oのオーバーヘッドも解消されています。

この知見は、lm_headチャンク分割と組み合わせることで初めてスループット改善に結びつきました。全層4bit化だけではMB=14のままで速度は変わらず、チャンク分割でMBの上限が解放されて初めて効果が現れました。

実装上の注意点

GPU→CPU転送の効率化

チャンク分割ではlm_headをMB回呼び出すため、GPU→CPU転送も同回数発生します。転送効率を上げるため、GPU上でTop-Kの結果を蓄積してから一括転送する方式を採用しました。

all_indices = []
all_values = []
for i in range(MB):
    logit_i = teacher.lm_head(hidden[i:i+1])
    top_i = torch.topk(logit_i, 128, dim=-1)
    all_indices.append(top_i.indices)     # GPU上に蓄積
    all_values.append(top_i.values)       # GPU上に蓄積
    del logit_i, top_i

# 一括CPU転送（PCIe効率向上）
indices = torch.cat(all_indices, dim=0).cpu().to(torch.int32).numpy()
values = torch.cat(all_values, dim=0).cpu().to(torch.float16).numpy()

adaptive MBの禁止

過去の経験から、実行中にバッチサイズを動的に変更するとCUDAコンテキストが破壊されることが分かっています。MB=56固定で運用しています。

TDR設定

WindowsのTDR（Timeout Detection and Recovery）がデフォルト（8秒）だと、大きなバッチのforward中にGPUがリセットされる場合があります。レジストリでTdrDelayとTdrDdiDelayを60秒に設定しています。

まとめ

項目	旧構成	新構成
device_map	"auto" + offload 16GB fp16	{"": 0} 全層4bit
lm_head処理	一括（MB×512×196608）	チャンク分割（1サンプルずつ）
MB	14	56
tok/s	1,060	2,809（2.65倍）
GPU利用率	37%	67%
VRAM peak	21.0 GB	23.3 GB

今回の構成では、主要な制約要因はMoEブロックのforward演算ではなく、lm_headが生成する大規模logitテンソルによるVRAM圧迫でした。

大語彙モデル（vocab > 100K）でlogitsを事前計算する場合、全語彙のlogitsをバッチ全体で一括生成するのではなく、1サンプルずつlm_headを適用してTop-Kを即座に抽出する方が、結果として大きなバッチサイズが使え、スループットが向上します。

この最適化は、教師モデルのアーキテクチャや量子化方式に依存しないため、他のMoEモデルや大語彙モデルにも適用可能と考えられます。

より一般的な知見として、大語彙モデルの蒸留では、教師モデルの演算性能よりもlogit保持方式がGPU利用効率を決める場合があるということが分かりました。教師モデルのforward速度をいくら最適化しても、出力テンソルがVRAMを圧迫してバッチサイズを上げられなければ、GPUは遊んだままです。