[0z5a]

Closed. This draft targets 0z5a-branch and is too large to review in one pass. I will re-land it against main in smaller PRs:

1. Weight-sync hardening (init timeout, full resend after a failed update) — independent of FSDP2.
2. Backend registry and the two shared hooks, with CPU tests.
3. The FSDP2 actor, model, packing, LR schedule, checkpoint and weight push, with the GPU validation re-run on main.

The branch 0z5a:fsdp2-training-backend stays up for reference. Design discussion continues in #478.

Implements the design proposed in #478: --train-backend fsdp2, a second training backend next to Megatron that shards HF transformers causal LMs with PyTorch FSDP2 (fully_shard) over all training ranks.

The new code is in vime/backends/registry.py and vime/backends/fsdp_utils/. Shared training files get two small hooks:

• _pre_parse_mode() takes the --train-backend choices from the registry and merges the selected backend's own flags (--fsdp-*). The Megatron parser runs with ignore_unknown_args=True and would otherwise drop them silently.
• RayTrainGroup resolves the actor class through the registry.

The backend reuses vime's existing training pieces rather than copying them:

• the rollout side's DP/micro-batch schedule (each scheduled micro-batch is packed into one padding-free row with restarting position_ids);
• megatron_utils/loss.py and train_metric_utils, with Megatron's gradient scaling (the actor registers a pure data-parallel layout in megatron.core.mpu);
• the vLLM weight-transfer trainers, the same ones the Megatron updaters on this branch use: CUDA IPC for colocated engines, NCCL otherwise. They are fed by a source that materializes HF-named parameters bucket by bucket with full_tensor();
• vime's checkpoint layout, written with torch.distributed.checkpoint.

The LR / weight-decay schedule counts samples exactly like Megatron's OptimizerParamScheduler, and AdamW uses Megatron's no-decay grouping. The reference model shares the sharded module through host copies of each rank's shards. Megatron stays a dependency: fsdp2 launches go through Megatron's argument parser and validation and pass the usual MODEL_ARGS.

Scope:

• Supported: dense causal LMs, colocate and separate GPUs, GRPO-style actor with an optional reference model / KL loss, sleep/wake, checkpoint save/resume and HF export.
• Rejected with an explicit error: critic, Megatron OPD teacher, DSpark/MTP draft training, routing replay, multimodal or quantized checkpoints, TP/PP/CP/EP > 1, disk/delta weight sync, and async save.

Model loading, checkpointing, packing and the LR schedule are adapted from slime's former FSDP backend (Apache-2.0); provenance is noted in each file header.

Commits:

1. the backend;
2. removal of thin wrappers and unused state;
3. CPU tests against real Megatron when it is installed;
4. a small adaptation to this branch's weight-transfer helpers: create_nccl_trainer, and pause + flush on every sync because this branch has no --flush-cache-interval;
5. fix: push colocated weights with torch_memory_saver tracking disabled (found on H100, see below);
6. a CPU test for that fix, and a non-vacuous gradient in the alignment test;
7. a CPU test that checks the step gradient, including the KL loss and the reference-model switch, against an unsharded computation of the same loss.

The diff contains only these changes (19 files, +2903/−5) on top of 0z5a-branch at 8c8b8e8e.

Validation

H100 runs (4× H100, Qwen3-0.6B)

All three runs used the equivalent code on top of main; this branch differs from it only in commit 4's weight-transfer helper calls. The second and third runs include commits 5 and 6; commit 7 only adds a CPU test. The third run only replays the KL variant's step-0 batch (see below). In the end-to-end cases vLLM starts with dummy weights, so rollouts can only see weights the trainer pushed.

Check (second run unless noted)	Result
The three CPU test files inside the GPU image, with real Megatron and torch_memory_saver	36 passed, 0 skipped
Colocated GRPO through train.py: 3 rollouts × 2 steps, KL loss with a reference model, a DCP checkpoint per rollout	passed. CUDA IPC weight push works after commit 5; per-step train-vs-rollout log-prob diff 0.014–0.019
Separate training (2 GPUs) and rollout (2 GPUs), NCCL weight push	passed; per-step train-vs-rollout log-prob diff 0.016–0.021
Sleep/wake: the colocated run replayed without ever sleeping	bitwise identical dumps, parameters, Adam moments, step count and LR schedule. Sleep moves the actor's allocated memory from 1.73 GB to 0.06 GB
DCP resume: restart from the colocated run's rollout-1 checkpoint and train rollout 2	bitwise identical to the continuous run's rollout-2 checkpoint
Megatron vs fsdp2 alignment test, same weights and batch (flash-attention 2)	aligned. Grad norms 0.082196 vs 0.082150 and 0.086195 vs 0.086167 (0.056%, 0.033%). Old-policy log-probs: mean per-sample abs diff 1.84e-2, max 2.48e-2, which is 1.11× Megatron's own train-vs-rollout diff on the same batch (1.66e-2)
The same alignment with a k1 KL loss (coefficient 1.0) added	aligned within bf16 numerics (third run). Step-0 grad norm Megatron 2.518 vs fsdp2 2.491 (1.07%); the test's built-in 1% check stopped the second run there. The fp32 replay below puts the exact value between the two. A bf16 rerun of the same batch without that check: step 1 fsdp2 2.5265 vs Megatron 2.5225 (0.16%); old-policy log-probs mean per-sample abs diff 1.92e-2, 1.09× Megatron's own train-vs-rollout diff on this batch (1.76e-2)
Packing with real rollout tokens (first run)	flash-attention 2 packed vs per-sample identical; sdpa fp32 4.7e-6

How the alignment is judged. The criterion is relative to bf16 noise and was fixed before the second run:

• log-prob: the fsdp2-vs-Megatron mean per-sample abs diff must be ≤ 1.25 × Megatron's own mean train-vs-rollout diff on the same batch;
• grad norm: both sides nonzero, and the relative diff ≤ 5% at every step.

The test's built-in limits are stricter (log-prob ≤ 2e-2; grad norm within 1%). By those, the plain alignment also passes, and the KL variant fails at step 0. The test itself has no KL term, so its limits are unchanged.

KL variant on CPU. The plain alignment, without the KL term, agrees within 0.06%, and the KL term shares its loss scaling with the rest of the loss, so the gap comes with the k1 term. At coefficient 1.0 that term dominates the gradient (grad norm about 2.5 with it, about 0.08 in the plain alignment run). fsdp2's KL loss and reference-model switch were therefore checked on CPU: real fully_shard over gloo, 2 ranks, fp32, a tiny Qwen3 with tied embeddings. The fsdp2 trainer's step-0 gradient was compared with an independent plain-PyTorch computation of the same loss (zero advantages, entropy coefficient 0.01, k1 KL coefficient 1.0):

Case	Gradient-norm rel diff	Worst per-parameter rel diff
Entropy term only	2.4e-7	—
KL, reference model == actor	9.4e-8	5.1e-7
KL, reference model ≠ actor	9.4e-8	5.1e-7

The k1 gradient does not depend on the reference values, and the last two cases agree. So fsdp2's KL loss and its switch between actor and reference weights compute the intended gradient; commit 7 keeps this as a test. This rules out a logic error in fsdp2's KL / reference-model path.

KL variant: fp32 replay on H100. The second run's step-0 batch was replayed without bf16:

Step-0 grad norm	Value	vs fp64
Independent single-GPU HF computation, fp64, one unpacked sequence at a time	2.5028826	—
fsdp2 trainer, fp32 (--fsdp-param-dtype fp32, sdpa, TF32 off, 4 ranks)	2.5028820	−2.5e-7
Same, with --fsdp-gradient-checkpointing	2.5028818	−3e-7
Independent unsharded HF computation on CPU, fp32	2.5028807	−8e-7
Megatron, bf16 (second run)	2.5180326	+0.61%
fsdp2, bf16, flash-attention 2 (second run / third-run rerun)	2.4909747 / 2.4909532	−0.48% / −0.48%
fsdp2, bf16, sdpa	2.4887922	−0.56%
HF, bf16, unpacked	2.4995402	−0.13%

Without bf16, the fsdp2 trainer reproduces the exact gradient. The 1.07% gap is the bf16 numerics of the two implementations, with opposite signs, and needs no code change. At step 1, which already carries step 0's update, both bf16 values sit above the fp32 replay (2.5077): Megatron by 0.59%, fsdp2 by 0.75%. Each bf16 configuration was run once; the fsdp2 rerun matches the second run within 1e-5 relative.

The first run found two problems, both fixed here:

• Colocated weight push (commit 5). The first push failed in vLLM's packed IPC producer: storage._share_cuda_() raised CUDA error: invalid argument. With --offload-train, the train group preloads torch_memory_saver with tracking enabled from process start, so every CUDA allocation of the actor is a VMM mapping that CUDA IPC cannot export. The Megatron actor already runs its push under torch_memory_saver.disable(). In the second run a standalone probe reproduced the failure under tracking and showed the export succeeding under disable().
• Vacuous gradient check (commit 6). In the first run every response was truncated, so rewards, advantages and both gradients were zero, and the alignment test compared zero against zero. Commit 6 adds a small entropy term to the test's loss and stops the test if Megatron's gradient is zero.

CPU, run on this branch

• tests/test_fsdp2_backend_registry.py: 20 passed, 1 skipped without Megatron; 21 passed with Megatron-LM importable.

  • The Megatron-only case runs a full parse_args() for an fsdp2 launch.
  • The new weight-push test checks that torch_memory_saver.disable() is entered only under --offload-train. Removing commit 5 makes it fail.

• tests/test_fsdp2_data_packing.py: 9 passed.

• tests/test_fsdp2_cpu_smoke.py: 9 passed, with both the mpu stub and real megatron.core.parallel_state. It runs real fully_shard over gloo on a tiny random Qwen3 and checks:

  • packed log-probs against an unsharded reference;
  • identical updates for dp=1 and dp=2 (GRPO, tied embeddings with per-token loss, KL with a reference model, single-step dump);
  • weight export;
  • the step-0 gradient against an unsharded fp32 computation of the same loss: entropy term only, KL with the reference equal to the actor, and KL with a different reference;
  • reference-model switching, DCP round trip and HF export.

  Removing the micro-batch or world-size loss scaling makes the dp comparison fail; removing the switch back to the actor weights after the reference log-probs makes the differing-reference gradient case fail.

• Existing tests touched by the shared hooks: 236 passed, 2 skipped.

• pre-commit hooks passed on the changed files.

The real-Megatron runs on CPU used a local Megatron-LM v0.19.2 checkout (4b4acac9), not the commit pinned in this branch's Dockerfile. Buildkite's CPU steps passed on the current head.

Not covered yet:

• the policy-gradient term with nonzero advantages in the Megatron comparisons: every response in both alignment batches was truncated, so all advantages were zero and only the entropy and KL terms carried gradient;
• --fsdp-cpu-offload;
• gradient checkpointing beyond the single fp32 step above;
• multi-node runs;
• larger models;
• a throughput or memory comparison with Megatron.

AI assistance: implementation and tests were prepared with Claude Code at the submitter's request, with attribution and DCO sign-off in each commit. Kept as a draft pending the submitter's line-by-line review; no end-to-end training-quality or performance result is claimed.

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #35172927 | 📁 Comparing 4394eb9 against latest (d7e40fe)

  🔍 Preview build  

68 files changed · + 9 added · ± 38 modified · - 21 deleted

+ Added

• advanced/draft-feature-collection.html
• examples/deepseek-r1.html
• examples/glm4-9B.html
• examples/glm4.7-30B-A3B.html
• examples/glm4.7-355B-A32B.html
• examples/glm5.2-744B-A40B.html
• examples/qwen3-30B-A3B.html
• examples/qwen3-4B.html
• _examples_synced/ouro/README.html

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 28 more...

- Deleted

• library.html
• advanced/parallelism-memory.html
• advanced/policy-mismatch.html
• advanced/rl-systems.html
• advanced/rollout-scheduling.html
• advanced/straw.html
• blogs/index.html
• blogs/introducing_slime.html
• blogs/release_v0.1.0.html
• developer_guide/index.html
• and 11 more...

[gemini-code-assist]

Code Review

This pull request introduces a new FSDP2 training backend (--train-backend fsdp2) to shard Hugging Face transformers models using PyTorch's fully_shard. It adds the necessary Ray training actor, trainer, distributed checkpointing, data packing, learning rate scheduler, and weight updater to push weights to vLLM rollout engines, along with comprehensive CPU unit and smoke tests. Feedback on the changes highlights two key issues: a critical race condition in FSDPWeightSource._materialize where calling to_local() on a DTensor before wait() prevents waiting on the asynchronous redistribute collective, and a potential ZeroDivisionError in FSDPParamScheduler.get_wd if self.wd_incr_steps is zero.

[gemini-code-assist]

In _materialize, calling tensor.to_local() immediately after redistribute(..., async_op=True) on a DTensor returns a standard torch.Tensor. Since standard PyTorch tensors do not have a wait() method, the subsequent hasattr(tensor, \"wait\") check evaluates to False, and the code never waits for the asynchronous redistribute collective to complete. This creates a critical race condition where incomplete or uninitialized weights are yielded and sent to the rollout engines.\n\nTo fix this, keep the DTensor in the pending list, call wait() on it first, and only then call to_local() to retrieve the physical local tensor.

            if isinstance(tensor, DTensor):
                tensor = tensor.redistribute(placements=[Replicate()] * tensor.device_mesh.ndim, async_op=True)
            pending.append((name, tensor))
        for name, tensor in pending:
            if hasattr(tensor, \"wait\"):
                tensor = tensor.wait()
            if isinstance(tensor, DTensor):
                tensor = tensor.to_local()
            yield name, tensor

[gemini-code-assist]

If self.wd_incr_steps is 0 (which can happen if args.train_iters is 0 due to a small dataset or large global batch size), the division float(self.num_steps) / float(self.wd_incr_steps) at line 86 will raise a ZeroDivisionError for non-constant weight decay styles.\n\nAdding a guard to check if self.wd_incr_steps == 0 and returning self.end_wd immediately prevents this crash.

Suggested change

	if self.num_steps > self.wd_incr_steps:
	return self.end_wd
	if self.wd_incr_steps == 0 or self.num_steps > self.wd_incr_steps:
	return self.end_wd