Skip to content

[megatron] feat: delta_sharded on Megatron-Bridge param mappings (TP+EP, hybrid-Mamba) - #7181

Merged
wuxibin89 merged 23 commits into
verl-project:mainfrom
ChangyiYang:delta-megatron
Jul 31, 2026
Merged

[megatron] feat: delta_sharded on Megatron-Bridge param mappings (TP+EP, hybrid-Mamba)#7181
wuxibin89 merged 23 commits into
verl-project:mainfrom
ChangyiYang:delta-megatron

Conversation

@ChangyiYang

@ChangyiYang ChangyiYang commented Jul 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

What does this PR do?

The mcore backend joins the sharded-delta contract (#7144). #7085 (veomni EP) is merged; this branch is rebased onto latest main, so the diff is the megatron/engine increment only. Everything rides Megatron-Bridge's own per-param machinery — no legacy converters, no hand-written index math:

  • Export index: AutoBridge.get_conversion_tasks enumerates parameters; each becomes a ShardSpec (TP partition dims → BlockPlacement over the TP group; expert tensors → virtual [ep, *etp] blocks over the etp×ep group; replicated → the engine's rank-0 direct path). No hand-derived geometry: the engine only consumes gather_group + the replicated-rank0 rule.
  • Comm-stub probe: each mapping's megatron_to_hf runs on a probe copy whose process groups are _ProbeGroups (real size()/rank(), any collective raises) and whose gather helpers are stubbed to pure local semantics — NaN sentinels through the REAL converter code give final HF coordinates with zero collectives and zero converter forks. This keeps converters with real value math (Mamba's TP de-interleave, GQA fusions) bit-exact, which the earlier "degenerate world" probe could not.
  • Idempotence verify sweep (verify_every=K engine kwarg): every K-th sync appends a full re-export in the same receive session; the rollout loader snapshots every copy_ destination before its real load path, re-applies, and fails loud on any bitwise divergence — a production-grade end-to-end oracle that tolerates multi-stage loaders (e.g. sglang's composed A_log loader).
  • Perf: shape-pooled NaN probe buffers (Opt-A) and an empty-delta short-circuit reusing the slot cache (Opt-B); snapshots default to pageable host memory on mcore (delta_pin_snapshots = False class attr — pinning whole shard sets surfaces as CUDA OOM at 30B/235B scale).

Scope: TP+EP+ETP, PP=1/VPP=1, no LoRA, Megatron-Bridge flavor only.

Test

check result
D1 equivalence (tiny Qwen2, real bridge, 1 GPU) probe entries == bitwise diff of the bridge's own full HF export, 22/22 HF tensors
TP2 differential oracle (tests/special_distributed/test_mcore_probe_differential.py) real megatron_to_hf vs merged probe outputs, bitwise, on qwen2 / qwen3_moe / nemotron_h (Mamba) / falcon_h1
7B TP2 e2e uw 3.8–4.9s steady; idempotence sweeps (v=2/4/6) zero mismatch
30B-A3B TP1×EP8 + moe_grouped_gemm uw 14.5–17.6s steady (hand-written-converter baseline: 17.2–19.7s)
235B-A22B TP4×EP16×PP1 + grouped_gemm (16 trainer GPUs + TP16 rollout) seed 107.6s/437.9GB; steady uw 23.6/23.7s, payload ~20GB/step; same-shape NCCL full sync: 235.0–237.6s/step → delta is 8.2–9.7×, seed < half one NCCL step
hybrid-Mamba e2e (NVIDIA-Nemotron-Nano-9B-v2, TP1 and TP2) 6 steps + 3 full idempotence sweeps (3.39B elems each) zero mismatch; 50-step GRPO W&B pair mechanically green
Opt-A regression comm-stub refactor cost +25–40% at 235B; pooled buffers recovered to 23.6/23.7s (baseline 24.3–29.5s)
Opt-B empty-delta short-circuit validated 7B TP2 e2e + sweeps
CPU suites tests/checkpoint_engine 32 passed

Additional Info

🤖 Generated with Claude Code

@gemini-code-assist

gemini-code-assist Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

ChangyiYang and others added 23 commits July 30, 2026 02:41
@ChangyiYang @claude
[checkpoint_engine] log the seed's one-off cost at warning level
bcc5f4b 
Worker default log level swallows info, so the FULL-SEED line never reached
any console; the seed duration is the number people ask for when sizing a
run. Also stamps elapsed time and wire GB on the line.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] delta export on Megatron-Bridge param mappings (TP+EP, PP=1)
0f6363f 
The mcore backend joins the HF-delta contract:

- delta_export.py: build_export_index enumerates parameters through
  AutoBridge.get_conversion_tasks and declares each one's geometry as a
  ShardSpec (TP partition dims -> BlockPlacement over the TP group, expert
  tensors -> virtual [ep, *etp] blocks over the etp x ep group, replicated ->
  the engine's rank-0 direct path). mcore_hf_delta_entry probes the bridge's
  own megatron_to_hf with NaN sentinels: make_probe copies each mapping and
  installs degenerate single-rank process groups (recursively -- composite
  mappings like QKVMapping delegate their TP gather to an inner _tp_mapping
  that does not receive the outer injection), so the converter short-circuits
  every collective and degrades into a pure permutation the sentinels survive.
- transformer_impl: get_per_tensor_param_shard / get_per_tensor_param_delta_
  shard / _hf_delta_entry wired on the index; seed and prime come free
  (bridge full export + BaseEngine).

Scope: TP+EP, PP=1/VPP=1, no LoRA, Megatron-Bridge flavor only (the
deprecated vanilla mbridge is rejected).

Validated on a tiny Qwen2 through the real bridge (1 GPU, container):
probe entries match the bitwise diff of the bridge's own full HF export on
all 22 touched HF tensors (0 mismatched, 0 extra); CPU suite 24/24.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] probe injection must reach AutoMapping's lazy delegate
656ea2f 
QKVMapping._tp_mapping is an AutoMapping whose concrete Column/Row/Replicated
delegate is created lazily on first use, snapshotting whatever process groups
exist at that moment -- when that first use happened inside the probe, the
delegate was born with the REAL groups and gathered for real (TP2 smoke:
qkv-bias buffer arrived at split_qkv_biases doubled, 9216 vs 4608). Warm the
lazy delegates from the task's module before copying, and inject the
degenerate groups recursively into every nested mapping.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[checkpoint_engine] return cupy staging blocks to CUDA after each send
38eafb9 
The cupy pool holds its blocks; after the full seed streams up to 2x
bucket_size through it, that memory is invisible to the torch allocator and
raw cudaMalloc (NCCL workspaces, optimizer buffers) OOMs on tight mcore
shapes. Free the pool at the end of the seed and of every steady send.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[engine] VERL_DELTA_PIN toggle for the snapshot pinning
7c8edde 
Pinning a whole shard set per rank (cudaHostAlloc) competes with everything
else that pins on the node, and its failure surfaces as a CUDA out-of-memory
during unrelated allocations. VERL_DELTA_PIN=0 keeps the delta snapshots in
pageable host memory (slower diff read-back, cannot OOM the device) -- also
the isolation knob for the mcore 30B OOM triage.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] expert partition-dim fallback by layout rule
7b555a4 
mcore expert weights do not always carry the standard tensor_model_parallel
attribute (PR#2 hit the same); without it the ETP-full shape never got its
etp factor and the expert slot shapes/coordinates were declared at 1/etp
size -- the rollout's fused-MoE loader then narrowed out of bounds
(start 192 + length 192 exceeds 192). Fall back to the mcore layout rule
(fc1 column dim 0, fc2 row dim 1) and treat '.mlp.experts.' names as expert
params regardless of the mapping flag.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[checkpoint_engine] measure the staging pool release
b6004a6 
Log (warning level) how much the cupy staging pool held at the end of each
send and the device-free delta across the release -- the e2e evidence for
the pool-holding mechanism, and a permanent observability line for the
staging footprint.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] delta export: drop dead is_expert name fallback
2ce8a62 
The bridge's MegatronParamMapping.is_expert is a base-class property and is
itself a name test broader than our ".mlp.experts." check, so the or-clause
could never fire; keep the mapping property as the single source. Reword the
expert pdim comment to state what the mcore source shows: tensor_model_parallel
is deliberately not stamped on expert weights, so for experts the layout rule
is the primary source rather than a fallback.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] delta export: drop unused hf_path parameter
67a522f 
No caller passes it: the engine builds its bridge from_hf_pretrained, which
carries the HF config, so get_conversion_tasks never needs the path. Leftover
from an early probe draft.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] warm the lazy mapping self-only; document the probe's safe…
e2c8c7f 
…ty boundary

_inject already warms every child right before copying it, so each node of the
mapping tree is warmed exactly once -- the children scan inside
_warm_lazy_mappings was redundant (and used a different filter than _inject's
isinstance). The ordering invariant that matters (delegate must exist before
the child is copied) is now stated at the call site. Module docstring gains
the degenerate-PG trick's safety boundary: group state must stay
communication-only inside megatron_to_hf; the D1 bitwise probe is the
regression oracle on Megatron-Bridge upgrades.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] comm-stubbed probe: real group sizes, no communication
59708ab 
A group carries two separable meanings -- SIZE/RANK (value math such as
Mamba's per-rank de-interleave dims consumes it) and COMMUNICATION. The old
degenerate single-rank groups collapsed both, which made mappings whose
transform depends on tp_size (MambaInProj, GDN, Chunked) silently wrong.
Probe copies now keep sizes faithful via _ProbeGroup (real size/rank; any
actual communication through the group raises) and stub the bridge's comm
helpers with local synthesis: gather_from_tp_ranks returns this rank's shard
at its true rank index with NaN placeholders elsewhere, the ep fan-out reduces
to the task-bound global name. megatron_to_hf now runs its REAL tp>1 code
paths on the LOCAL shard buffer, so the entry no longer builds full-shape
buffers or translates indices, and unknown communication patterns fail loud
instead of skewing silently.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] audit leftovers + TP>1 probe differential oracle
6b7f2d9 
make_probe's module=None path was unreachable (every caller passes the task's
module) -- module is required now. spec.py's to_hf_chunk comment described the
pre-backend-export receiver-side segment rebuild; it documents the sender-side
NaN row probe now. New tests/special_distributed/test_mcore_probe_differential.py:
torchrun TP>1 oracle asserting real megatron_to_hf == merged comm-stubbed probe
outputs bitwise per conversion task (MODEL_KIND=qwen2|qwen3_moe|nemotron_h),
the regression gate for Megatron-Bridge upgrades and the Mamba acceptance.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] drop the hand-computed shard geometry from the export index
3c06129 
With the comm-stubbed probe emitting final HF coordinates straight from the
local shard, the engine consumes only each parameter's wire merge group and
the replicated-from-rank-0 rule -- the hand-built BlockPlacements (TP slice
offsets, the virtual all-experts tensor, the expert partition-dim layout rule)
had no reader left. Expert entries merge over the joint etp x ep group; the
former ep_size==1-with-TP expert gap is covered by the same branch.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] sweep stale probe wording from the engine docstrings
23b82e6 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[test] falcon_h1 kind for the TP2 probe differential (hub-kernel stub…
d4ed4dd 
…s, versioned tiny-model cache)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] tolerate NoPE configs in make_megatron_module
f9cc705 
NemotronH-family hybrids carry no rope at all; the provider/bridge owns rotary
setup, so a missing rope_theta is not an error on the bridge path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] pool the gather stub's NaN placeholder blocks
d8d6a59 
The stub minted (tp_size-1) fresh all-NaN chunks per parameter per sync; they
are read-only constants, so one pooled block per (shape, dtype, device) serves
the run's lifetime. Recovers most of the comm-stub probe's 235B steady-sync
regression (the per-param cudaMalloc+fill bandwidth).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[checkpoint_engine] on-wire state verification sweep
3176158 
VERL_DELTA_VERIFY_EVERY=K appends the trainer's FULL export to every K-th
steady sync inside the same receive session (the steady bucket withholds
is_last; the sweep's final flush carries it). The receiver bit-compares each
copy_ destination -- the exact internal slice sglang's own load_weights
name-mapping resolves -- against what the server already holds before
overwriting, and raises on any mismatch. Zero mismatched elements proves the
delta-accumulated server state is bit-identical to the trainer's, end to end
through sglang's fusion/sharding, independent of generation or log-probs and
therefore usable on any model. Loader unit tests cover the pass and the
fail-loud paths.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[checkpoint_engine] attribute verify-sweep mismatches to parameter names
f399a2e 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[checkpoint_engine] policy knobs move to their owners (review)
0b3e974 
Snapshot pinning is the ENGINE's memory-profile call, not an environment
variable: prime_delta_snapshots takes an explicit pin flag, BaseEngine resolves
it from a class-level delta_pin_snapshots default, and MegatronEngine overrides
it to False (pinning the whole shard set on top of mcore's resident params and
per-step host traffic starves the node's cudaHostAlloc pool -- observed at
30B/235B). The verification sweep cadence likewise becomes the delta engine's
verify_every constructor kwarg (engine_kwargs.delta_sharded.verify_every)
instead of an env var. The full seed/verify export now fails loud on duplicate
HF tensor names -- two source params mapping to the same output silently
desyncs the diff base (observed: NemotronH A_log).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[checkpoint_engine] verify sweep judges idempotence, not per-copy equ…
3206f1c 
…ality

sglang loads some params through multi-stage loaders (mamba's composed loader
writes raw A_log, then overwrites with -exp in place); comparing at every
copy_ call flags raw-vs-transformed at each stage -- a guaranteed false
positive (the Nano9 A_log 2x128-per-layer signature). The sweep now snapshots
every copy_ destination before the load, runs the real load path untouched,
and bit-compares post-load state against the snapshot: replaying the trainer's
full state onto an in-sync server must be a no-op, whatever the loader does
in between.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[test] closed-loop sglang delta oracle (on-demand)
ac0612e 
perturb -> sparse delta flush -> engine.update_weights_from_tensor (production
load_format path) -> idempotence verify flush asserting zero mismatches, on a
tiny qwen2 / nemotron_h against a REAL sglang engine. On-demand
special_distributed script for sglang/bridge upgrades, not wired into per-PR
CI; the equivalent standing check for e2e jobs is the delta engine's
verify_every knob. Validated logic-wise via the loader CPU tests; full runs in
our enroot environment still fight infra quirks (kernel dims, offline caches).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@ChangyiYang @claude
[megatron] skip the probe entirely for empty deltas
4358b94 
The slot table freezes in slot_cache after each param's first probe, so a
zero-delta steady entry (frequent at small lr, late training, or with frozen
modules) needs only the cached slots and zero counts to stay in lockstep --
not a NaN buffer, a transform run and a full-output scan. First sync always
probes (cache miss), so the table is guaranteed before any short-circuit.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@wuxibin89 wuxibin89 changed the title [megatron] delta_sharded on Megatron-Bridge param mappings (TP+EP, hybrid-Mamba) [megatron] feat: delta_sharded on Megatron-Bridge param mappings (TP+EP, hybrid-Mamba) Jul 30, 2026
@ISEEKYAN

ISEEKYAN commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator

hi @ChangyiYang, do we have a protocol file defining the format of delta, I prefer a training backend agnostic protocol (like HF safetensor for full parameter) and difference backends have their impl

@ChangyiYang

ChangyiYang commented Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

hi @ChangyiYang, do we have a protocol file defining the format of delta, I prefer a training backend agnostic protocol (like HF safetensor for full parameter) and difference backends have their impl

hi @ISEEKYAN Yes, that's exactly how it works today: each training backend implements get_per_tensor_param_delta_shard, whose output format is the protocol — final HF-coordinate sparse entries, defined in verl/workers/engine/spec.py (wire encoding: delta_checkpoint_engine.py header). Everything past that boundary — gather, bucketing, the receiver — is backend-agnostic, and each impl fully controls how it produces its delta (snapshot diff, converters, or even diffing against a stored checkpoint).

@wuxibin89

wuxibin89 commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

@ChangyiYang Does pipeline parallelism supported?

@ChangyiYang

ChangyiYang commented Jul 31, 2026

Copy link
Copy Markdown
Contributor Author

@wuxibin89 PP is left as unsupported and will raise immediately. @gxlvera will support it

@ChangyiYang ChangyiYang mentioned this pull request Jul 31, 2026
Open
2 of 5 tasks
@wuxibin89
wuxibin89 merged commit 48ca9ee into verl-project:main Jul 31, 2026
95 of 127 checks passed
@ChangyiYang ChangyiYang mentioned this pull request Jul 31, 2026
Open
@attack204 attack204 mentioned this pull request Aug 7, 2026
Closed
6 of 8 tasks
@HaoWang113 HaoWang113 mentioned this pull request Sep 7, 2026
Draft
5 of 8 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@wuxibin89 wuxibin89 wuxibin89 approved these changes
@ISEEKYAN ISEEKYAN Awaiting requested review from ISEEKYAN ISEEKYAN is a code owner
@vermouth1992 vermouth1992 Awaiting requested review from vermouth1992 vermouth1992 is a code owner
@Luosuu Luosuu Awaiting requested review from Luosuu
@FoolPlayer FoolPlayer Awaiting requested review from FoolPlayer
@eric-haibin-lin eric-haibin-lin Awaiting requested review from eric-haibin-lin eric-haibin-lin is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ChangyiYang @ISEEKYAN @wuxibin89