fix(reshard): map global QKV intervals when KV heads are below TP - #650
fix(reshard): map global QKV intervals when KV heads are below TP#650KavinKrishnan wants to merge 1 commit into
Conversation
Computing num_kv_heads // tp_size gives zero for a model with fewer KV heads than trainer TP ranks -- Nemotron Ultra has 2 KV heads at TP8 -- so publishing a local KV head count either divided by zero or rejected the layout outright. The premise was wrong rather than the arithmetic. Megatron does not hand every TP rank a whole KV head. It slices the globally interleaved fused QKV tensor by raw rows, so most ranks legitimately own query rows and no K or V rows at all. Map each rank's real row interval instead of inventing a local head count. _QkvLayout derives the group geometry from the global head counts and bands() states Megatron's row order once, as a sequence of runs; the mapping is then an interval intersection, and a rank owning no K/V rows simply matches no K/V band. The sparse per-rank offers merge into complete tensors downstream. Requires the companion NeMo-RL change that publishes global, per-layer QKV geometry. Descriptors carrying only the old divisible local-head fields keep working through the legacy path, and half-specified metadata fails closed. Qualified on CPU and representative CUDA tensors: Q=64/KV=2/head_dim=128 at logical TP8 reconstructs byte-exact Q/K/V with no gaps or overlaps, divisible layouts stay byte-identical to the legacy path, 24Q/6KV/TP4 and heterogeneous per-layer geometry are covered, and BF16 CUDA same-weight and changed-weight refits match on parameters and projection outputs. Real Megatron TP8 to vLLM TP8 end-to-end is still outstanding. Split out of #635 so the interval mapping can be reviewed on its own. Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review. WalkthroughMegatron QKV aliasing now supports global interleaved head metadata, sparse KV ownership, validation, and legacy compatibility. New tests verify reconstruction, transfer planning, CUDA refits, logits, and fail-closed behavior. ChangesMegatron GQA QKV aliasing
Estimated code review effort: 4 (Complex) | ~60 minutes Merge Risk: 🔵 Low · up to The new global QKV mapping requires the companion emitter to provide complete compatible metadata; incomplete metadata can cause descriptor rejection rather than legacy fallback. The PR is mergeable with explicit owner awareness and follow-up on that integration contract. Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
|
Closing — folding this back into #635 rather than reviewing it separately. The split would have needed three separate approvals, and reviewer availability is the real constraint here, not diff size. One PR with a clear review guide is the faster path. Nothing is lost: the change is in #635 unchanged, and the readability concern that prompted the split is addressed there by giving the fused QKV row layout a name ( |
1 participant
The bug
Publishing a Megatron fused QKV tensor computed a local KV head count as
num_kv_heads // tp_size. Nemotron Ultra has 64 query heads but only 2 KV heads, so at TP8 that is2 // 8 = 0— either aZeroDivisionErrororinvalid local Q/KV head geometry, depending on the path.The arithmetic wasn't the problem. The premise was.
Why the premise was wrong
Megatron does not hand every TP rank a whole KV head. It slices the globally interleaved fused QKV tensor by raw rows. When KV heads are fewer than TP ranks, most ranks legitimately own query rows and no K or V rows at all. There is no local KV head count to compute, so no amount of fixing the division helps.
Megatron's row order is one block per KV group — that group's query rows, then its single K head, then its single V head — repeated per group. Q, K and V therefore interleave rather than forming three contiguous regions:
flowchart LR subgraph G0["KV group 0"] Q0["Q rows"] --> K0["K head"] --> V0["V head"] end subgraph G1["KV group 1"] Q1["Q rows"] --> K1["K head"] --> V1["V head"] end G0 --> G1The fix
Map each rank's real row interval through that layout instead of inventing a local head count.
flowchart LR E[Global per-layer Q/KV metadata] --> L[_QkvLayout] L --> B["bands(): projection runs<br/>in global row order"] S[This rank's raw fused-row interval] --> X[Interval intersection] B --> X X --> Q[q_proj shards] X --> K[k_proj shards] X --> V[v_proj shards]A rank that owns no K/V rows simply matches no K/V band and publishes Q only. The sparse per-rank offers merge into complete tensors downstream.
Review guide
Two files, and they are the whole change:
modelexpress_rl/train/engines/megatron/aliases.pytests/test_reshard_megatron_gqa.pyIn
aliases.py, the code is deliberately arranged so the layout is a named thing rather than index arithmetic inside a loop:_QkvLayout— holds the head counts, derives the group geometry.bands()— states Megatron's row order exactly once, as a sequence of runs. If our understanding of Megatron's layout is ever wrong, or upstream changes it, this is the single place to correct._read_qkv_layout— validates the published metadata; half-specified metadata fails closed rather than guessing._build_global_qkv_aliases— a plain interval intersection against those runs.Worth scrutinising: the per-group overlap math, that every source row is mapped exactly once (there's an explicit coverage assertion), that empty K/V tensors are omitted rather than published as zero-row tensors, and that divisible layouts still take the legacy path unchanged.
Compatibility
Descriptors carrying only the old divisible local-head fields keep working through the legacy path, so this is not a breaking change for existing topologies. The new path activates only when global head metadata is present.
Requires the companion NeMo-RL change that publishes global, per-layer QKV geometry: NVIDIA-NeMo/RL#3632. Neither half is useful alone — NeMo-RL publishes the geometry, ModelExpress maps the intervals.
Evidence, and what is still missing
Qualified on CPU and representative CUDA tensors:
Not yet covered: real Megatron TP8 to vLLM TP8 end-to-end for Q=64/KV=2, and full Nemotron Ultra qualification. I'd rather state that plainly than imply this is fully qualified.
Known limitation
This reverse-engineers Megatron's fused-QKV row order rather than obtaining it from a supported interface, so it is sensitive to upstream Megatron changes.
bands()exists to keep that assumption in one auditable place. Asking Megatron to expose the layout properly is the right long-term fix and is worth a follow-up.Test plan
PYTHONPATH=. python3 -m pytest tests/test_reshard_megatron_gqa.py -q— 15 passedContext
Split out of #635 at review request, to grow this surface gradually rather than land it all at once. Independent of the other split PRs — it touches only
aliases.pyand its test.Summary by CodeRabbit
New Features
Bug Fixes