Repository navigation
p2p: update speculative drafts from trained MTP weights - #3985
Conversation
With a draft model, each rollout engine runner registered its weights under the rank alone, so a p2p sender read the draft's weight table for the target, and the draft itself was never written: the trained MTP layer never reached it. sglang now publishes each runner under its role. When the update covers the draft (the trainer holds MTP layers) and the engines draft with the target model's own MTP layer, a sender keeps replicas, a stager and write targets per runner; the draft replica is built from the draft model config under the draft's speculative MoE backends, as EAGLE v2 builds its draft. A bucket's HF tensors go to the first runner whose replica loads them, the target before the draft, so the weights a draft shares with its target are written once. Multi-layer EAGLE, drafts from their own checkpoint and engines that run different runners are rejected at connect.
Four training GPUs (TP1 DP4 EP4) write the target and its MTP draft into two TP2 EAGLE engines over Mooncake; the equality check covers both.
The loader probe maps the Qwen3.5 target, which ParameterMapper never did, and its Gemma norms' derived buffers. Four training GPUs write the target and its MTP draft into one TP4 EAGLE engine; the equality check covers both, vision weights skipped.
yueming-yuan
added
run-ci-weight-update
yueming-yuan
added
run-on-hopper
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
Resolve overlapping docstrings using the shorter contracts while retaining target-first routing and per-role replica lifetime. Keep the draft documentation row alongside the corrected parameter-group buffer sizing.
yueming-yuan
removed
the
run-on-hopper
There was a problem hiding this comment.
I reviewed this PR and didn't find any bugs. Because it's a large change to the core P2P weight-transfer path spanning speculative-decoding draft runners, Mooncake transport, and write-failure handling, a human look is still worthwhile.
What was reviewed (new since the last run, commits 5496f76/43a9bd2/04a7485): the double-load zero-vs-0xFF probe in model_replica.py that rejects a loader leaving param bytes unwritten — the byte-diff and unwritten-byte-count logic checks out against its tests; the move of write-failure/timeout handling into TransferBuffers.acquire()/wait_for_all_writes so a broken write fails the update before more writes queue behind it; the previously-flagged candidate that a draft's shared embed/head weights could stay stale if its registered address differs from the target's remains open and worth human attention.
Extended reasoning...
This review's scope diff (b056fc6..HEAD) reduces, net of what I already reviewed twice before, to the delta between commit 088fc53 (my last review) and HEAD: a refactor moving write-failure/timeout handling into TransferBuffers and a new double-load probe in model_replica.py that detects a loader leaving param bytes unwritten. The surface touched is purely internal weight-transfer plumbing (no auth/crypto/injection exposure); the new probe logic is subtle (compares results of loading over zero-filled vs 0xFF-filled buffers to detect untouched bytes) but is exercised by matching unit tests whose expected byte counts I verified against the test fixtures. Given the overall PR remains large and touches a critical correctness path (RL rollout weight synchronization) with prior non-author review comments addressed via follow-up commits, and the previously surfaced candidate about stale shared draft weights is still outstanding, a human review is still warranted even though this run found no new bugs.
This review covers commit 04a7485, which is no longer the latest commit on this pull request; later commits are not covered by it.
# Conflicts: # docs/advanced/p2p-weight-transfer.md
# Conflicts: # miles/backends/training_utils/weight_update/protocols/p2p.py
send_bucket writes each HF tensor through the first runner that loads it, the target, which is right only while the draft loads such a tensor into params it shares with the target, as drafts share its embed and head. A draft holding its own copy would never be written, and its acceptance would fall as training moved the target. begin_sync now checks that each such draft param is published at an address the target publishes on the same engine rank.
Routing each HF tensor to the first runner that loads it never updated a draft's own copy of a target weight. The engine's own weight updates run every runner's loader over the whole stream; p2p now does the same, so the check rejecting such drafts goes. A weight the draft shares with the target is written by both with the same bytes.
… sends Its weight would keep the value the engine started with. Each sender loads only its own engine ranks (EP splits the experts), so every rank ORs its view of its PP stage's tensors in one all_reduce, once per connect. Exempt: towers (resent unchanged; a language-only engine has none), lm_head on a tied engine, and MTP layers when no draft runner is written. Replaces the info log of skipped tensors.
… trainer sends" This reverts commit ee00ab4. Coverage rests on --check-weight-update-equal, as on the non-p2p path. Telling at runtime which tensors the rollout engines accept needs a maintained list of exceptions (towers, tied heads, MTP layers without a draft runner), which does not hold up as models and formats change.
|
Known issue, to be fixed in a follow-up PR:
Pre-existing and not p2p-specific: the same Planned fix: move the trainer-has-MTP rule into |
train/grad_norm jitters between runs and is exactly 0 on steps where every GRPO group's rewards are uniform, so the old default band (0.5x up, 0.8x down, 0.1 floor) failed healthy runs. Twenty tests had opted out with enforce=False, but a failing shadow check still untrusts the run: 480 of their 490 untrusted stored runs carry such a failure, so their baselines for every metric starved. The P2P MTP tests (#3985) still enforced it. The gate only needs to catch an order-of-magnitude jump. The default now allows up to 10x the trusted mean (1.0 when that mean is near zero) and any value down to 0, and every grad_norm declaration is enforced again, which also reverts the previous commit's per-test opt-out. Replaying all stored grad_norm history with every run as a baseline, the old band fails 262 of 3347 comparisons with at least ten baselines and the new band fails none.
Update both target and MTP draft weights through P2P. Previously, the draft replaced the target's weight registration, causing connect failures, and trained MTP weights never reached the draft (sgl-project/sglang#35312).
Create a sender and replica per runner role, using the draft's model config and speculative MoE backends. Route shared embedding/head inputs to the target once. The companion SGLang change registers weights by
(role, tp_rank).Supports the model's own single EAGLE MTP layer with selector
all. Separate draft checkpoints, multi-layer EAGLE and a draft with its own MoE runner or a2a backend (sglang reloads a draft under the target's) are rejected; NGRAM/UNO update only the target. Adds GLM-4.7-Flash and Qwen3.5 MTP e2e coverage, including PP-stage writes to multiple engines.Stacked on #3969; requires sgl-project/sglang#42820 and sgl-project/sglang#43251, which rebuilds Qwen3.5's
gemma_weightand runs the GLM-lite NextN post-load after p2p's writes.Recorded validation
Target/draft equality passed with
--check-weight-update-equal:Target-only GB300 Qwen3-30B-A3B TP2 passed in bf16/TRT-LLM and fp8-block/UE8M0. The GB300 equivalence suite passed bf16, fp8 block with UE8M0 and fp32 scales, mxfp8 and nvfp4.
Failure controls reproduced the missing per-role registration at connect,
Attempting to copy from device cpu to device metawithout #3969's buffer materialization, and UE8M0 loader failures with meta probe inputs.304 weight-update/API-client tests passed on H200 with sglang (296 on #3969 alone). Cleanup passed pre-commit and 64 before/after context-scope cases. The GLM timing row was measured on cleanup revision
ddcfab372a; other table/equivalence results predate that cleanup.H200 CI run: GLM test step 758 s, Qwen3.5 1132 s; registered estimates 1000/1500 s.
ci-sglang-pr: #43251