2
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

フレームワークが動くのを待つな、FSDPを書け — リリース8日後に397B MoEを学習した基盤を全公開

2
Last updated at Posted at 2026-10-17

はじめに:フレームワークが動けば、たしかに簡単。動けば。

LLMのファインチューニングというと、verl・ms-swift・axolotl・DeepSpeedといった学習フレームワークを使うのが定石とされています。

しかし、日本の共有計算資源(Slurmクラスタ、固定されたドライバ・CUDAバージョン、root権限なし、ノード構成や台数がフレームワークの想定と違う環境)では、この定石が成立しないことが珍しくありません。リリースから時間が経った既存モデルですら、フレームワーク側の依存関係やクラスタとの相性問題を解決する作業が、学習そのものより重くなることがあります。

具体例を挙げます。私が評価した範囲では、DeepSpeedはMoEモデルで集合通信が噛み合わずハングすることがありました。MoEは構成が複雑で、ランク間で集合通信の呼び出しに少しでもズレがあると、エラーを吐かずに停止し得ます。ms-swiftは並列設定がシビアで、クラスタのノード構成と噛み合わせるのに苦労しました。さらに、私が触れた2026年初頭の版ではパラメータ凍結が用意されたパターン中心で、「特定のexpertのみ学習する」「凍結層との境界で活性化を切る」といった、私がやりたかった構成はそのままでは表現しづらく、モデルによっては学習自体のハードルが上がりました(フレームワークは進化が速いので、現状は要確認)。アーキテクチャが特殊なモデルほど、フレームワークの想定から外れていきます。

とはいえ、ハングのリスク自体は素のFSDPにもあります。違いは、ループが自分のコードなので原因を特定でき、対策を自分で書けること。実際に私が踏んだのは、データをランクに分配したときに端数バッチの数がランク間でずれ、イテレーション数の不一致でFSDPがハングするというものでしたが、これは端数バッチをdropして揃えれば解決しました(該当は1か所)。

先に断っておくと、これはチーム全体の方針ではなく、私個人の開発戦略です。私自身もms-swiftやaxolotlは使っていましたし、ms-swiftやverlを主軸にしていたメンバーもいて、フレームワークが有効に機能する場面は多くあります。つまり私が担当した一部の構成では、フレームワークより素のFSDPの方が早かった、というのが本記事の立場です。

加えて個人的な好みとして、私はLoRAや量子化での学習をあまり好みません。できるならBF16で回したい。私の経験上、そのほうが品質を読みやすいからです(もちろんLoRA/QLoRAや量子化が最適な場面は多々あります)。巨大MoEをこの好みのまま学習するには工夫が要り、それが後述するexpert選択学習(低ランク近似なしに学習対象を絞る)や、FP8で動かすモデルでも学習対象だけはBF16に戻す扱いにつながっています。選んだのは、PyTorch標準のFSDP(Fully Sharded Data Parallel)の上に学習ループを一から書くこと。一見遠回りですが、私にとっては結果的にこれが最短経路でした。本記事ではその理由、つまり素のFSDPの強みを、私が実際に構築した4つの学習基盤(いずれもリポジトリとして公開)を題材に紹介します。

構築した4つの学習基盤

基盤 モデル 手法 開発時期
qwen3-235b-fsdp Qwen3-235B-A22B (MoE) SFT / DPO / GRPO・GDPO・GSPO(attention/expert選択) 2025年12月頃に基盤完成 ←ベース
qwen-3-5-fsdp Qwen3.5-397B-A17B SFT / DPO / GRPO(β) 〜2/24(リリース2/16の8日後に学習スクリプト完成)
glm-5-fsdp GLM-5(744B、FP8運用) SFT / DPO 〜3/2(SFT完成2/24=リリース2/11の13日後、DPO完成3/2)
qwen-next-80b-fsdp Qwen3-Next-80B-A3B フルパラメータ SFT / DPO / GRPO(β) 3/4〜

いずれもH100クラスタ、Slurm運用、40〜96 GPU規模です。

1. 依存:小さく、自分で制御できる

SFT / DPO のスタックは、本質的には PyTorch + Transformers + Accelerate の3つです(リポジトリによってはflash-attn等を追加でインストールします)。学習フレームワーク特有の長大な依存ツリーがなく、conda環境の構築スクリプトは1ファイルで完結します。

最新モデルでは「何も足さなくていい」わけではありません。GLM-5のglm_moe_dsaはtransformersのdev版(git main)を要求しますし、GRPOで生成にvLLMを使う場合はvLLM nightlyが要ります。ただ決定的に違うのは、足すものが少なく・見通せて・自分で制御できる点です。入れるのはモデル自身のmodelingコードと、せいぜい数個の先端wheel。フレームワークのように「対応待ちの固定された依存ツリー全体」を相手にするのとは、詰まる確率が桁で違います。

依存の分離も自分で設計できます。たとえばGRPOでは、vLLM nightlyをNCCL重み同期のユーティリティ専用として入れ、モデルのロードには使いません。そのため学習側とvLLM側でtransformersのバージョンが食い違っても問題になりません。「一つのenvに全部」のフレームワークでは、こうした切り分けは難しい。

フレームワークが動かない原因の多くは、モデル対応の遅れではなくクラスタ環境との不整合です。だからこそ、制約の強い共有クラスタでは(既存モデルであっても)フレームワークの環境・相性デバッグのコストが、素のFSDPループを書くコストを上回ることがあります。

2. 資産化:一度書けば、235Bから3.5へ8日で移植

学習基盤を自分のコードとして持っていると、新モデルへの対応は「フレームワークの対応待ち」ではなく「自分の基盤の移植作業」になります。

実際の系譜はこうです。まずQwen3-235B-A22B向けに、FSDPベースの学習基盤(SFT / DPO / GRPO、expert選択学習、vLLM重み同期)を2025年12月頃までに作り込みました。このカスタマイズ可能なベースを起点に、その後:

  • Qwen3.5-397B-A17B(2/16リリース):Gated DeltaNetとGated Attentionのハイブリッド注意機構、512 expertのMoE、ViTエンコーダ付きという、既存フレームワークが苦手とする構成です。それでも235Bの基盤をベースに、リリース8日後の2/24にはSFT/DPO/GRPOの学習スクリプトが揃い、同日中に最初のGRPO学習が実際に回り始めています。
  • GLM-5(2/11リリース。本記事の「リリース」はいずれもHugging Faceでの重み公開日を指します):235Bベースで13日後の2/24にSFTが完成(DPOは3/2に完成)。
  • Qwen3-Next-80B-A3B:同系統のハイブリッド構成のため、ほぼ横展開で完了。

FSDPはnn.Moduleをラップする仕組みなので、移植の中心作業は「auto_wrap_policyに指定するdecoder layerクラスの差し替え」と「そのモデル固有の癖への対処」に絞られます。推論用のmodelingコードは(多くの場合transformersのmainやdev版として)リリース直後に出るため、そこを起点に着手できる。公式の「学習」スクリプトを待つ必要はないです。

そして次の新モデルが出たときも、ゼロからではなく、この基盤から差分修正するところから始められます。

補足:4つの基盤ができるまでの実際の歩み

参考までに、4つの基盤がどう積み上がってきたかを時系列で振り返ります。

  • 2025年12月9日:Qwen3-235B-A22B向けに、expert選択学習とFSDP↔vLLM重み同期を含む基盤の原型が稼働。9ノード(vLLM 1+トレーナー8)でのRLと、4ノードでのSFTをこの構成で回し始めました(公開版リポジトリではvLLM 2ノード+トレーナー8ノードの10ノード構成に整理しています)。
  • 2026年1月27日:GRPO/GDPO/GSPOの3つの目的関数を実装し、Qwen3-235Bのベースモデル・SFTモデル・派生モデルに対して比較実験。ベースモデルが誤答した問題を集めたサブセットで正答率が数%向上し、全体のスコアは維持できることを確認しました。
  • 2026年2月24日:GLM-5とQwen3.5のSFTが完成し、Qwen3.5のGRPOも完成——Qwen3.5リリースの8日後です。
  • 2026年3月2日:GLM-5のDPOが完成しました。3月4日からはQwen3-Next-80B-A3Bへの横展開に着手。

公開した4リポジトリは、この積み重ねを整理して切り出したものです。「235Bの基盤から8日で移植」という系譜は、こうして段階的に伸びてきたものです。

3. 壁がない:モデルの癖に正面から対処する

移植が速いもう一つの理由は、問題が起きたとき全コードが自分の手の中にあることです。フレームワークの抽象化レイヤーの奥で起きる問題は追跡も修正も困難ですが、自前の数ファイルなら直接手を入れられます。実例を挙げます。

FP8重みをFSDPに通す10行のmonkey-patch(GLM-5)

GLM-5にはBF16版もありますが、744Bという規模ではrank 0でモデル全体をCPUに展開する段(後述のrank 0ロード方式)でRAMに乗り切りません。そこでFP8(約760GB)で動かすことにしたところ、今度はFSDPのflat paramのパディング処理がFP8非対応という壁に当たりました。FSDP1の_construct_padding_tensorはパディングを乗算(torch.onesどうしの積)で作るため、FP8では mul_cuda not implemented for Float8_e4m3fn で落ちます。素のFSDPなら、該当関数を差し替えて解決可能です:

# 対象: PyTorch 2.x 系の FSDP1 (torch.distributed.fsdp._flat_param)。
# この _construct_padding_tensor は private API。FSDP2 には存在しないため、
# torch を上げる際は要再確認。元関数はパディングを乗算で作るため、FP8 では
# "mul_cuda not implemented for Float8_e4m3fn" で落ちる。
import torch.distributed.fsdp._flat_param as flat
_orig = flat._construct_padding_tensor

def _fp8_safe(padding_numel, dtype, requires_grad, device):
if dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
# BF16 でゼロパディングを作ってから FP8 にキャスト
return torch.zeros(padding_numel, dtype=torch.bfloat16,
device=device).to(dtype)
return _orig(padding_numel, dtype, requires_grad, device)

flat._construct_padding_tensor = _fp8_safe

さらにforward hookでFP8→BF16のブロック単位dequantをall-gather直後に挟み、重みはFP8のままシャーディングして、学習対象レイヤーだけBF16化する構成にしました。hookの登録先をFSDP unitではなくその内側のsubmoduleにする(unshard完了後にhookが走る位置を選ぶ)。こうしたFSDPの実行順序を踏まえた細工は、自前コードだからこそ可能です。

凍結層の活性化を捨てる detach hook

部分レイヤー学習では、凍結層のbackwardは不要です。凍結層→学習層の境界にforward-pre-hookでdetachし、autogradグラフを切断すると、凍結層は活性化を保持せずforwardだけを通るので、そのメモリコストはほぼゼロになります(フル学習と比べたときの差です)。なお、FSDPでwrapされた単一のforwardの途中でグラフを切るには、torch.no_grad()では届かず、この種のhookが要るという事情もあります。

ひとつ補足すると、attnのみを学習し層内のMoEを凍結する構成では、MoE出力もdetachして逆伝播を止めています。これは残差経路を通る勾配が近似になる意図的なトレードオフで、メモリ削減のために割り切ったものです(学習対象がattnに限られる前提での選択)。

MoEのexpertを選択的に学習(Qwen3-235B)

特定のexpertだけを学習対象にし、残りを凍結する構成です。FSDPはwrapping unit内でdtype統一を要求するため、凍結パラメータの扱い(FP8のGLM-5では凍結スケールのbuffer化など)に調整が要りますが、すべて自分のコード内で完結します。前述のとおり、凍結が用意されたパターン中心のフレームワークでは、この種の構成は最初から選択肢に入りにくい。detach hookも同様で、「どこを凍結し、どこでグラフを切るか」を自由に設計できることが、巨大MoEを限られたGPUで学習する鍵になりました。

GRPOのためのFSDP↔vLLMリアルタイム重み同期

RLでは生成にvLLMを使いたい。学習ノード群(FSDP)と推論ノード群(vLLM。Qwen3.5ではノード跨ぎのTP=16、235B/80BではTP=8)を分け、重み同期はTRLのWeightSyncWorkerExtension(NCCLでのbroadcastの受け口)をベースにしました。その上で、FSDP側でsummon_full_paramsをFSDP unitごとに回してフルパラメータを集約し、別クラスタのvLLMサーバーへNCCLでストリーム同期する部分を自前で実装しています。学習対象も自由に選べます。235Bでは、attention-only(FFN/expertを凍結し、対象とする層の範囲も指定可)か、マスクで選んだ(layer, expert)単位のexpertのみか、をフラグで切り替えられ、いずれもLoRAではなくフルパラメータ(BF16)で学習します(参考としてLoRA版GRPOも同梱しています)。報酬関数はLLM-as-Judge(医療精度+有用性の2段ジャッジ、reward hacking対策のゲーティング付き)。GRPO→GDPO/GSPOの目的関数の差し替えも、損失計算が自分の関数なので数十行で済みます。

※ GRPOの安定度はモデルで差があります。ベースとなった235BのGRPO(attention-only / expert選択などの学習対象切り替え・GDPO/GSPO・LLM-as-Judge)は作り込んであり安定しています。一方、それをQwen3.5 / Qwen3-Nextに移植したGRPOは、開発の最後で十分なデバッグ時間が取れず、現状β扱いです。これらの移植版で私が見ている不具合は、大きく2系統です。(1) 重みがvLLMに渡らない:同期が効かず、vLLMが初期重みのまま生成し続けるケースです。これは静かに壊れる(学習は回っているように見えるのにRLが効かない)ので、能動的な検知が要ります。学習が進んでもvLLMの生成物が変化・改善せず、学習側で再計算したlogpの乖離が同期直後でも縮まらない、という兆候で分かります(この診断は手元で行っていたもので、リポジトリには含めていません)。(2) vLLMが落ちる:別の環境ではvLLMサーバー自体がクラッシュします(原因は環境依存で、NCCLの初期化や重み同期の周辺と見ていますが未特定)。こちらはすぐ気づく失敗です。アーキテクチャ固有というより、移植時のテスト不足が主因です。

4. 可観測性:何が起きているかが全部見える

OOMやhangの原因がフレームワーク内部のどこにあるかを追う作業は消耗します。素のFSDPでは、メモリ・分散の制御点がすべてスクリプトに明示されています:

  • rank 0だけがモデルをCPUにロードし、sync_module_states=Trueで全rankへ配布(CPU RAM節約。397BクラスではこれがないとノードのRAMが足りません)
  • limit_all_gathers=Trueでall-gatherの同時実行数を制御し、ピークメモリを削減(FSDP1)
  • 端数バッチをdropして、rank間のイテレーション数不一致によるFSDP hangを防止
  • checkpoint保存時のoffload_to_cpuでGPU OOMを回避

40〜96 GPU規模を安定運用できたのは、この見通しの良さによるところが大きいです。

「全部自分で書く」のは大変では?

たしかに、素のFSDPは「全部自分で書く」ことの裏返しです。SFT/DPOスクリプトは各900〜2,000行規模になり、データローダー、sequence packing、checkpointのマージ、W&Bログまで自前です。

ただし、このコストを払うのは基盤を最初に作る場合だけです。本記事のリポジトリは、その「最初の1回」を済ませた成果物です。そのため、必要となる作業は:

  • モデルパス・データパス・Slurm設定を自分の環境に合わせる
  • 必要ならdecoder layerクラス名と学習対象の指定を差し替える

つまり、フレームワークを使う場合の「設定ファイルを書く」のと大差ない手間で、フレームワークでは表現しにくい自由度が手に入ります。

まとめ

  • 制約の強い共有クラスタでは、学習フレームワークの環境構築・相性問題が最大のボトルネックになりがち。依存が小さく自分で制御できる素のFSDPは、その点で強く、既存モデルですら「一から書く方が結局は速い」ことがある
  • 学習基盤はコードとして持てば資産になる。235Bで作った基盤から、リリース8日後のQwen3.5-397B、13日後のGLM-5(FP8)へと移植できた
  • 推論用のmodelingコードはリリース直後に出る(多くはtransformers main/dev)。FSDPはnn.Moduleをラップするだけなので、公式の「学習」スクリプトを待つ必要がない
  • FP8対応のmonkey-patch、detach hook、expert選択学習、vLLMへのNCCL重み同期。フレームワークの壁の向こうにある改造が、全部自分の手でできる

このクラスタで、このモデルを、今すぐ好きなように学習したいという方のお役に立てたら幸いです。

謝辞

この成果は、NEDO(国立研究開発法人新エネルギー・産業技術総合開発機構)の委託業務(JPNP25006)の結果得られたものです。

付録:公開リポジトリ(いずれもMITライセンス)

2
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
2
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?