test: cover fused clamped SwiGLU Triton autograd (PRPUNDIT-23) - #1093
test: cover fused clamped SwiGLU Triton autograd (PRPUNDIT-23)#1093jiagaoxiang wants to merge 3 commits into
Conversation
Closes PRPUNDIT-23. One GPU-gated suite for unweighted and weighted ClampedSwiGLUFunction paths; not the eager test_clamped_swiglu.py module.
There was a problem hiding this comment.
🟡 Changes recommended
The new fp8_input_store=True test variants are not gated on float8-cast support and can hard-fail on CUDA/HIP devices/builds without FP8 casting enabled.
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 DeepSeek-V4 Triton clamped-SwiGLU autograd wrappers in primus.backends.megatron.core.fusions.fused_bias_swiglu, comparing forward/backward behavior to an eager clamp/silu/mul reference and checking impl routing.
Changes:
- Introduces
test_fused_clamped_swiglu.pycoveringClampedSwiGLUFunctionandClampedWeightedSwiGLUFunction. - Adds forward/backward correctness checks against eager references plus routing checks for
swiglu_impl/weighted_bias_swiglu_impl. - Adds an
fp8_input_storeforward-path check (saved tensor dtype) for the weighted variant.
File summaries
| File | Description |
|---|---|
| tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py | New GPU-gated tests for Triton clamped SwiGLU autograd wrappers (weighted + unweighted), including routing validation. |
Review details
Suppressed comments (2)
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:50
- This parametrize currently always runs fp8_input_store=True when CUDA is available, which can hard-fail on devices without float8 cast support. Use a gated fp8_input_store case list so only the supported configuration runs fp8_input_store=True.
@pytest.mark.parametrize("fp8_input_store", [False, True])
def test_forward_matches_eager_reference(self, alpha: float, fp8_input_store: bool):
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:89
- Same as above: fp8_input_store=True should be skipped when float8 casts are unsupported on the current build/device; otherwise this test can fail even though the kernels themselves are unrelated to the fp8 storage option.
@pytest.mark.parametrize("fp8_input_store", [False, True])
def test_forward_matches_eager_reference(self, alpha: float, fp8_input_store: bool):
- 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.
| import pytest | ||
| import torch | ||
| import torch.nn.functional as F | ||
|
|
||
| cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="clamped SwiGLU Triton kernels need CUDA/HIP") |
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.
|
TestIntel PR Steward: Follow-up on this poll: CI ( Copilot review ( I'll keep following this PR for CI status, further review comments, and merge-readiness (open, green required checks, human approval, no outstanding change requests). |
There was a problem hiding this comment.
🟡 Changes recommended
The FP8 gating currently probes casting on CPU (potentially skipping valid CUDA FP8 coverage) and the new backward tests don’t exercise the fp8_input_store=True backward path.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
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:90
- The backward test only exercises fp8_input_store=False, so the fp8_input_store=True backward path (saving input as float8 and casting back in backward) is currently untested. Parameterize this test over fp8_input_store and at least assert gradients are finite for the FP8-storage case.
This issue also appears on line 126 of the same file.
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:130
- This backward test only covers fp8_input_store=False. Since ClampedWeightedSwiGLUFunction has a distinct fp8_input_store=True path (stores input in float8 but still needs to backprop through both input and weights), add fp8_input_store to the parametrization and at least assert grads are finite in the FP8-storage case.
def test_backward_matches_eager_autograd(self):
torch.manual_seed(8)
alpha = 7.0
y_fn = (torch.randn(6, 16, device="cuda", dtype=torch.float32) * 4.0).requires_grad_(True)
w_fn = (torch.rand(6, device="cuda", dtype=torch.float32) + 0.1).requires_grad_(True)
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
| def _fp8_cast_works() -> bool: | ||
| if not hasattr(torch, "float8_e4m3fn"): | ||
| return False | ||
| try: | ||
| torch.zeros(4).to(torch.float8_e4m3fn).to(torch.float32) | ||
| return True | ||
| except (RuntimeError, TypeError): | ||
| return False |
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.
|
TestIntel PR Steward: Follow-up on this poll (new Copilot review + persisting CI status). Copilot review ( Copilot review (same review, summary text) — backward CI ( Still watching for CI, further review comments, and merge-readiness. |
There was a problem hiding this comment.
🔵 Needs a closer look
The new tests don’t currently exercise fp8_input_store=True through backward and the impl-routing paths, leaving key behavior described in the PR summary uncovered.
Review details
Suppressed comments (4)
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:96
- The backward test only exercises fp8_input_store=False, so the fp8 saved-tensor/cast path in ClampedSwiGLUFunction.backward isn’t covered. Consider parametrizing this test over fp8_input_store and, for the fp8 case, compare against clamped_swiglu_back run on an fp8-cast copy of the input (to match what the autograd wrapper actually uses).
This issue also appears in the following locations of the same file:
- line 97
- line 132
- line 146
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:103
- The swiglu_impl routing test hard-codes fp8_input_store=False, so it doesn’t verify that the fp8_input_store plumbing is correctly forwarded through swiglu_impl → ClampedSwiGLUFunction.apply. Parameterizing over _fp8_input_store_cases would cover that without changing the reference math (forward computation is independent of the saved-tensor dtype).
def test_swiglu_impl_routes_to_clamped_function(self):
torch.manual_seed(9)
x = torch.randn(3, 5, 16, device="cuda", dtype=torch.float32)
out = self.swiglu_impl(x, None, fp8_input_store=False, clamp_value=7.0)
ref = _eager_clamped_swiglu(x.view(-1, 16), 7.0).view(3, 5, 8)
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:145
- The weighted backward test only covers fp8_input_store=False, so it doesn’t exercise the fp8 saved-input path (input stored as float8 then cast back) in ClampedWeightedSwiGLUFunction.backward. Consider parametrizing over fp8_input_store and, for fp8_input_store=True, compare grads against clamped_weighted_swiglu_back using an fp8-cast copy of the input tensor.
def test_backward_matches_eager_autograd(self):
torch.manual_seed(8)
alpha = 7.0
y_fn = (torch.randn(6, 16, device="cuda", dtype=torch.float32) * 4.0).requires_grad_(True)
w_fn = (torch.rand(6, device="cuda", dtype=torch.float32) + 0.1).requires_grad_(True)
y_e = y_fn.detach().clone().requires_grad_(True)
w_e = w_fn.detach().clone().requires_grad_(True)
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)
tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_clamped_swiglu.py:151
- The weighted_bias_swiglu_impl routing test hard-codes fp8_input_store=False, so it doesn’t cover the fp8_input_store forwarding behavior for the weighted wrapper path. Parameterizing the test over _fp8_input_store_cases would cover that at minimal cost.
def test_impl_routes_to_clamped_weighted_function(self):
torch.manual_seed(3)
x = torch.randn(4, 16, device="cuda", dtype=torch.float32)
weights = torch.rand(4, device="cuda", dtype=torch.float32) + 0.1
out = self.impl(x, None, weights, fp8_input_store=False, clamp_value=7.0)
torch.testing.assert_close(out, _eager_clamped_weighted_swiglu(x, weights, 7.0), atol=1e-5, rtol=1e-5)
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
|
Superseded by same-repo PR #1106 (head |
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