[fsdp, recipe] feat: support BF16 LoRA training for Qwen3.8-Flash-Next

[aoshen02]

Summary

Keep this PR focused on FSDP2 LoRA initialization and a compact BF16 training/rollout recipe:

• Preserve frozen CPU-resident parameters marked by Transformers _no_placement_params during FSDP initialization and device transfers.
• Initialize sharded parameters from the full state without materializing the entire unsharded model on GPU; preserve runtime buffers and tied-parameter aliases.
• Avoid overlapping parent/child FSDP wrap targets.
• Preserve three-dimensional position-ID layout when reconstructing nested tensors.
• Provide one compact recipe: one shared checkpoint, FSDP16, vLLM TP4, CUDA Graph + MTP, adapter-only synchronization, dynamic batching, and packing disabled.

Removed from this PR: rollout.model_path, separate checkpoint-path plumbing, Slurm/Enroot launchers, checkpoint-lineage tooling, standalone adapter-export changes, and FP8-specific rollout modifications. Scheduling and checkpoint preparation are outside this recipe. The Ray address override and model-specific AutoModel selection are also removed: the actual BF16 checkpoint declares ForConditionalGeneration and uses upstream multimodal selection unchanged.

Why this is needed

The selected model implementation keeps a large frozen embedding on CPU even when the rest of the model is sharded or moved to GPU. Moving these parameters with the full model defeats that placement and can exhaust GPU memory. Separately, initializing FSDP from meta tensors must restore buffers as well as parameters, and nested wrapping must not select both an enclosing module and its child.

Tied-weight regression: plain Module.to_empty(device="meta") can replace a shared parameter with independent parameters. After an optimizer step, input embeddings and the tied output head diverge; loading an HF checkpoint re-ties them, so the checkpoint comparison fails. This was missed by the untied-model smoke test. The initialization path records parameter aliases and restores them only at the pre-FSDP meta conversion. After sharding, regular to_empty is retained: PyTorch preserves the DTensor parameter objects, so a second alias scan/rebinding is unnecessary. No model-specific naming or global PyTorch future flags are used.

Both trainer and rollout use the same BF16 checkpoint derived from the FP8 source. Existing VERL adapter-only synchronization is reused; no separate rollout checkpoint option or new quantization/export path is introduced. Model-specific Transformers and vLLM support must already be installed.

CPU parameter initialization now reuses VERL's existing CPU/Gloo process group, without a custom group or process-global weight cache. When VERL_NO_PLACEMENT_MMAP_DIR is configured on Ray workers, ranks on each host map one node-local file-backed copy of each frozen CPU parameter. Each initialization uses a fresh temporary file; it is unlinked after all readers map it, without invalidating the live mappings. Independently constructed models are not cached across invocations. The directory must be backed by node-local disk and visible at the same path to that host's ranks. This recipe disables the reference-policy KL paths.

The model's placement declaration already exists in upstream Transformers; this PR implements the FSDP-side handling, not a new model declaration.

Validation

Current CPU regression suite:

python -m pytest -q -o addopts='' \
  tests/test_protocol_v2_on_cpu.py \
  tests/utils/test_fsdp2_model_transfer_on_cpu.py \
  tests/utils/test_fsdp2_peft_wrapping.py \
  tests/utils/test_model_on_cpu.py

Result: 53 passed in a uv-managed environment. The latest simplification removes 18 net lines from fsdp_utils.py: redundant empty-input branches, aggregated missing-pattern bookkeeping, and API signature inspection. CPU shared storage, tied-parameter preservation, and buffer restoration remain intact. The final registration cleanup additionally passed all 5 CPU placement tests and the two-GPU distributed regression below. Ruff and git diff --check passed; the full-size training smoke test below predates this latest cleanup and was not rerun. Shell syntax was checked previously.

The updated distributed regression passed on two H200 GPUs (PyTorch 2.11.0+cu130):

PYTHONPATH=. python -m torch.distributed.run --standalone --nproc-per-node=2 \
  tests/special_distributed/test_fsdp2_full_state_load.py

It verifies exact reconstructed weights, restoration of a nonpersistent buffer, CPU placement and values of the frozen embedding on both ranks, and forward/backward with an embedding nested inside a wrapped block. The extended test also checks that both ranks map the same backing inode and that temporary files are removed. After the shared-storage change, this distributed test and the 11 CPU placement/wrapping tests passed again.

Additional tied-weight regression coverage uses a tiny BF16 Qwen2 with tied input/output embeddings on two H200 GPUs: meta conversion, FSDP2 loading, one AdamW update, full-state gathering, and HF save/reload. Tied identity is checked after initialization and every reloaded tensor is compared exactly (atol=0, rtol=0). A negative control replacing the helper with plain to_empty fails after the update: 489/512 embedding/head elements differ. The fixed version passes, including after removing the redundant post-sharding alias restoration. The full 53-test CPU suite also passes after that simplification; the distributed runtime check used PyTorch 2.11, while the DTensor conversion path was additionally inspected in the PyTorch 2.4 and 2.6 source (not runtime-tested on those versions). This is a focused regression test, not a claim that every outstanding CI failure is fixed.

Before/after checks against the PR base:

• Original position-ID handling fails to recover two equal-length samples after assigning the ragged-axis label; the corrected implementation recovers both exactly. Historical execution also reported split_with_sizes expects split_sizes to sum exactly to 204 ... but got split_sizes=[4].
• Original wrapping selection fails on a block containing an embedding with TypeError: Invariant encountered: value was None when it should not be; the retained selection completes forward/backward on both ranks.
• In the model-support container, loading the actual BF16 checkpoint configuration selects AutoModelForImageTextToText without the removed special case.
• The compact launcher's Hydra configuration composition passed.

Fresh one-step validation of the simplified code

A full-size one-step smoke test completed on 16 H200 GPUs with exit code 0: rollout, logprob computation, actor update, checkpoint saving on all 16 ranks, and post-step adapter synchronization all completed. Configuration remains BF16/BF16, FSDP16, TP4, CUDA Graph + MTP3, dynamic batching enabled and packing disabled. The recipe also restores the expert LoRA target parameters, optimizer settings, seed, and Triton compile-time autotune setting used by the historical successful run.

Regression caught during simplification: retaining a separate approximately 95.37 GiB frozen CPU embedding per rank caused host-memory exhaustion during the initial rollout sleep. Each inference rank additionally backed up approximately 90.30 GiB of weights to CPU. Ray killed a worker at 1892.96 / 1929.21 GiB host usage, exceeding its 98% threshold, before any training step. The minimal file-backed sharing restores one embedding copy per host without restoring the old global cache or custom process group. Same-inode mappings were verified in the actual training processes, not only the toy test.

Metric	New step	Historical median
Mean absolute rollout probability difference	0.006274	0.006409
Gradient norm	0.021118	0.023315
Reported host memory (GiB)	1299.38	1319.97
Actor update (seconds)	401.55	381.47
Post-step weight synchronization (seconds)	38.24	34.08

The new metrics fall within the historical observed ranges. Reward was 0.09375, with finite loss and gradient metrics. The measured step took 768.52 seconds; total launch-to-exit time was 46m21s including cold initialization. Rollout took 85.75 seconds and included first-use kernel compilation. Different checkpoints, sampled responses and token counts mean this is a sanity comparison, not a controlled speed or reward ablation.

A DataLoader-worker termination traceback appeared in Python exit finalization after training reached 1/1; the job still completed with exit code 0. This is recorded rather than claiming warning-free shutdown. The smoke test does not establish checkpoint-resume behavior or renewed long-run convergence.

Both this test and the historical learning run used separately installed model-specific serving and adapter-loading overlays, plus a local compatibility override for the image's development-version string. These are not added to this PR, and the evidence does not imply this PR alone provides complete model support.

Extended learning validation (BF16 train and rollout)

This follow-up experiment completed 159 global training steps. On held-out AIME2024, accuracy increased from 11/30 (36.7%) at step 0 to 21/30 (70.0%) at the latest evaluation, step 150, a 33.3 percentage-point gain. The best observed checkpoint was step 140 at 23/30 (76.7%); this peak is not the final score. The run was stopped before the planned 200 steps; the latest saved checkpoint is step 150.

Actual configuration: 16 H200 GPUs across two nodes; FSDP2 with a 16-GPU shard group; LoRA rank 16 / alpha 32; learning rate 1e-5; GRPO token-mean loss; 8 prompts x 8 responses per rollout; 4,096-token response limit; vLLM TP4 with CUDA Graph and three-token inline MTP; dynamic batching enabled and remove-padding/packing disabled. Training uses temperature 1; evaluation uses greedy decoding, one response per question, on the same 30 AIME2024 questions with the same 4,096-token response limit.

Scope: both trainer and rollout in this long run use BF16 weights derived from the same FP8 checkpoint source, with adapter-only synchronization. This is historical combined-stack validation, not a fresh run of the simplified launcher. It does not establish an FP8-rollout long-training result. The run also used model-specific serving and adapter-loading overlays; the curves are combined-stack evidence, not an isolated ablation of this PR's diff.

Train–rollout mismatch

The logged metric is training/rollout_probs_diff_mean: mean absolute probability difference on sampled response tokens, not mean absolute log-probability difference. Its 10-step average is 0.00637 at steps 31–40 and 0.00645 at steps 150–159. The corresponding average rollout_corr/k3_kl is 0.00159 and 0.00165. These observations show no sustained divergence over the displayed interval; they do not imply exact numerical equivalence.

Training reward

Raw per-rollout reward and a trailing 10-rollout mean are shown. Reward is -1 for an incorrect answer and +1 for a correct answer. The mean rises from -0.2656 over steps 31–40 to 0.5500 over steps 150–159 (36.7% to 77.5% sample accuracy). Training prompts vary between batches; this curve is not a fixed held-out evaluation.

Held-out evaluation reward

The axis and annotations both show mean test reward: step 0 = -0.2667, step 140 = 0.5333, and step 150 = 0.4000. These correspond to 36.7%, 76.7%, and 70.0% accuracy via accuracy = (reward + 1) / 2. Markers are recorded evaluations, not repeated-sampling estimates. One question changes accuracy by 3.33 percentage points, and the latest result is below the peak. All plots start at step 0 on the horizontal axis: training/mismatch include every completed step 1–159 (there is no training-batch metric at step 0), and evaluation includes step 0 and every 10 steps through 150. The first 30 steps were recovered from the original worker logs and joined to the resumed run by global step. At resume boundaries, the original completed-step evaluation is retained rather than overwritten by a startup re-evaluation: step 10 originally scored 11/30; its separate startup re-evaluation scored 12/30 and is not substituted into this curve. No missing values are synthesized.

Download the plotted metrics (CSV). Images and numerical data are pinned to an immutable asset commit on a separate evidence branch; no chart assets are added to this PR's code diff.

AI assistance

AI assistance was used to implement and review this change. The human submitter is responsible for reviewing every changed line and the validation evidence.

Related work

PR #7483 addresses DeepSeek V4 LoRA weight-name reconciliation and MXFP4 synchronization. This PR instead concerns FSDP2 parameter placement, initialization, and a BF16 adapter-only recipe.