Skip to content

[PyTorch] Grouped MLP: emit FP32 scale gradient with no cast to BF16 #3870

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sraman-rgb wants to merge 11 commits into
base: main
Choose a base branch
from
+6 −9
Open

[PyTorch] Grouped MLP: emit FP32 scale gradient with no cast to BF16 #3870

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
9550305
[PyTorch] Grouped MLP: emit FP32 scale gradient in one kernel for FP3…
sraman-rgb Nov 9, 2026
423987a
[PyTorch] Grouped MLP: keep the FP32 scale gradient, without rounding it
sraman-rgb Nov 13, 2026
3d72adc
Merge branch 'main' into sraman/grouped-mlp-fp32-scales-grad
sraman-rgb Nov 13, 2026
6aa6f81
[PyTorch] Grouped MLP: FP32 router-probability scales end to end, no …
sraman-rgb Nov 14, 2026
0ee0cfa
Merge branch 'main' into sraman/grouped-mlp-fp32-scales-grad
vthumbe1503 Nov 14, 2026
7dc8e0c
[PyTorch] Test: the MXFP8 GLU forward now gets FP32 probabilities
sraman-rgb Nov 14, 2026
f871cab
Cast only for NVFP4
vthumbe1503 Nov 14, 2026
55152d3
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 14, 2026
8f24294
Move fp32 scale conversion for NVFP4 condition
vthumbe1503 Nov 14, 2026
5c3553d
revert
vthumbe1503 Nov 14, 2026
4e11b5e
Revert properly
vthumbe1503 Nov 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion tests/pytorch/test_grouped_mlp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1804,7 +1804,7 @@ def test_grouped_mlp_glu_mxfp8_real_cudnn_fusion(
"a_dtype": torch.float8_e4m3fn,
"b_dtype": torch.float8_e4m3fn,
"alpha_dtype": torch.bfloat16,
"prob_dtype": torch.bfloat16,
"prob_dtype": torch.float32,
"c_dtype": torch.bfloat16,
"d_dtype": torch.float8_e4m3fn,
"act_func": act_func,
Expand Down
13 changes: 5 additions & 8 deletions transformer_engine/pytorch/ops/fused/grouped_mlp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1589,11 +1589,11 @@ def fuser_forward(

fc1_prob_tensor = None
if not unit_activation_scale:
fc1_prob_tensor = (
scales.detach()
.to(dtype=torch.float32 if use_nvfp4 else dtype)
.reshape((-1,) if canonical else (-1, 1, 1))
)
fc1_prob_tensor = scales.detach().reshape((-1,) if canonical else (-1, 1, 1))
if use_nvfp4:
# cuDNN NVFP4 kernel only accepts fp32 scales.
fc1_prob_tensor = fc1_prob_tensor.to(torch.float32)

if use_nvfp4:
# cuDNN receives NVFP4 block-scaled inputs without TE's per-group
# global scale factors, so alpha supplies the product of the two
Expand Down Expand Up @@ -2701,9 +2701,6 @@ def fuser_backward(
else:
fc2_bias_grads = [fc2_dbias_packed[idx] for idx in range(num_groups)]

if grad_scales is not None:
grad_scales = grad_scales.to(dtype=dtype)

fc1_bias_grads: Optional[list[Optional[torch.Tensor]]] = None
fc1_bias_grad_packed: Optional[torch.Tensor] = None
if fc1_op.has_bias:
Expand Down
Loading