Add qk_clip (MuonClip) patch for hybrid models and the distributed optimizer - #1164
vanshbhatia-amd wants to merge 5 commits into
Conversation
![pr-pundit-code-review[bot]](https://avatars.githubusercontent.com/u/16860860?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
✅ Passed
All 8 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, attention_mechanism, training, parallelism
Blockers
🔴 A reproducer test or command that fails before the qk_clip/distributed-optimizer patch is applied (e.g., a hybrid KDA/Mamba layer stack without a clip-capable self_attention triggering AttributeError, or a NaN eta denominator case)
This is a bug_fix PR; without a failing reproducer there is no demonstration that the described defects (crash on layers lacking self_attention, missing fp32 master-param shard offsets, NaN eta denominator) actually occurred and are addressed by the patch.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Confirmation that the reproducer passes after the fix (pytest output, log excerpt, or CI run for the new patch module)
No test signals were detected in the diff and no pytest/log output is attached, so there is no evidence the patched behavior is correct post-change.
Rule: bug_fix required: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
🔴 Multi-GPU or simulated multi-rank test of the distributed optimizer path exercising the new fp32 master-param shard offset bookkeeping (at least 2 GPUs, or a mocked multi-rank sharding test)
The PR is tagged parallelism and modifies the distributed optimizer's parameter sharding/offset accounting, which only manifests incorrectly when parameters are sharded across ranks; a single-rank run cannot detect wrong offsets.
Rule: parallelism required: Test with at least 2 GPUs or simulated multi-rank run
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description covering why qk_clip failed on hybrid (KDA/Mamba) stacks, why shard offsets for fp32 master params were required in the distributed optimizer, and why the eta denominator could become NaN
The summary lists the symptoms/changes but does not articulate the underlying mechanism (e.g., which Megatron source string is monkey-patched and why the original code assumed a uniform attention type), which is required for reviewers to validate a source-string patch that can silently drift with upstream Megatron versions.
Rule: bug_fix required: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Unit test that applies the source-string patch and asserts the patched Megatron source actually matches/replaces the expected upstream snippet (guard against silent no-op patching on version drift)
Source-string patching is brittle; without a test asserting the target string is found and substituted, the fix can silently stop applying after a Megatron-LM bump while tests and training still appear green.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🟡 Confirmation of no hangs or deadlocks on the relevant parallelism configurations (DP/TP/PP/EP combos used with the distributed optimizer and qk_clip)
qk_clip changes touch per-layer collective/skip logic; skipping layers without a clip-capable self_attention on some ranks but not others can cause rank divergence and collective hangs, so an explicit no-hang confirmation is required.
Rule: parallelism required: Confirmation no hangs or deadlocks on the relevant parallelism config
Source: repo evidence requirement
🟡 Small-model training correctness run showing loss/convergence before vs. after the qk_clip (MuonClip) change, with an explicit statement of whether outputs are numerically equivalent or intentionally changed
The PR is tagged training and alters clipping math (eta denominator guard) and which layers get clipped, which directly changes optimizer updates and therefore training outputs; a convergence/loss comparison is needed to show the fix does not regress or silently alter training dynamics.
Rule: training_correctness required: Loss curve or convergence check before/after; Correctness test on a small model run; Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
🟡 Focused unit test of the eta denominator NaN guard (zero/near-zero denominator and NaN input producing a finite, no-op clip factor)
Numerical guards are cheap to unit test and are the exact class of bug that silently regresses; without a targeted test the NaN path is unverified.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
✅ Passed
All 8 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, training, attention_mechanism, parallelism
Blockers
🔴 A reproducer test or command that fails before the qk_clip/distributed-optimizer patch is applied (e.g., a hybrid KDA/Mamba stack with MuonClip enabled raising an index/shape error or producing NaN eta)
This PR is categorized as a bug_fix; without a failing-before reproducer there is no proof the patched source strings actually address the reported qk_clip failure on hybrid layer stacks or the fp32 master-param shard offset bug.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Confirmation (pytest output, log excerpt, or CI run) that the reproducer passes after the fix
No test signals were detected in the diff and no logs/pytest output are attached, so the fix's effectiveness is unverified.
Rule: bug_fix required: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
🔴 Multi-rank / multi-GPU test exercising the distributed optimizer path with fp32 master param shard offsets (at least 2 GPUs or a simulated multi-rank run), covering the DP-sharded qk_clip weight update
The fix records shard offsets for fp32 master params in the distributed optimizer — a correctness issue that only manifests when parameters are sharded across data-parallel ranks and cannot be validated single-rank.
Rule: parallelism required: Test with at least 2 GPUs or simulated multi-rank run
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description covering why non-attention (KDA/Mamba) layers broke qk_clip, why shard offsets for fp32 master params were missing in the distributed optimizer, and why the eta denominator could reach zero/NaN
The summary states what changed but does not document the underlying failure mechanism; this matters especially because the change is a monkey-patch on Megatron-LM source strings, which is brittle against upstream versions.
Rule: bug_fix required: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Unit test asserting that qk_clip skips non-attention layers in a hybrid stack (i.e., layer-type dispatch correctness) and that the eta denominator floor prevents NaN for zero/degenerate max-logit inputs
The change alters attention-mechanism numerics (MuonClip qk scaling); a focused correctness test on a small model or synthetic tensors is required to show the clipping math is unchanged for attention layers and NaN-free at the boundary.
Rule: training_correctness required: Correctness test on a small model run
Source: repo evidence requirement
🟡 Loss curve or convergence check before/after the qk_clip change on a small training run
MuonClip qk_clip directly modifies weights during training; flooring the eta denominator and skipping layers changes training trajectories, so before/after convergence evidence is required.
Rule: training_correctness required: Loss curve or convergence check before/after
Source: repo evidence requirement
🟡 Explicit statement of whether training outputs are numerically equivalent to pre-patch behavior or intentionally changed (e.g., attention layers bit-identical, non-attention layers now untouched, eta floored only in degenerate cases)
Reviewers need to know whether this is a pure bug fix (no numeric drift for previously-working configs) or a deliberate numerical change.
Rule: training_correctness required: Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
🟡 Confirmation of no hangs or deadlocks on the relevant parallelism configs (DP-sharded distributed optimizer, plus TP/PP/EP combinations used with hybrid KDA/Mamba stacks)
Skipping non-attention layers inside a collective-participating qk_clip path risks rank divergence (some ranks skipping a collective), which is a classic hang source in hybrid/pipeline-parallel stacks.
Rule: parallelism required: Confirmation no hangs or deadlocks on the relevant parallelism config
Source: repo evidence requirement
Advisory
🔵 Regression/guard test that the Megatron-LM source-string patch still applies (patch target string matches) against the pinned Megatron-LM version
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
31486fb to
97836a4
Compare
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
1 architecture concern, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
Architecture Concerns
These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.
gate-workaround-removal-on-runtime-version-detection
When a downstream patch/monkey-patch exists to work around a bug in an external dependency (compiler, runtime, driver, ISA), and that bug is fixed in a specific upstream version, the removal of the workaround must be gated on a runtime detection of the dependency's version — not removed unconditionally nor left permanently in place. This preserves correctness on both old and new versions of the dependency within the same codebase.
Findings:
- This file installs three source-string rewrites against Megatron-LM (an external dependency) to compensate for upstream bugs in
optimizer/qk_clip.py,optimizer/distrib_optimizer.py, andtransformer/multi_latent_attention.py, butpatch_qk_clip_hybridapplies them purely based onargs.qk_clip/args.log_max_attention_logitwith no runtime detection of the Megatron-LM version. Add a runtime version query (e.g.importlib.metadata.version('megatron-core')ormegatron.core.__version__, read from the actually-installed package, not a build-time constant or env var) and gate each of the three sub-patches on the version range where the bug is known to be present. - Record, in code rather than only in the docstring, the upstream version window each anchor was validated against (e.g.
_QKCLIP_BUG_RANGE = ('0.x', '<0.y')) and skip that sub-patch with alog_rank_0warning once the installed Megatron-LM is at or past the version whereclip_qkhandles hybrid layers / the distributed-optimizer master shard. Otherwise users on a fixed Megatron will get either a double-applied rescale or a hard failure. - Replace the unconditional
assert _DISTOPT_ORI.strip() in source, "[qk_clip_hybrid] distrib_optimizer anchor not found"(and the equivalent anchor checks insidepatch_function_source/patch_method_source) with version-aware handling: if the installed dependency version is newer than the last known-buggy version and the anchor is absent, treat it as 'upstream already fixed' and no-op; only hard-fail when the version says the bug should still be there. As written, a Megatron bump that fixes the bug turns into a crash atbefore_train. - The three sub-patches target three independent upstream defects that may well be fixed in different Megatron-LM releases; split the single
_PATCH_KEYguard into per-spot guards each with its own runtime version predicate, so the hybrid-layer skip, themain_param_shard_startoffset, and the MLA eta/sharded-rescale fix can be retired independently as upstream lands them. - Avoid deriving the gate from anything other than the live runtime: do not key off a pinned submodule SHA, a hardcoded tuple, or a
PRIMUS_MEGATRON_VERSION-style env override; query the installedmegatron.coreat patch-application time so containers with a different Megatron than the one the repo was built against behave correctly.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
🟡 Changes recommended
Three critical and one moderate issue remain unresolved, with focused regression tests also requested.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds a gated Megatron source-rewrite patch for qk_clip compatibility with hybrid models and distributed optimizers.
Changes:
- Skips unsupported hybrid layers and resets log-only buffers.
- Tracks distributed master-parameter shard offsets.
- Rescales model and master weights with safer MLA handling.
File summaries
| File | Review findings |
|---|---|
primus/backends/megatron/patches/qk_clip_hybrid_patches.py |
Critical (2 votes): The clip_qk rewrite can generate duplicate else clauses and fail compilation. Critical (2 votes): Dedenting breaks the distributed-optimizer anchor, so shard offsets remain unset. Critical (1 vote): The strict-positive eta assertion rejects valid zero values from infinite logits. Moderate (1 vote): RoPE query rows are incorrectly scaled. Nit (1 vote): Add focused regression tests for hybrid layers, shard offsets, and master-weight handling. |
Review details
Suppressed comments (2)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:177
- The upstream block being replaced explicitly keeps the
qk_pos_emb_head_dimportion of the query projection unchanged, but this assignment scales those RoPE rows byeta. The corresponding positional key projection is not rescaled, so this changes the intended MLA qk-clipping behavior; leave the positional entries at one and applyeta ** alphaonly to the firstqk_head_dimrows.
flat = head_factor.expand(n, rows_per_head, cols).reshape(-1)
mp.data.mul_(flat[start : start + numel].to(mp.dtype).view_as(mp.data))
# q side: content (nope) part *= eta^alpha, rotary (pe) part *= eta
b_pe = self.config.qk_pos_emb_head_dim
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:171
- This new source rewrite has no focused tests for the hybrid layer skip/log-only reset, the nonzero distributed-optimizer shard offset, or the per-head master-weight factor slicing. The repository does test analogous source patches for anchor matching and idempotency (for example
tests/unit_tests/backends/megatron/test_gdn_rocm_gate_patches.py:43-79), so these high-risk runtime paths can regress or compile invalid rewritten code without detection. Add tests that exercise both a full master and a sharded master.
a = self.config.qk_head_dim
alpha = self.config.qk_clip_alpha
def _rescale(model_weight, rows_per_head, head_factor):
# head_factor: [n, rows_per_head, 1] fp32 multiplicative factor
w = model_weight.data
cols = w.numel() // (n * rows_per_head)
# (1) replicated bf16 model weight (full)
w.view(n, rows_per_head, cols).mul_(head_factor.to(w.dtype))
# (2) fp32 master: flat shard (distributed) or full 2-D (regular)
mp = getattr(model_weight, 'main_param', None)
if mp is not None:
start = int(getattr(model_weight, 'main_param_shard_start', 0))
- Files reviewed: 1/1 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if not log_max_only: | ||
| self_attn.clip_qk() | ||
| else: | ||
| # log-only: clip_qk() (which resets the buffer) isn't | ||
| # called, so reset here for a per-step (not running) max. | ||
| self_attn.core_attention.current_max_attn_logits = None""" |
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, attention_mechanism, training, parallelism
Blockers
🔴 A reproducer (unit test, script, or command) that demonstrates the qk_clip/MuonClip failure on a KDA/Mamba hybrid layer stack (and the distributed-optimizer main-param offset issue) before the fix
The PR is categorized as a bug_fix; no test, log excerpt, or reproducer command is present in the diff, so there is no way to verify the reported failure (e.g., index/shape mismatch when non-attention layers are traversed, or zero/negative eta denominator) actually occurs before the change.
Rule: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Confirmation (pytest output or training log excerpt) that the reproducer passes / the failure disappears after applying the patch
bug_fix PRs must show the post-fix pass evidence; the diff contains no test signals and no attached run logs proving the hybrid KDA/Mamba stack now skips non-attention layers correctly.
Rule: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description covering why the previous qk_clip indexing broke for hybrid layer stacks, why main param shard offsets must be recorded, and why the eta denominator needed flooring
The summary states what was changed but not the underlying failure mechanism (e.g., mismatch between attention layer indices and global layer indices, division by ~0 producing inf/NaN eta), which is required for bug_fix review.
Rule: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Correctness/convergence evidence: a small-model training run comparing loss curves before/after the qk_clip fix, plus an explicit statement of whether outputs are numerically equivalent or intentionally changed
Rescaling both bf16 and fp32 master weights and flooring the eta denominator directly alters optimizer state and attention weights, i.e., training numerics; without a loss/convergence check there is no evidence the change does not silently degrade or change training outputs.
Rule: Loss curve or convergence check before/after; Correctness test on a small model run; Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
🟡 Multi-GPU (>=2 rank) or simulated multi-rank run exercising the distributed optimizer path with main param shard offsets, confirming no hangs/deadlocks and correct shard-local rescaling across ranks
The PR is categorized as parallelism and modifies the Megatron distributed optimizer (param shard offsets used for qk_clip rescaling); single-process evidence cannot validate per-rank shard offset arithmetic, and no multi-rank test or log is provided.
Rule: Test with at least 2 GPUs or simulated multi-rank run; Confirmation no hangs or deadlocks on the relevant parallelism config
Source: repo evidence requirement
Advisory
🔵 A unit test asserting the source-string patch actually applies (target substring found / patched source compiles) against the pinned Megatron-LM version
Source-string patching silently no-ops when upstream text drifts; a guard test is the only cheap way to detect that the qk_clip/distributed-optimizer fix stopped being applied.
Rule: Reproducer command or test that fails before the fix
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
There was a problem hiding this comment.
🟡 Changes recommended
The installer anchor mismatch and unenforced FLA max-logit dependency leave required behavior unapplied; regression coverage is also missing.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (2)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:205
inspect.getsource(func)is dedented before this check, but_DISTOPT_ORIis written with the original 20-column indentation; after dedenting, the anchor cannot match (and.strip()only removes the first line's leading spaces). This assertion therefore aborts the installer before the distributed-optimizer and MLA patches are applied. Keep the raw source for the anchor replacement and dedent only after replacing it.
source = textwrap.dedent(inspect.getsource(func))
assert _DISTOPT_ORI.strip() in source, "[qk_clip_hybrid] distrib_optimizer anchor not found"
source = source.replace(_DISTOPT_ORI, _DISTOPT_NEW)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:79
- The current
FLAFlashAttentionimplementation only initializescurrent_max_attn_logitstoNone(primus/backends/megatron/core/transformer/fla_flash_attention.py:167-169) and never updates it. With this condition,clip_qk()immediately skips every FLA attention layer, so enablingqk_cliporlog_max_attention_logitsilently does nothing unless the separate #1162 change is guaranteed to be present. Include/gate the required max-logit accumulation or declare and enforce that dependency.
self_attn = getattr(transformer_layer, 'self_attention', None)
if self_attn is not None and hasattr(self_attn, 'clip_qk'):
if self_attn.core_attention.current_max_attn_logits is None:
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
| new_clip_qk = patch_function_source(qk_clip_mod, "clip_qk", _QKCLIP_ORI, _QKCLIP_NEW) | ||
| training_mod.clip_qk = new_clip_qk | ||
|
|
||
| # 2. distributed-optimizer shard offset. | ||
| _patch_distopt_classmethod() |
97836a4 to
47dbc02
Compare
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
✅ Passed
All 8 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, attention_mechanism, training, parallelism
Blockers
🔴 A reproducer test or command that fails before the fix (e.g., a unit test that builds a KDA/Mamba hybrid layer stack without a qk_clip-capable self_attention and shows the pre-patch qk_clip path raising AttributeError/IndexError, plus a case where the eta denominator is zero/NaN)
PR is categorized as bug_fix; without a failing-before reproducer there is no demonstration that the reported qk_clip/distributed-optimizer breakage is actually triggered and addressed by the source-string patch.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Confirmation the reproducer passes after the fix (pytest output, training log excerpt, or CI run showing qk_clip applied correctly on hybrid stacks and no NaN eta)
bug_fix category requires post-fix verification evidence; none is present in the diff or PR description.
Rule: bug_fix required: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
🔴 A multi-GPU or simulated multi-rank test exercising the distributed optimizer path with shard offsets for fp32 master params (at least 2 GPUs, DP/TP sharding enabled)
PR is categorized as parallelism and directly modifies distributed-optimizer shard-offset bookkeeping; single-process testing cannot validate per-rank shard mapping correctness.
Rule: parallelism required: Test with at least 2 GPUs or simulated multi-rank run
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description covering why hybrid (KDA/Mamba) layers lack self_attention.qk_clip attributes, why fp32 master params need shard offsets recorded in the distributed optimizer, and why the eta denominator can become NaN
The summary lists what changed but does not explain the underlying failure mechanism; source-string patching of Megatron-LM is fragile and needs explicit root-cause justification.
Rule: bug_fix required: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Confirmation of no hangs/deadlocks with the relevant parallelism config (e.g., distributed optimizer + MuonClip enabled, DP>1, with hybrid attention layers that skip qk_clip on some ranks/layers)
Skipping layers without a qk_clip-capable self_attention can create rank-divergent collective participation; a run log confirming completion without hang is needed.
Rule: parallelism required: Confirmation no hangs or deadlocks on the relevant parallelism config
Source: repo evidence requirement
🟡 Training-correctness evidence: short convergence/loss-curve comparison before vs after the patch on a small hybrid (KDA/Mamba) model with MuonClip enabled
qk_clip scaling and the eta NaN guard directly alter optimizer/attention numerics during training; a small-model loss check is required to show the fix restores (or intentionally changes) training behavior.
Rule: training_correctness required: Loss curve or convergence check before/after; Correctness test on a small model run
Source: repo evidence requirement
🟡 Explicit statement of whether outputs are numerically equivalent to pre-patch behavior for non-hybrid (pure attention) models, or intentionally changed
The patch rewrites Megatron-LM source strings that affect all models using qk_clip and the distributed optimizer, so regression risk for previously working configs must be explicitly addressed.
Rule: training_correctness required: Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
Advisory
🔵 A unit test verifying the source-string patch itself applies cleanly (i.e., the target substrings are found in the pinned Megatron-LM version and the patched module still imports/compiles)
String-based monkeypatching silently no-ops or breaks when upstream source drifts; a guard test detecting a failed match is standard for this patching pattern in this repo.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved compatibility, installation, correctness, and test-coverage issues remain.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (4)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:229
- This adds three
exec-based source rewrites plus shard-offset-dependent weight math, but the PR adds no tests that install the patch against the pinned Megatron revision or exercise hybrid-layer skipping, log-only buffer reset, and a non-zero partial master shard. Please add focused tests for anchor compilation/idempotency and the rescale factors on both full and offset shards; otherwise source drift or an offset/layout regression will only be detected during a training run.
new_clip_qk = patch_function_source(qk_clip_mod, "clip_qk", _QKCLIP_ORI, _QKCLIP_NEW)
training_mod.clip_qk = new_clip_qk
# 2. distributed-optimizer shard offset.
_patch_distopt_classmethod()
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:229
- The first rewrite is committed to
qk_clip_modbefore the distributed-optimizer and MLA rewrites run. If either later anchor/import fails,run_patchescan continue after logging the exception, leaving a live process with only part of the fix; a retry also cannot reapply this step because_QKCLIP_ORIis no longer present while the guard was never marked. Preflight all three anchors or make installation rollback/guard each target atomically before rebinding any of them.
new_clip_qk = patch_function_source(qk_clip_mod, "clip_qk", _QKCLIP_ORI, _QKCLIP_NEW)
training_mod.clip_qk = new_clip_qk
# 2. distributed-optimizer shard offset.
_patch_distopt_classmethod()
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:95
- This replacement ends before the upstream
elsebranch, but also adds its ownelsefor resetting the buffer. On Megatron versions whoseclip_qkalready contains the upstream log-only reset,str.replacematches only this prefix and leaves the originalelseafter the newly inserted one, producing invalid generated Python and aSyntaxErrorwhen the patch is installed. Match the complete conditional or select the replacement based on whether the reset branch is already present.
self_attn.core_attention.current_max_attn_logits = None"""
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:180
- The upstream helper in the replaced
_MLA_ORIblock explicitly leaves theqk_pos_emb_head_dimrows unchanged, but this assignment scales those q/rotary rows byeta. That changes MuonClip's clipping rule and shrinks the positional contribution even though the matching k positional projection is not rescaled; keep these rows at the factor-one initialization.
# q side: content (nope) part *= eta^alpha, rotary (pe) part *= eta
b_pe = self.config.qk_pos_emb_head_dim
q_factor = torch.ones(n, a + b_pe, 1, device=eta.device, dtype=torch.float32)
q_factor[:, :a, :] = torch.pow(eta, alpha)
q_factor[:, a:, :] = eta
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
| self_attn = getattr(transformer_layer, 'self_attention', None) | ||
| if self_attn is not None and hasattr(self_attn, 'clip_qk'): |
…timizer Upstream megatron-core's qk_clip assumes every transformer layer has a .self_attention and that the fp32 master weight is a full 2-D tensor. Both break for KDA/Mamba hybrid stacks and for the distributed optimizer, so enabling qk_clip there raises AttributeError or fails to rescale the master. Add a source-rewrite patch (gated on qk_clip / log_max_attention_logit) that fixes all three spots at runtime, keeping the Megatron-LM submodule pristine: - clip_qk: skip layers without a qk_clip-capable self_attention, and reset the per-head buffer in log-only mode (per-step, not running, max). - distrib_optimizer: record each shard's flat offset (main_param_shard_start) so clip_qk can locate the sharded fp32 master. - MLASelfAttention.clip_qk: floor the eta denominator (avoids eta**alpha = NaN for heads with max logit <= 0) and rescale both the DP-replicated bf16 weight and the fp32 master (flat shard under the distributed optimizer, else 2-D). With no qk_clip-capable attention present the paths are no-ops, so existing non-hybrid / non-qk_clip configs are unaffected. Co-authored-by: Cursor <cursoragent@cursor.com>
47dbc02 to
ed5a4c1
Compare
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
✅ Passed
All 8 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, attention_mechanism, training, parallelism
Blockers
🔴 A reproducer (pytest test, script invocation, or training command) that fails before the qk_clip/distributed-optimizer patch is applied — e.g., a hybrid KDA/Mamba stack where a layer lacks self_attention (AttributeError/crash), or a NaN-producing eta computation
This PR is explicitly a bug_fix to MuonClip qk_clip on hybrid attention stacks; without a failing-before reproducer there is no demonstration that the described bugs (missing self_attention, unrecorded shard offsets, NaN eta) actually occur or that the patch targets them.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Confirmation the reproducer passes after the fix (pytest output, log excerpt showing the hybrid model trains and eta is finite/clip applied)
No test signals of any kind were detected in the diff, so there is no post-fix pass evidence for the three distinct code paths changed (layer skipping, fp32 master-param shard offset recording, NaN guard).
Rule: bug_fix required: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
🔴 Multi-GPU / multi-rank test (at least 2 GPUs or simulated multi-rank) exercising the distributed optimizer path with qk_clip enabled, since the shard-offset recording for fp32 master params is only meaningful under distributed optimizer sharding
The PR is categorized parallelism and modifies distributed-optimizer master-parameter shard bookkeeping; a single-rank run cannot exercise the sharded offset logic, so correctness of the offset math is unverified.
Rule: parallelism required: Test with at least 2 GPUs or simulated multi-rank run
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description for each of the three fixes: why layers without self_attention were reached, why fp32 master params needed shard offsets recorded in the distributed optimizer, and what produced NaN in eta
The summary lists the symptoms/fixes but not the underlying cause (e.g., how the distributed optimizer flattens/shards master params such that qk_clip previously wrote to the wrong offsets). Source-rewrite patches against Megatron-LM are especially fragile and need documented root cause to justify the monkey-patch surface.
Rule: bug_fix required: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Confirmation of no hangs or deadlocks on the relevant parallelism configs (e.g., TP/PP/DP with distributed optimizer + qk_clip on a hybrid KDA/Mamba model), since layers are now conditionally skipped and could desynchronize collectives across ranks
Skipping layers without self_attention is rank/layer-dependent under pipeline parallelism; if any collective (e.g., all-reduce of qk max logits) is inside the skipped branch, ranks can diverge and hang. This must be explicitly validated.
Rule: parallelism required: Confirmation no hangs or deadlocks on the relevant parallelism config
Source: repo evidence requirement
🟡 Statement of whether training outputs are numerically equivalent or intentionally changed, plus a short convergence/loss check before vs. after on a small model
qk_clip directly rescales query/key weights and the NaN guard changes the eta value applied, so this alters training numerics. A loss-curve or small-model correctness comparison is needed to show the clip still takes effect (i.e., the NaN guard is not silently disabling clipping) and that convergence is unaffected.
Rule: training_correctness required: Loss curve or convergence check before/after; Correctness test on a small model run; Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
🟡 Focused unit test for the NaN guard in the eta computation asserting the fallback behavior (e.g., eta defaults to 1.0 / clip is skipped) when max logits are NaN/Inf, rather than propagating NaN into weights
A guard against NaN is a silent-failure-prone branch; without a targeted test it can regress unnoticed and either mask real divergence or disable qk_clip entirely.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
✅ Passed
All 8 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, attention_mechanism, training, parallelism
Blockers
🔴 A reproducer (test or command) that demonstrates the qk_clip/MuonClip failure on KDA/Mamba hybrid stacks and the distributed-optimizer master-param rescaling bug before the patch is applied
PR is categorized as bug_fix; no test signals were detected in the diff, so there is no failing-before reproducer showing that qk_clip incorrectly processed non-attention (KDA/Mamba) layers or mis-scaled master weights.
Rule: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Confirmation the reproducer passes after the fix (pytest output or training log excerpt showing qk_clip skipping non-attention layers and correct eta flooring/rescaling)
bug_fix evidence requires post-fix verification output; the diff contains no tests and no attached logs demonstrating the corrected behavior.
Rule: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description covering why non-attention layers were previously clipped, why the eta denominator needed flooring, and why fp32 master weights were previously missed
The summary describes what changed but not the underlying failure mechanism (e.g., division-by-zero/NaN from unfloored eta, shard-offset mismatch in the distributed optimizer) needed to justify a source-rewrite patch of Megatron internals.
Rule: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Unit test asserting that qk_clip skips layers without attention modules (KDA/Mamba hybrid layer list) and applies clipping only to attention layers
This is the core logic change (attention_mechanism category) and is unit-testable without GPUs by constructing a hybrid module list; no such test exists.
Rule: Correctness test on a small model run
Source: repo evidence requirement
🟡 Numerical correctness test on a small model run plus a statement of whether training outputs are numerically equivalent or intentionally changed (loss curve / convergence check before vs after)
The change rescales bf16 and fp32 master weights and floors the eta denominator, which directly alters optimizer state and attention logits — i.e., training outputs. No loss curve, convergence check, or equivalence statement is provided.
Rule: Loss curve or convergence check before/after; Correctness test on a small model run; Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
🟡 Multi-rank / ≥2-GPU test (or simulated multi-rank run) exercising the distributed optimizer master-param shard offset recording, with confirmation of no hangs or deadlocks
PR is categorized as parallelism and modifies the distributed optimizer's per-shard master-param offsets; correctness of shard offsets and cross-rank collectives can only be validated with a multi-rank run, and none is shown.
Rule: Test with at least 2 GPUs or simulated multi-rank run; Confirmation no hangs or deadlocks on the relevant parallelism config
Source: repo evidence requirement
Advisory
🔵 Test/verification that the Megatron source-rewrite patch applies cleanly (patch target strings still match upstream Megatron source) and is idempotent/no-op when the targeted code is absent
Source-rewrite patch modules silently break on upstream drift; a guard test asserting the patch applied (or raised a clear error) protects against silent loss of the fix.
Rule: Reproducer command or test that fails before the fix
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved correctness, compatibility, patching, and test-coverage issues remain.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (4)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:79
DeepseekV4Attentionalso subclassesMLASelfAttentionfor type identity, so it inheritsclip_qkeven though its V4 layout has nolinear_kv_up_projand its eager/HCA/CSA variants can havecore_attention = None. Thishasattrcheck therefore treats V4 as qk-clip-capable and can dereferenceNonehere (or fail later inclip_qk) when the flag is enabled. Check the required MLA projection/core attributes, or explicitly exclude V4, before entering this block.
if self_attn is not None and hasattr(self_attn, 'clip_qk'):
if self_attn.core_attention.current_max_attn_logits is None:
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:166
- This file adds three dynamic source rewrites and several correctness-sensitive branches, but no regression tests cover them. Existing source-patch modules have anchor/idempotency tests; add coverage for a mixed hybrid layer list, log-only buffer reset, a non-zero shard offset, and the <=0-logit rescale path so upstream layout drift cannot silently leave the model/master weights inconsistent.
def _rescale(model_weight, rows_per_head, head_factor):
# head_factor: [n, rows_per_head, 1] fp32 multiplicative factor
w = model_weight.data
cols = w.numel() // (n * rows_per_head)
# (1) replicated bf16 model weight (full)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:229
- If
_patch_distopt_classmethod()or the MLA rewrite raises after the first rewrite succeeds,qk_clip_modandtraining_mod.clip_qkhave already been rebound but_PATCH_KEYis not marked. The patch runner normally continues, leaving a partially applied feature; a later retry also cannotinspect.getsourcethe exec-created function. Apply these rewrites transactionally and restore the original bindings on failure, asmlp_fla_swiglu_patches.pydoes.
new_clip_qk = patch_function_source(qk_clip_mod, "clip_qk", _QKCLIP_ORI, _QKCLIP_NEW)
training_mod.clip_qk = new_clip_qk
# 2. distributed-optimizer shard offset.
_patch_distopt_classmethod()
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:72
- This anchor stops immediately after the
if not log_max_onlybody. Current Megatron-LM already has anelsereset branch after this line; replacing only this prefix with_QKCLIP_NEWleaves the originalelseafter the new block's ownelse, producing an invalidelse ... elseand aSyntaxErrorwhenexecruns. Replace the complete upstream branch or support both source variants explicitly.
transformer_layer.self_attention.clip_qk()"""
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
The fp32 master is owned by the fused optimizer (no main_param), so the rescale would only touch the bf16 weight and be overwritten next step. Fail fast instead of silently under-clipping; log-only mode is unaffected.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
✅ Passed
All 8 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, attention_mechanism, training, parallelism
Blockers
🔴 A reproducer test or command that fails before the qk_clip/hybrid-attention fix (e.g., a unit test constructing a hybrid stack with layers lacking self_attention, or a NaN-producing eta computation) and demonstrably errors/produces NaN without the patch
This PR is explicitly a bug_fix to MuonClip/qk_clip for hybrid (KDA/Mamba) stacks and the distributed optimizer; without a failing reproducer there is no proof the described defects (AttributeError on missing self_attention, wrong fp32 master-param shard offsets, NaN eta) actually existed or are addressed.
Rule: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Pytest output or log excerpt confirming the reproducer now passes after applying the patch module
No test signals were detected in the diff at all, so there is no post-fix confirmation for any of the three separate fixes bundled here.
Rule: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
🔴 Multi-GPU (>=2 GPU) or simulated multi-rank run exercising the distributed optimizer path where shard offsets for fp32 master params are recorded and consumed by qk_clip
The distributed-optimizer shard-offset change only manifests when parameters are sharded across DP ranks; a single-rank run cannot exercise the offset arithmetic, and an off-by-one offset would corrupt weights silently on multi-rank configs.
Rule: Test with at least 2 GPUs or simulated multi-rank run
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description for each of the three fixes (layer skipping, shard offset recording for fp32 master params, NaN guard in eta)
The summary lists what changed but not why the original Megatron-LM source path breaks for hybrid attention stacks or how the missing shard offsets caused incorrect qk_clip scaling under the distributed optimizer; a source-rewrite patch needs the upstream defect articulated so reviewers can validate the rewrite stays in sync.
Rule: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Correctness check on a small model run showing qk_clip behavior before/after (loss curve or convergence check), plus an explicit statement of whether training outputs are numerically equivalent or intentionally changed
qk_clip/MuonClip directly rescales attention query/key weights during optimization, so this patch changes training numerics. The NaN guard in eta in particular silently alters update magnitudes; no loss curve, convergence check, or equivalence statement is provided.
Rule: Loss curve or convergence check before/after; Correctness test on a small model run; Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
🟡 Confirmation of no hangs or deadlocks on the relevant parallelism configurations (DP with distributed optimizer, and any TP/PP config where some pipeline stages contain only KDA/Mamba layers with no self_attention)
Skipping layers without
self_attentioncan make ranks/stages take divergent code paths; if the qk_clip path contains any collective, unbalanced participation across ranks causes a hang. This needs explicit confirmation on a hybrid stack layout.
Rule: Confirmation no hangs or deadlocks on the relevant parallelism config
Source: repo evidence requirement
🟡 Focused unit test asserting the layer-skipping predicate for hybrid stacks: modules lacking self_attention (KDA/Mamba blocks) are skipped and attention layers are still clipped, covering a mixed module list
This is the core attention-mechanism behavior change of the PR and is trivially unit-testable with a fake module list; no such test exists in the diff.
Rule: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🟡 Unit test for the NaN guard in the eta computation, asserting that NaN/Inf inputs produce a safe (non-NaN) scaling factor and that normal inputs are unchanged relative to the pre-patch formula
A NaN guard is a pure-function numerical change that silently alters updates; without a test pinning both the guarded and unguarded branches, regressions or over-broad clamping would go undetected.
Rule: Correctness test on a small model run
Source: repo evidence requirement
Advisory
🔵 Drift/sync test or assertion that the vendored Megatron-LM source rewrite matches the upstream function it patches (e.g., signature/version guard)
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
There was a problem hiding this comment.
🟡 Changes recommended
Critical and moderate review findings remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (4)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:266
- Raising
NotImplementedErrorhere does not actually reject this combination:TrainRuntimeinvokesrun_patcheswith its defaultstop_on_error=False, and the runner catches handler exceptions and continues (primus/core/runtime/train_runtime.py:222-234,primus/core/patches/patch_runner.py:124-140). Therefore qk_clip + precision-aware optimizer proceeds without installing this patch and can still under-clip/stale the master weight. Move this validation to an argument-validation path or otherwise make it fail before training starts.
if getattr(args, "qk_clip", False) and getattr(args, "use_precision_aware_optimizer", False):
raise NotImplementedError(
"qk_clip is not supported together with use_precision_aware_optimizer: "
"the fp32 master is owned by the fused optimizer and cannot be rescaled "
"by this patch. Disable one of the two."
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:180
- The upstream fragment immediately above says to keep the
qk_pos_emb_head_dimportion unchanged, but this assignment multiplies the rotary/PE Q rows byeta. The corresponding K positional projection is not inlinear_kv_up_proj, so this changes the relative positional contribution to the logits rather than only applying the intended content rescaling. Leave the initialized PE factor at one and remove this assignment.
# q side: content (nope) part *= eta^alpha, rotary (pe) part *= eta
b_pe = self.config.qk_pos_emb_head_dim
q_factor = torch.ones(n, a + b_pe, 1, device=eta.device, dtype=torch.float32)
q_factor[:, :a, :] = torch.pow(eta, alpha)
q_factor[:, a:, :] = eta
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:79
- Checking only for an inherited
clip_qkmethod does not establish that this attention layout supports it.DeepseekV4AttentionsubclassesMLASelfAttention, so it passes this test, but its implementation has nolinear_kv_up_proj(seeprimus/backends/megatron/core/transformer/deepseek_v4_attention.py:468-477); enablingqk_clipthen reachesself_attn.clip_qk()and raisesAttributeError. Gate this branch on the projections/core-attention buffer that the patched method requires, or use an explicit capability flag.
if self_attn is not None and hasattr(self_attn, 'clip_qk'):
if self_attn.core_attention.current_max_attn_logits is None:
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:166
- Please add focused tests for this new source-rewrite and rescaling path. There is currently no coverage for the hybrid-layer skip, the non-positive-logit eta floor, or slicing the fp32 master by
main_param_shard_start; an incorrect rewrite or shard factor would only surface during distributed training. Similar source patches in this repository already test anchors and idempotence, so this patch needs equivalent tests for its numerical behavior.
def _rescale(model_weight, rows_per_head, head_factor):
# head_factor: [n, rows_per_head, 1] fp32 multiplicative factor
w = model_weight.data
cols = w.numel() // (n * rows_per_head)
# (1) replicated bf16 model weight (full)
- Files reviewed: 1/1 changed files
- Comments generated: 2
- Review effort level: Lite
run_patches swallows exceptions (stop_on_error=False), so a failed install would silently fall back to upstream and either crash on hybrid layers or leave the master unscaled. Escalate incompatibility/install failures to a hard abort. Also restrict the clip dispatch to MLA attention with linear_kv_up_proj, skipping standard SelfAttention (whose clip_qk mishandles the sharded master) and MLA variants without it.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
✅ Passed
All 8 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, attention_mechanism, training, parallelism
Blockers
🔴 A reproducer (unit test, script, or command) that demonstrates the qk_clip/MuonClip failure on a KDA/Mamba hybrid model (non-attention layers present) and fails before the patch is applied
This is categorized as a bug_fix; without a failing-before reproducer there is no proof the described crash/incorrect-scaling on non-attention layers actually occurs or that skipping them is the correct remedy.
Rule: bug_fix.required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Confirmation (pytest output, training log excerpt, or CI run) that the reproducer passes after applying the Megatron-LM source-rewrite patch
No test signals of any kind were detected in the diff; a source-rewrite patch against upstream Megatron-LM is especially fragile (it can silently no-op if the upstream source text drifts), so post-fix verification output is required.
Rule: bug_fix.required: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
🔴 Correctness check on a small model run showing qk_clip rescaling produces the intended effect on both bf16 and fp32 master weights (e.g., assert scaled weight values / q-k logit norms before and after clip)
The patch directly mutates master/optimizer weights during training, so it changes training outputs; a numerical correctness check on a small model is required to show the rescale is applied to the right shard slice and dtype.
Rule: training_correctness.required: Correctness test on a small model run
Source: repo evidence requirement
🔴 Multi-GPU (>=2 rank) or simulated multi-rank run exercising the distributed optimizer path with main param shard offsets, covering the parallelism configs (TP/PP/DP-shard) the patch touches
The change records and uses main param shard offsets inside the distributed optimizer, which is inherently rank-dependent; a single-process run cannot validate that each rank rescales its own shard slice correctly.
Rule: parallelism.required: Test with at least 2 GPUs or simulated multi-rank run
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description covering why qk_clip mis-handled KDA/Mamba layers and why main param shard offsets were previously unavailable/incorrect in the distributed optimizer
The summary states what changed (skip non-attention layers, record shard offsets, rescale bf16/fp32 master weights) but does not explain the underlying failure mechanism (e.g., index misalignment, wrong shard slice, dtype mismatch) that makes these changes necessary.
Rule: bug_fix.required: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Loss curve / convergence comparison before and after the patch for a KDA or Mamba hybrid model with MuonClip enabled
qk_clip modifies weights every step; without a before/after convergence check there is no evidence that skipping non-attention layers and rescaling master weights does not destabilize or silently degrade training.
Rule: training_correctness.required: Loss curve or convergence check before/after
Source: repo evidence requirement
🟡 Explicit statement of whether training outputs are numerically equivalent to the pre-patch behavior or intentionally changed (and for which model families)
The fix intentionally alters which layers get clipped and how master weights are scaled; reviewers need an explicit equivalence/intentional-divergence statement to judge impact on existing attention-only models.
Rule: training_correctness.required: Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
🟡 Confirmation of no hangs or deadlocks on the relevant distributed-optimizer / hybrid-model parallelism configuration (e.g., completed N-step multi-rank run log)
Conditionally skipping non-attention layers can cause rank-divergent control flow around collective operations in the distributed optimizer, a classic deadlock source; a completed multi-rank run log is needed.
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved critical correctness issues and moderate reliability risks remain.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (4)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:88
- On the current branch, the default hybrid spec selects
FLAFlashAttention, whosecurrent_max_attn_logitsis alwaysNone; thiscontinuetherefore makes both qk clipping and log-only mode no-ops on the default FLA path. PR #1162 adds the producer, but it is a separate open PR, so this patch should require/include that capability or fail clearly instead of silently skipping the requested behavior.
if self_attn.core_attention.current_max_attn_logits is None:
continue
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:182
- For every distributed-optimizer clip,
expand(...).reshape(-1)materializes a full FP32 per-element factor for the entire projection before taking this rank's shard. On large MLA projections that transient can be as large as the full weight on every DP rank, defeating the memory benefit of sharding and potentially causing OOM during clipping. Compute the factor only for the[start:start + numel]shard (or apply row-wise factors without constructing the full flattened tensor).
flat = head_factor.expand(n, rows_per_head, cols).reshape(-1)
mp.data.mul_(flat[start : start + numel].to(mp.dtype).view_as(mp.data))
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:235
- This adds three version-sensitive source rewrites and new tensor/shard arithmetic, but no tests exercise them. The analogous
test_mlp_fla_swiglu_patches.pycovers source anchors, installation, idempotence, and rollback; without equivalent tests here, regressions in hybrid-layer skipping, log-only reset, or sharded-master rescaling will only surface during distributed training. Add focused unit tests for those paths and the anchor/installation failure cases.
new_clip_qk = patch_function_source(qk_clip_mod, "clip_qk", _QKCLIP_ORI, _QKCLIP_NEW)
training_mod.clip_qk = new_clip_qk
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:212
- Unlike
_source_patch_utils.patch_method_source_multi, this anchor check usesassert, so running Python with-Oremoves the guard. If the upstream method drifts in that mode,raw.replace(...)leaves it unchanged, the patch still marks itself installed, and sharded masters later default to offset 0 and can be rescaled with the wrong factors. Use an explicit runtime check (and ideally require exactly one match) before recompiling.
assert _DISTOPT_ORI in raw, "[qk_clip_hybrid] distrib_optimizer anchor not found"
- Files reviewed: 1/1 changed files
- Comments generated: 2
- Review effort level: Lite
| b_pe = self.config.qk_pos_emb_head_dim | ||
| q_factor = torch.ones(n, a + b_pe, 1, device=eta.device, dtype=torch.float32) | ||
| q_factor[:, :a, :] = torch.pow(eta, alpha) | ||
| q_factor[:, a:, :] = eta |
The linear_kv_up_proj gate was too broad: it disabled qk_clip for every standard/GQA SelfAttention (which works fine outside the distributed optimizer). Restore the minimal hybrid-skip so non-MLA attention keeps upstream clip_qk; MLA still routes to the offset-aware patched version.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 1 file reviewed.
Rules checked: 8 (8 passed, 0 failed).
✅ Passed
All 8 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
PR Pundit — Test Coverage ❌
Verdict: Missing | Categories: bug_fix, attention_mechanism, training, parallelism
Blockers
🔴 A reproducer test or command that fails before the qk_clip / distributed-optimizer patch is applied (e.g., a unit test constructing a KDA/Mamba hybrid layer stack where qk_clip previously indexed non-attention layers, or a case where the eta denominator becomes 0/NaN)
This PR is explicitly a bug_fix to MuonClip qk_clip and the distributed optimizer; the bug_fix evidence policy requires a failing-before reproducer so the regression cannot silently return. The diff contains no test files or reproduction script at all.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
🔴 Confirmation (pytest output, log excerpt, or CI run) that the reproducer passes after the fix
No test output, log excerpt, or CI evidence is attached showing the patched qk_clip skips non-attention (KDA/Mamba) layers correctly and that the NaN-guarded eta path produces finite values.
Rule: bug_fix required: Confirmation it passes after the fix (pytest output or log excerpt)
Source: repo evidence requirement
🔴 Multi-rank / distributed test (>=2 GPUs or simulated multi-rank) exercising the distributed optimizer path that records shard offsets for fp32 master params, confirming shards map to the correct parameter slices and that qk_clip applies consistently across ranks
The PR is tagged parallelism and directly modifies the Megatron distributed optimizer's master-param sharding; single-process testing cannot catch incorrect per-rank offsets or cross-rank divergence in clipped weights.
Rule: parallelism required: Test with at least 2 GPUs or simulated multi-rank run
Source: repo evidence requirement
Strong suggestions
🟡 Root cause explanation in the PR description covering (a) why qk_clip mis-handled KDA/Mamba hybrid layers, (b) why fp32 master param shard offsets were previously wrong/absent in the distributed optimizer, and (c) under what condition the eta denominator became NaN
The summary states what was changed but not why the original Megatron-LM code failed; for a source-rewrite/monkey-patch of upstream Megatron, the root cause must be documented so the patch can be re-validated or dropped on upstream version bumps.
Rule: bug_fix required: Root cause explanation in PR description
Source: repo evidence requirement
🟡 Confirmation of no hangs/deadlocks on the relevant parallelism configuration (DP with distributed optimizer, plus TP/PP as used by the KDA/Mamba hybrid model) after the patch
Skipping non-attention layers inside a collective-bearing optimizer/clip routine risks rank-divergent control flow (some ranks entering a collective others skip), which manifests as a hang; the PR provides no run evidence on a real parallel config.
Rule: parallelism required: Confirmation no hangs or deadlocks on the relevant parallelism config
Source: repo evidence requirement
🟡 Correctness evidence on a small KDA/Mamba hybrid model run: loss/convergence curve before vs after, plus a statement of whether outputs are numerically equivalent or intentionally changed
qk_clip (MuonClip) directly rescales attention query/key weights during training and the eta guard changes the clip magnitude, so this alters training numerics; the PR must show a short convergence check and declare whether the change is intentionally non-equivalent (it fixes wrong clipping) to distinguish a fix from a regression.
Rule: training_correctness required: Loss curve or convergence check before/after; Correctness test on a small model run; Statement of whether outputs are numerically equivalent or intentionally changed
Source: repo evidence requirement
Advisory
🔵 Focused unit test for the eta NaN/zero-denominator guard asserting finite output and unchanged behavior when the denominator is well-conditioned
The NaN guard is a small, purely numerical branch that is trivially unit-testable without GPUs; without it there is no protection against the guard being removed or mis-thresholded in a future refactor.
Rule: bug_fix required: Reproducer command or test that fails before the fix
Source: repo evidence requirement
PR Pundit test coverage — grounded in this repo's merged PR history.
There was a problem hiding this comment.
🔵 Needs a closer look
Four moderate findings remain involving safety checks, memory usage, test coverage, and standard-attention distributed clipping.
Review details
Suppressed comments (4)
Previously missed (1) — in code that hasn't changed since the last review.
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:177
head_factor.expand(...).reshape(-1)materializes a full FP32 factor for every element of the projection before the code slices out this rank's master shard. For a large MLA q projection this can add hundreds of MB of temporary memory on each rank at the optimizer-step boundary, defeating the benefit of the sharded master and risking OOM; apply the factor by head/row chunks or slice the shard without expanding the full projection.
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:207
- The installation path claims to hard-fail on an upstream mismatch, but this anchor check uses
assert. Underpython -Othe check is removed, so a changed distributed-optimizer body is recompiled unchanged and the patch is marked installed; the same problem exists in the sharedpatch_function_source/patch_method_sourceassertions used below. Replace these guards with explicitRuntimeErrorchecks before relying on theSystemExitsafety fallback.
assert _DISTOPT_ORI in raw, "[qk_clip_hybrid] distrib_optimizer anchor not found"
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:235
- This adds three version-sensitive source rewrites and the correctness-critical shard/factor logic without a corresponding unit test. Existing analogous source-patch tests cover anchor matching, installation, idempotency, and conditions (for example
tests/unit_tests/backends/megatron/test_gdn_rocm_gate_patches.py:43-111), but there is no coverage here for skipping hybrid layers, log-only reset, nonzero shard offsets, or the fp32/bf16 scaling parity. Add focused tests before relying on this in distributed training.
# 1. clip_qk: rebind the module attr and the name already imported by training.
new_clip_qk = patch_function_source(qk_clip_mod, "clip_qk", _QKCLIP_ORI, _QKCLIP_NEW)
training_mod.clip_qk = new_clip_qk
# 2. distributed-optimizer shard offset.
_patch_distopt_classmethod()
# 3. MLASelfAttention.clip_qk (inherited unchanged by PrimusMLASelfAttention).
patch_method_source(MLASelfAttention, "clip_qk", _MLA_ORI, _MLA_NEW)
primus/backends/megatron/patches/qk_clip_hybrid_patches.py:79
- This deliberately leaves standard
SelfAttention.clip_qkon the upstream implementation, butclip_qk()dispatches every attention object that exposes that method. In the pinned Megatron-LM, standard attention also applies_clip_linear_qkvtolinear_qkv.weight.main_param; with the distributed optimizer that tensor is only a flat shard, so the full-weight reshape/head factors are wrong (or fail) and standard-attention qk clipping remains broken. Patch that path too, or explicitly narrow/reject distributedqk_clipfor non-MLA attention.
# do have a ``clip_qk``-capable attention keep upstream dispatch
# (MLA routes to the offset-aware clip_qk patched below; standard
# attention keeps its own upstream clip_qk unchanged).
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
Upstream megatron-core's
qk_clip(MuonClip attention-logit clipping) assumes every transformer layer has a.self_attentionand that the fp32 master weight is a full 2-D tensor. Both break for KDA/Mamba hybrid stacks and for the distributed optimizer, so enablingqk_clipon those configs either raisesAttributeErroror silently fails to rescale the master weight.Rather than modifying the Megatron-LM submodule, this adds a source-rewrite patch (same style as
mamba_fused_ce_patches), gated onqk_clip/log_max_attention_logit, that fixes all three spots at runtime:clip_qk- skip layers without aqk_clip-capableself_attention(hybrid models interleave attention with linear/Mamba/KDA layers), and reset the per-head buffer in log-only mode so it reports a per-step (not running) max.distrib_optimizer._build_model_and_main_param_groups- record each shard's flat offset (main_param_shard_start) soclip_qkcan locate the sharded fp32 master.MLASelfAttention.clip_qk- floor theetadenominator (a head with max logit <= 0 would give a negative base andeta ** alpha == NaN) and rescale both the DP-replicated bf16 weight and the fp32 master (flat shard under the distributed optimizer, or full 2-D otherwise).With no
qk_clip-capable attention present the paths are no-ops, so non-hybrid / non-qk_clipconfigs are unaffected. Complements #1162, which exposes the per-head max attention logit on the FLA fused backend thatclip_qkconsumes.