test: pin AttentionResidualMixer against an independent mix (PRPUNDIT-22) - #1092
test: pin AttentionResidualMixer against an independent mix (PRPUNDIT-22)#1092jiagaoxiang wants to merge 2 commits into
Conversation
Closes PRPUNDIT-22. Independent RMS-score / unnormalised-value oracle, plus num_blocks=0 identity, hidden-size guard, and gradcheck.
There was a problem hiding this comment.
🟡 Changes recommended
The current gradcheck test does not actually validate gradients w.r.t. norm_weight/proj_weight because the lambda ignores the passed weight inputs.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR adds missing unit coverage for the Kimi K3 AttentionResidualMixer.forward implementation, aiming to lock in key correctness details (unnormalised mixing, rank-1 scorer factorisation, internal fp32+ compute with dtype cast-back) and close a test gap tracked in PRPUNDIT-22.
Changes:
- Add an independent oracle-style unit test suite for
AttentionResidualMixercovering correctness, dtype behavior, shape validation,num_blocks=0identity behavior, andgradcheck. - Add
__init__.pyto ensure thetests/unit_tests/backends/megatron/kimi_k3directory is treated as a package.
File summaries
| File | Description |
|---|---|
| tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py | New unit tests validating AttentionResidualMixer.forward against an independent reference and edge cases. |
| tests/unit_tests/backends/megatron/kimi_k3/init.py | Adds package marker and licensing header for the test directory. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # norm_weight / proj_weight are passed through so gradcheck perturbs the | ||
| # exact tensors the mixer reads (they are the same objects, not copies). | ||
| assert torch.autograd.gradcheck( | ||
| lambda ps, br, nw, pw: mixer(ps, br), | ||
| (prefix_sum, block_residual, mixer.norm_weight, mixer.proj_weight), | ||
| eps=1e-6, | ||
| atol=1e-4, | ||
| ) |
| mcore_module = pytest.importorskip( | ||
| "megatron.core.transformer.module", | ||
| reason="megatron.core is not importable in this environment", | ||
| ) |
pytest.importorskip's return value was never used; call it for its side effect (skip if megatron.core is unimportable) without binding it.
jiagaoxiang
left a comment
jiagaoxiang
left a comment
There was a problem hiding this comment.
TestIntel PR Steward: addressed the Copilot review findings — see inline replies.
| pytest.importorskip( | ||
| "megatron.core.transformer.module", | ||
| reason="megatron.core is not importable in this environment", | ||
| ) |
There was a problem hiding this comment.
TestIntel PR Steward: taken — fixed in 529ccbf by calling pytest.importorskip(...) without binding the result, matching your suggestion.
| (prefix_sum, block_residual, mixer.norm_weight, mixer.proj_weight), | ||
| eps=1e-6, | ||
| atol=1e-4, | ||
| ) |
There was a problem hiding this comment.
TestIntel PR Steward: skipping this one — the lambda's parameter list ignores nw/pw, but the tuple passed to gradcheck is (prefix_sum, block_residual, mixer.norm_weight, mixer.proj_weight), i.e. the exact same tensor objects the mixer reads internally (see the comment directly above the assert). gradcheck's numerical Jacobian perturbs each input tensor's storage in place, and its analytical Jacobian calls autograd.grad(output, inputs) against those same leaf objects — both work correctly by object identity even though the lambda's own argument names go unused. So the check does validate gradients w.r.t. norm_weight/proj_weight today. Happy to revisit if you have a concrete case where this breaks, but I don't want to change working gradient-check plumbing without one.
|
TestIntel PR Steward: follow-up on this run's findings. Copilot review (2 comments)
CI — This isn't caused by anything in this PR (only test files were touched) and isn't a plain flake either — it reads like the Docker Hub push step's PAT secret isn't available in this run's context (this PR runs from a fork branch), which is standard GitHub Actions behavior for I'll keep watching for review activity and CI updates, and will merge once there's a human approval and green (or maintainer-cleared) required checks. |
There was a problem hiding this comment.
🔵 Needs a closer look
The gradcheck currently ignores the nw/pw inputs and therefore doesn’t validate gradients w.r.t. the mixer’s scorer weights as intended.
Review details
Suppressed comments (1)
tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py:227
- The gradcheck lambda ignores the
nw/pwinputs and instead readsmixer.norm_weight/mixer.proj_weightfrom module state, so gradcheck won’t actually validate gradients w.r.t. the scorer weights (and may treat the unused inputs as having undefined grads). Use the backend call that takesnorm_weight/proj_weightexplicitly (or a stateless functional_call) so the perturbed inputs are the tensors used in the computation.
# norm_weight / proj_weight are passed through so gradcheck perturbs the
# exact tensors the mixer reads (they are the same objects, not copies).
assert torch.autograd.gradcheck(
lambda ps, br, nw, pw: mixer(ps, br),
(prefix_sum, block_residual, mixer.norm_weight, mixer.proj_weight),
- Files reviewed: 2/2 changed files
- Comments generated: 0 new
- Review effort level: Lite
This PR was created by dougljia via Test Gap Resolver.
Summary
tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py, the file the module docstring already claimed existed.num_blocks=0identity, hidden-sizeValueError, dtype cast-back, andgradcheck.Closes test gap PRPUNDIT-22.
Test plan
megatron.core:pytest tests/unit_tests/backends/megatron/kimi_k3/test_attention_residual.py(skipped on this host;megatron.corenot importable)black --check/isort --profile blackon the new files