Skip to content

[PyTorch] Fix Test's Reference Implementation of Clamped Swiglu - #3424

Merged
vthumbe1503 merged 1 commit into
NVIDIA:mainfrom
vthumbe1503:clamped_swiglu_fix_v2
Aug 25, 2026
Merged

[PyTorch] Fix Test's Reference Implementation of Clamped Swiglu#3424
vthumbe1503 merged 1 commit into
NVIDIA:mainfrom
vthumbe1503:clamped_swiglu_fix_v2

Conversation

@vthumbe1503

@vthumbe1503vthumbe1503 commented Aug 25, 2026•
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Pytorch recently changed their clamp function's autograd behavior in this PR: https://github.com/pytorch/pytorch/pull/191142/changes#diff-e1a04c95c2809af34915952d3c7cc6210fadd2a277af1cfdb4924c74a5795295R121

torch.clamp method's original behavior was derivative = 1 at the clamp boundary. Now they changed it to be 0 instead.(I think one of the reasons is consistency of relu operation which has also produces zero derivative at boundary=0)

This PR changes the reference implementation of clamped swiglu in pytorch to use torch.where() to explicitly define the boundary gradient

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@vthumbe1503
fix clamped swiglu test
06847e8
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503

vthumbe1503 commented Aug 25, 2026

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci pytorch

@greptile-apps

greptile-appsBot commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Updates two PyTorch test reference implementations for clamped SwiGLU to express upper and lower clamping with explicit torch.where operations.

  • Replaces the GLU branch’s upper clamp in both references.
  • Replaces the linear branch’s two-sided clamp in both references.

Confidence Score: 5/5

The PR appears safe to merge because the updated test references preserve the production clamped SwiGLU semantics.

The explicit conditional clamps preserve forward values and pass gradients through at the exact limits, matching the production kernel behavior in both changed test paths.

Important Files Changed

FilenameOverview
tests/pytorch/test_fusible_ops.pyThe two reference implementations now match the production kernel’s inclusive clamp-boundary gradient semantics without introducing an actionable issue.

Reviews (1): Last reviewed commit: "fix clamped swiglu test" | Re-trigger Greptile

timmoon10
timmoon10 reviewed Aug 25, 2026
View reviewed changes

@timmoon10timmoon10 left a comment •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we know why this error started appearing?

I'm suspicious of #3402, which touched the logic for GLU backwards. However, the existing logic also preserved gradients at boundaries. Also, the CI passed for that PR.

#3359 landed exactly when the CI pipelines started failing, but it doesn't seem relevant and it passed CI.

I see that the errors are caused by a recent change in PyTorch. This is a reasonable fix.

Comment threadtests/pytorch/test_fusible_ops.py
@vthumbe1503

vthumbe1503 commented Aug 25, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Errors are pre-existing in main

@vthumbe1503
vthumbe1503 merged commit 32adc64 into NVIDIA:mainAug 25, 2026
18 of 28 checks passed
cael-ling added a commit to cael-ling/TransformerEngine that referenced this pull request Aug 26, 2026
@cael-ling
[PyTorch] Check the grouped scaled SwiGLU against TE's SwiGLU reference
66aa446
The grouped kernel is the fused recompute counterpart of ScaledSwiGLU and
ScaledClampedQGeGLU, but nothing checked that the two compute the same
function. The Python tests covered only shapes and rejections, and the C++
suite compares against its own CPU reference, so both sides could be
internally correct while disagreeing with each other.
Reuses the reference from test_fusible_ops.py verbatim, including the
torch.where formulation from NVIDIA#3424, so this tracks that definition rather than
restating it. Per expert the reference comes from that expert's own rows and is
quantized with the quantizer the kernel used, so a wrong per-group offset
cannot average away. Being at the Python level it also covers the bindings,
quantizer flags, grouped split and scale layout that the C++ test skips.
The comparison is resolution-level rather than bit-exact, mainly because SiLU
goes through tanh.approx.f32, so the two sides already differ before either is
quantized. Asserts that e8m0 scales agree to within one step and that
dequantized values sit within two MXFP8 code steps of their own 32-row block.
Two steps rather than one covers a block whose e8m0 exponents straddle a
boundary; measured on GB200 the worst deviation is 0.889 steps, so the kernel
stays inside a single step everywhere.
Also switches the tests added here to the public te.DType.
Signed-off-by: Cael Ling <caell@nvidia.com>
@cael-lingcael-ling mentioned this pull request Aug 26, 2026
Merged
13 tasks
Oleg-Goncharov added a commit that referenced this pull request Aug 28, 2026
@cael-ling@pre-commit-ci
@vthumbe1503@greptile-apps@Oleg-Goncharov
[Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel (#3315)
174caef
* [Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel
Add nvte_group_swiglu_quantize, which fuses the weighted SwiGLU
activation with columnwise MXFP8 quantization over grouped (MoE)
tensors.
Signed-off-by: Cael Ling <caell@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [PyTorch] Require contiguous, colocated operands in group_swiglu_quantize
Signed-off-by: Cael Ling <caell@nvidia.com>
* [PyTorch] Require grouped metadata tensors on the current CUDA device
Signed-off-by: Cael Ling <caell@nvidia.com>
* [PyTorch] Scope the grouped metadata device check to group_swiglu_quantize
Signed-off-by: Cael Ling <caell@nvidia.com>
* [Common/PyTorch] Rename to group_scaled_swiglu and speed up the kernel
lines up with TE's existing ScaledSwiGLU op,
so nvte_group_swiglu_quantize becomes nvte_group_scaled_swiglu, and the
kernel header, dispatch, bindings and tests follow.
Also makes the kernel actually faster than the unfused path it replaces.
Per-token scales are staged through shared memory once per chunk, the
output buffer is single-buffered to raise occupancy, and SiLU uses an
approximate exp and divide. The activation is quantized to MXFP8
immediately, so the approximation stays far below one FP8 ULP.
Signed-off-by: Cael Ling <caell@nvidia.com>
* [PyTorch] Add a benchmark for group_scaled_swiglu
Times the fused kernel against the unfused path it replaces, with the
activation half computed three ways: eager PyTorch, torch.compile, and
TE's existing ScaledSwiGLU op.
Signed-off-by: Cael Ling <caell@nvidia.com>
* [Common/PyTorch] Halve the silu MUFU cost and add clamped scaled SwiGLU
Signed-off-by: Cael Ling <caell@nvidia.com>
* [Common] Cut the clamped scaled SwiGLU ALU cost
Signed-off-by: Cael Ling <caell@nvidia.com>
* [PyTorch] Benchmark the clamped kernel against a clamped baseline
Signed-off-by: Cael Ling <caell@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update transformer_engine/pytorch/csrc/extensions/cast.cpp
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
* [Common][PyTorch] Promote the approximate SwiGLU operators to math.h
Signed-off-by: Cael Ling <caell@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [PyTorch] Check the grouped scaled SwiGLU against TE's SwiGLU reference
The grouped kernel is the fused recompute counterpart of ScaledSwiGLU and
ScaledClampedQGeGLU, but nothing checked that the two compute the same
function. The Python tests covered only shapes and rejections, and the C++
suite compares against its own CPU reference, so both sides could be
internally correct while disagreeing with each other.
Reuses the reference from test_fusible_ops.py verbatim, including the
torch.where formulation from #3424, so this tracks that definition rather than
restating it. Per expert the reference comes from that expert's own rows and is
quantized with the quantizer the kernel used, so a wrong per-group offset
cannot average away. Being at the Python level it also covers the bindings,
quantizer flags, grouped split and scale layout that the C++ test skips.
The comparison is resolution-level rather than bit-exact, mainly because SiLU
goes through tanh.approx.f32, so the two sides already differ before either is
quantized. Asserts that e8m0 scales agree to within one step and that
dequantized values sit within two MXFP8 code steps of their own 32-row block.
Two steps rather than one covers a block whose e8m0 exponents straddle a
boundary; measured on GB200 the worst deviation is 0.889 steps, so the kernel
stays inside a single step everywhere.
Also switches the tests added here to the public te.DType.
Signed-off-by: Cael Ling <caell@nvidia.com>
---------
Signed-off-by: Cael Ling <caell@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10timmoon10timmoon10 approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@vthumbe1503@timmoon10