test: cover fused clamped SwiGLU Triton autograd (PRPUNDIT-23) - #1106
Conversation
Closes PRPUNDIT-23. One GPU-gated suite for unweighted and weighted ClampedSwiGLUFunction paths; not the eager test_clamped_swiglu.py module.
Copilot review noted the fp8_input_store=True parametrized cases can hard-fail on GPUs/torch builds where float8 casting is unsupported, even when CUDA is available. Mirror the FP8 cast probe used in test_v4_fp8_indexer.py, but scope the skip to just the True case so fp8_input_store=False coverage still runs on non-FP8 devices.
Copilot flagged that _fp8_cast_works() probed float8_e4m3fn casting on a CPU tensor even though every parametrized case in this suite runs on device="cuda". On builds where FP8 casting is CUDA-only, the CPU probe would incorrectly report unsupported and skip valid fp8_input_store=True coverage. Probe on the CUDA device when available instead.
There was a problem hiding this comment.
🟡 Changes recommended
The new tests parameterize fp8_input_store but do not reliably validate the FP8 path for the unweighted forward case and never exercise fp8_input_store=True through backward.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds a new GPU-gated unit test suite to validate the Triton-backed DeepSeek-V4 clamped SwiGLU autograd wrappers in primus.backends.megatron.core.fusions.fused_bias_swiglu, covering both the unweighted (ClampedSwiGLUFunction) and weighted (ClampedWeightedSwiGLUFunction) paths.
Changes:
- Introduces
test_fused_clamped_swiglu.pywith forward/backward parity checks against an eager clamp+SiLU reference for unweighted and weighted clamped SwiGLU. - Adds routing tests to ensure
swiglu_impl/weighted_bias_swiglu_impldispatch to the clamped implementations whenclamp_valueis provided. - Adds FP8-cast gating logic for the
fp8_input_storeparameterization.
File summaries
| File | Description |
|---|---|
| tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py | Adds GPU-only parity/routing tests for Triton clamped SwiGLU autograd wrappers, including FP8-cast gating. |
Review details
Suppressed comments (2)
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:95
- The backward test hard-codes
fp8_input_store=False, so thefp8_input_store=Trueautograd path is never exercised in backward. This means regressions in the FP8-save logic (casting + restore in backward) would go undetected.
self.Fn.apply(y_fn, False, alpha).backward(grad_out)
_eager_clamped_swiglu(y_eager, alpha).backward(grad_out)
torch.testing.assert_close(y_fn.grad, y_eager.grad, atol=1e-5, rtol=1e-5)
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:144
- Same as the unweighted case: this backward test hard-codes
fp8_input_store=False, sofp8_input_store=Trueis never run through backward for the weighted autograd wrapper.
grad_out = torch.randn(6, 8, device="cuda", dtype=torch.float32)
self.Fn.apply(y_fn, w_fn, False, alpha).backward(grad_out)
_eager_clamped_weighted_swiglu(y_e, w_e, alpha).backward(grad_out)
torch.testing.assert_close(y_fn.grad, y_e.grad, atol=1e-5, rtol=1e-5)
torch.testing.assert_close(w_fn.grad, w_e.grad, atol=1e-5, rtol=1e-5)
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| def test_forward_matches_eager_reference(self, alpha: float, fp8_input_store: bool): | ||
| torch.manual_seed(2024) | ||
| y = (torch.randn(17, 64, device="cuda", dtype=torch.float32) * 5.0).requires_grad_(True) | ||
| out = self.Fn.apply(y, fp8_input_store, alpha) | ||
| torch.testing.assert_close(out, _eager_clamped_swiglu(y, alpha), atol=1e-5, rtol=1e-5) |
…backward Copilot flagged that fp8_input_store was parametrized but only forward output was checked against the eager reference, so the FP8-save/restore path was never actually exercised through backward, and the unweighted forward test never validated which dtype got saved. - Unweighted forward test now asserts the saved tensor dtype matches the flag, mirroring the weighted forward test. - Both backward tests are now parametrized over fp8_input_store. When True, the eager reference dequantizes through float8_e4m3fn first, mirroring the exact lossy round trip ClampedSwiGLUFunction.backward / ClampedWeightedSwiGLUFunction.backward perform, so the comparison stays tight (atol/rtol 1e-5) instead of masking the FP8 path behind a loosened tolerance.
|
TestIntel PR Steward: Addressed the Copilot review finding on
Also for the record on CI: the Will keep watching CI and for a human approval before merging. |
There was a problem hiding this comment.
🔵 Needs a closer look
The new parity suite is fp32-only and should be extended to cover the mixed-precision dtypes (at least bf16, optionally fp16) used by the production call sites to avoid missing dtype-specific regressions.
Review details
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:84
- These Triton parity tests only exercise fp32 inputs/outputs. The fused clamped SwiGLU kernels are called from the model in mixed precision (typically bf16/fp16), so limiting to fp32 can miss dtype-specific casting/rounding issues and contract mismatches. Consider parametrizing over at least {torch.float32, torch.bfloat16} (and torch.float16 if supported) with dtype-appropriate tolerances, similar to other fused-kernel parity suites under deepseek_v4.
This issue also appears on line 129 of the same file.
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:137
- The weighted clamped SwiGLU coverage is fp32-only as well. Since
weighted_bias_swiglu_implis used in MoE expert paths that usually run in bf16/fp16, consider adding dtype parametrization (and matching tolerances) so forward/backward parity is validated under the dtypes used in training/inference.
def test_forward_matches_eager_reference(self, alpha: float, fp8_input_store: bool):
torch.manual_seed(2024)
M, half = 5, 16
y = (torch.randn(M, 2 * half, dtype=torch.float32, device="cuda") * 5.0).requires_grad_()
weights = torch.rand(M, dtype=torch.float32, device="cuda") + 0.1
out = self.Fn.apply(y, weights, fp8_input_store, alpha)
torch.testing.assert_close(
out, _eager_clamped_weighted_swiglu(y, weights, alpha), atol=1e-5, rtol=1e-5
)
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
|
TestIntel PR Steward: Follow-up on the second Copilot review (pullrequestreview-5115815432, submitted 16:58 UTC, commit That review suggests extending the new parity suite from fp32-only to also cover bf16 (and optionally fp16) for both the unweighted and weighted forward/backward tests, since the production kernels run in mixed precision. Declining for this PR, scope-only: PRPUNDIT-23's gap was untested clamped-SwiGLU autograd correctness against the eager reference — it didn't call for dtype-precision coverage across the training dtype matrix. Adding bf16/fp16 would double the parametrization surface and need new dtype-aware tolerances, which is a materially different (and larger) ask than what was scoped here. That's consistent with the scope line already in the PR description ("weighted is not a second ticket") — same principle applies to dtype expansion. If mixed-precision parity for this kernel is wanted, it should be its own follow-up ticket rather than folded into this gap-closure PR. No code change from this feedback (the earlier, first Copilot review — fp8_input_store not exercised through backward — was already addressed in Status: |
Same-repo refile of #1093 so Primus-CI-TAS runs against
AMD-AGI/Primus(fork PRs do not receive Docker Hub credentials, sobuild-docker/ torch unit tests never ran).This PR was created by dougljia via Test Gap Resolver.
Summary
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.pyfor the production Triton wrappers infused_bias_swiglu.py.ClampedSwiGLUFunctionand weightedClampedWeightedSwiGLUFunction(forward, backward,fp8_input_store, impl routing) against an eager clamp/silu/mul reference.test_clamped_swiglu.py; the weighted path is dormant in shipped V4 config but stays in this shared suite. Weighted is not a second ticket.Closes test gap PRPUNDIT-23.
Test plan
pytest tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py(skipped on this host; no CUDA)black --check/isort --profile blackon the new file