(fsdp2 dev.) Add training.shard_frozen_tables: frozen target tables in the root FSDP group

[yushengsu-thu]

Increment-only view of sgl-project#924 (base codex/fsdp2-backend = sgl-project#915).

Conclusion

• Size over [verifying] Add optional FSDP2 training backend sgl-project/SpecForge#915 (excluding tests and docs): +97 / −3 lines of Python (backend.py +54 shared scaffolding, fsdp2.py +15, launch.py +8, assembly.py +7, trainer.py +5, schema.py +4, disaggregated.py +4); tests +193, docs +2.
• FSDP1 → FSDP2 alone ([verifying] Add optional FSDP2 training backend sgl-project/SpecForge#915, no option): speed — EAGLE3 faster, −132 ms/step (offline e2e); DFlash2 / DSpark no change. Memory — EAGLE3 smaller, −1782 MB peak; DFlash2 / DSpark no change.
• FSDP2 + shard_frozen_tables: speed — no change: DFlash2 −1 ms/step (995 → 994 ms), DSpark +3 ms (726 → 729 ms). Memory — peak slightly larger: +595 MB (DFlash2, 37652 → 38247 MB) and +593 MB (DSpark, 25058 → 25651 MB); memory at rest smaller, −1782 MB per rank (allocated after wrap 4049 → 2267 MB at DP=4, grows with world size).
• Online e2e (disaggregated, real ShareGPT prompts), FSDP1 → FSDP2 alone: DFlash2 — speed no change (1074 → 1058 ms compute per step, -16 ms, -1.5%); memory no change (55794 → 56082 MB nvidia-smi max, +288 MB, within the granularity of reserved memory). DSpark — speed no change (720 → 720 ms compute per step, +0 ms, +0.0%); memory no change (41814 → 40840 MB nvidia-smi max, -974 MB, within the granularity of reserved memory).
• Online e2e, FSDP2 + shard_frozen_tables: DFlash2 — speed no change (1058 → 1060 ms compute per step, +2 ms, +0.2%); memory no change (56082 → 57996 MB nvidia-smi max, +1914 MB, within the granularity of reserved memory). DSpark — not run online for this option.
• EAGLE3: the option is a no-op (no composite-level frozen tables); step level +9 ms / +96 MB, within noise.
• Recommendation: keep off by default. It does not lower the training peak (the tables must be resident for the vocabulary objective and the gathered copy coexists with the shard) and costs one extra ≈2.5 GB all-gather per optimizer window; enable only when memory outside the training window matters.

Motivation

The DFlash-family wrappers carry the frozen target lm_head and embed_tokens (two 151936 × 4096 bf16 tables ≈ 1.24 GB each for Qwen3-8B) and both backends replicate them on every rank as ignored parameters. Under FSDP2 they can instead live in the root fully_shard group: sharded at rest, all-gathered once per accumulation window together with the root parameters, which keeps the DFlash fused head working (it reads lm_head.weight as an ordinary tensor inside forward). This PR adds the opt-in and measures the trade.

Stacks directly on sgl-project#915 (base branch codex/fsdp2-backend). The BackendOptions scaffolding is shared verbatim with the other FSDP2 option PRs so they can land in any order.

Modifications

• training/fsdp2.py: with shard_frozen_tables, the frozen tables are not listed in ignored_params, so they join the root FSDP2 group; everything else in _shard_model is unchanged.
• training/backend.py: BackendOptions / hooks (shared scaffolding). FSDP1 rejects the option.
• config/schema.py: training.shard_frozen_tables (requires backend: fsdp2); plumbing.
• training/disaggregated.py + assembly._backend_options(): the disaggregated runtime did not forward backend_options to its trainers, so the option was silently off in every online run until the online e2e of this series caught it; both call sites now forward it, with tests.
• tests/test_runtime/test_fsdp2_shard_frozen_tables.py: rejection/config tests, 2-GPU parity of sharded vs replicated tables (losses equal to 5 places, tables are DTensors, full state dict still carries them).

End-to-end comparison — offline (Trainer.fit)

Offline e2e = build_offline_runtime → Trainer.fit(): offline reader, FeatureDataLoader with 4 workers, TrainerController acks and logging, final checkpoint; 256 synthetic samples × 6 epochs = 48 optimizer steps, seq 2048, batch 2 × accum 4, DP=4 on 4 × B300 (torch 2.13.0+cu130); steady state = last 24 steps from the controller's log callbacks. Step-level = the real TrainerCore step on a resident batch, batch 2 × accum 8, seq 4096. Production-shaped Qwen3-8B drafts with random weights and synthetic features; deltas are per rank.

draft	fsdp → fsdp2 step time	fsdp2 → +shard_frozen_tables step time	Δ time	samples/s (4 GPUs)	peak allocated per rank	Δ peak	fit wall (48 steps)
DFlash2	978 → 995 ms (no change)	995 → 994 ms	−1 ms (no change)	32.2 → 32.2	37652 → 38247 MB	+595 MB	62.5 → 62.0 s
DSpark	740 → 726 ms (no change)	726 → 729 ms	+3 ms (no change)	44.1 → 43.9	25058 → 25651 MB	+593 MB	48.5 → 48.8 s
EAGLE3 (no-op)	796 → 664 ms (−132 ms)	664 → 966 ms	+302 ms, see note	48.2 → 33.1	21168 → 21168 MB	±0 MB	42.7 → 49.7 s

End-to-end comparison — online (disaggregated)

Online e2e = the disaggregated path exactly as specforge train runs it (managed-local): one SGLang 0.5.18 capture server with the spec-capture patch on GPU 0 (Qwen3-8B, mem_fraction_static 0.5), Mooncake store, producers feeding real ShareGPT prompts (6000 conversations, max_length 2048, qwen chat template), and a DP=3 trainer on GPUs 1–3; batch 2 × accum 4, 40 optimizer steps, num_anchors 512; steady state = the controller's perf/* counters (averaged over the ranks that logged them) over the last 24 optimizer steps. samples/s = 24 samples per optimizer step (batch 2 × accum 4 × 3 trainer GPUs) divided by the mean step time of that window; train compute and data wait are per optimizer step. GPU memory is the nvidia-smi maximum over the run for the trainer GPUs.

DFlash2

variant	steady samples/s (3 trainer GPUs)	train compute per step	Δ compute vs fsdp2	data wait per step	trainer GPU memory max	Δ memory vs fsdp2
fsdp (FSDP1)	22.22	1074 ms	+16 ms (+1.5%)	1 ms	55794 MB	-288 MB
fsdp2	22.62	1058 ms	–	1 ms	56082 MB	–
fsdp2 + shard_frozen_tables	22.59	1060 ms	+2 ms (+0.2%)	1 ms	57996 MB	+1914 MB

DSpark

variant	steady samples/s (3 trainer GPUs)	train compute per step	Δ compute vs fsdp2	data wait per step	trainer GPU memory max	Δ memory vs fsdp2
fsdp (FSDP1)	33.26	720 ms	-0 ms (-0.0%)	1 ms	41814 MB	+974 MB
fsdp2	33.26	720 ms	–	1 ms	40840 MB	–
fsdp2 + shard_frozen_tables	not run

DSpark was not run online for this option (its tables and trade-off are identical to DFlash2).

The first online matrix of this series ran with the options silently ignored: the disaggregated launch path did not forward them (fixed in this PR); the tables above are from the runs after the fix. Online memory is nvidia-smi used on the trainer GPUs, i.e. what the caching allocator reserved, a coarse measure; the precise peak-allocated numbers are in the offline table.

Notes: the "memory at rest −1782 MB" figure is the allocated memory right after wrapping (step-level harness, DP=4); everything else above is e2e. On the EAGLE3 offline e2e row: the option changes nothing in the EAGLE3 model (the frozen embed_tokens lives inside draft_model, not on the composite, and is already handled by FSDP2's frozen-parameter path); the slower run was the last of a long sequence of runs on the node and reads 67 MB of features per sample, so it needs a re-run before being attributed to anything but loader/page-cache contention.