Repository navigation
Conversation
Add `--train-backend fsdp2`: HF transformers causal LMs sharded with `torch.distributed.fsdp.fully_shard` over every training rank, as a second backend next to Megatron. Backend selection goes through a new `vime/backends/registry.py` (backend name -> actor class and backend-only flags, by import path). Shared training files change in two places only: - `_pre_parse_mode()` takes the `--train-backend` choices from the registry and merges the selected backend's own flags (`--fsdp-*`), which the Megatron parser would otherwise drop as unknown; - `RayTrainGroup` resolves the actor class through the registry. The backend lives in `vime/backends/fsdp_utils/` and reuses vime's shared pieces instead of copying them: - the rollout side's DP/micro-batch schedule (micro-batches are packed into one padding-free row with restarting position ids); - `megatron_utils/loss.py` (log-probs, advantages, policy/KL/entropy losses) and `train_metric_utils`, with Megatron's gradient scaling; these read their process groups from `megatron.core.mpu`, so the actor registers the pure data-parallel layout (TP = PP = CP = 1); - the vLLM weight-transfer trainers (CUDA IPC for colocated engines, NCCL otherwise) fed by a source that materializes parameters per bucket with `full_tensor()` under their HF names; - vime's checkpoint layout (`iter_<rollout_id>` + tracker), saved with torch.distributed.checkpoint, plus optional HF export. The LR/weight-decay schedule counts samples like Megatron's, AdamW uses Megatron's no-decay grouping, and the reference model shares the sharded module through host copies of each rank's shards. Scope: dense causal LMs (e.g. Qwen3), colocate and non-colocate, GRPO-style actor with optional reference model, sleep/wake (shards and optimizer state moved to the host), checkpoint save/resume. Critic, Megatron OPD teacher, draft/MTP training, routing replay, multimodal, TP/PP/CP and disk/delta weight sync are rejected with an explicit error. Megatron's argument parser and model-shape validation still run for fsdp2, so launches keep passing the model-arch flags (`scripts/models/*.sh`). Model loading, the rank-0 broadcast flow, checkpointing, packing and the LR schedule are adapted from slime's FSDP backend before its removal (THUDM/slime 0104a9e, Apache-2.0); provenance is noted in each file. Tests: - CPU (registered in the plugin-contracts step): backend registry and flag parsing, packing/schedule/LR, and a gloo `fully_shard` smoke test on a tiny random Qwen3 that checks packed log-probs against an unsharded model, identical updates for dp=1 and dp=2 (GRPO, per-token loss, KL with a reference model), weight export, model switching, checkpoint round trip and HF export. - GPU (registered in the vime-customized suite): Megatron vs fsdp2 on the same weights and batch (grad norms and old-policy log-probs). It has not been run yet; its tolerances are first estimates. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Follow vime's review preferences on the new backend: - inline single-use wrappers (pure data-parallel mpu registration, HF state dict loading, meta-device construction, rollout dtype) at their only call sites and remove `fsdp_utils/parallel_state.py`; - remove code only tests used (`register_train_backend`, `unpack_sequences`, `is_fsdp_checkpoint`) and state nothing read (`FSDPTrainer.active_tag`, the always-empty weight-update metrics); - reuse `ROLLOUT_TOP_P_TOKEN_KEYS` from the shared loss module; - read configuration attributes directly instead of `getattr` defaults, and reject disk weight sync with one predicate (delta sync already requires the disk transport). Behavior is unchanged. The weight source's contract was checked against the vLLM weight-transfer trainers at the pinned commit 43b4aaea: the IPC and NCCL trainers only call `metadata()` and iterate the source, and NCCL verifies each yielded name, dtype and shape against the metadata. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
The fsdp2 backend keeps Megatron as a dependency: launches go through Megatron's argument parser and validation, and the shared loss/metric helpers read their process groups from `megatron.core.mpu`. - `test_fsdp2_cpu_smoke.py` registers the actor's pure data-parallel layout with Megatron's `initialize_model_parallel` when Megatron is installed, and falls back to the stub only without it. - `test_fsdp2_backend_registry.py` adds a full `parse_args()` for an fsdp2 launch (skipped without Megatron): the `--fsdp-*` flags survive Megatron's parser, and the shared fields the trainer reads (load path, LR, clip, weight-decay schedule) are filled in by Megatron's validation. - Both files also run in the `upstream-sync-cpu` step, whose image ships the pinned Megatron, so CI covers the real path as well as the stubbed one. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Match how the Megatron updaters on this branch drive the vLLM weight transfer: create the NCCL trainer with `create_nccl_trainer`, shut every trainer down the same way on disconnect, and pause generation and flush the cache on every sync (there is no `--flush-cache-interval` here). Also keep the CPU test count in `.buildkite/README.md` in step with this branch's pipeline. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
…cking 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. CUDA IPC cannot export those: the first colocated weight push failed in vLLM's packed IPC producer with `storage._share_cuda_() -> CUDA error: invalid argument` (Qwen3-0.6B, 4x H100, colocate). The Megatron actor already runs its push under `torch_memory_saver.disable()`; do the same in the fsdp2 actor. The actor sleeps (and empties the cache) before every push, so the transfer buffers are fresh, untracked allocations. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
… gradient non-vacuous - Add a CPU test that the actor pushes weights with torch_memory_saver tracking disabled under --offload-train, and never touches it otherwise. It fails without the previous fix. - In the first H100 run of the Megatron/fsdp2 alignment test every response was truncated, so all rewards and GRPO advantages were zero and both backends reported a zero gradient: the 1% grad-norm comparison passed without comparing anything. Give the loss a small entropy term so the gradient does not depend on rewards alone, and stop the test when Megatron's gradient is zero. The CPU check of the script covers both. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
…nce loss Add a CPU test that compares the trainer's step-0 gradient on two gloo ranks with a plain, unsharded fp32 computation of the same loss (entropy term, optionally k1 KL against a reference model). It covers packing, loss scaling, the reduce-scatter and the switch between the actor and reference weights in one place. With all rewards zero the GRPO advantages vanish, so the expected gradient comes from the entropy and KL terms alone. With k1 the gradient does not depend on the reference log-probs, so a reference model that differs from the actor leaves the expected gradient unchanged. That case fails if the trainer keeps computing on the reference weights after the reference log-probs (checked by removing the switch back to the actor). All three cases (entropy only, KL with reference equal to the actor, KL with a different reference) match within rtol 1e-4. This rules out a logic error in the KL / reference-model path; it does not exercise bf16 numerics. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Incremental review layer for vllm-project#479: the FSDP2 training backend (7 commits), rebased onto the update to main.
Base:
sync-0z5a-branch-to-main(vllm-project#483,0z5a-branchupdated tomain), so Files changed shows this layer only. The cumulative upstream PR is vllm-project#479.🤖 Generated with Claude Code