Skip to content

[PyTorch] Fix NaN expert weight gradients at num_groups == 1 with SReLU #3485

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
GarlGuo wants to merge 1 commit into
base: main
Choose a base branch
from
+38 −14
Open

[PyTorch] Fix NaN expert weight gradients at num_groups == 1 with SReLU #3485

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
52 changes: 38 additions & 14 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 @@ -195,10 +195,15 @@ def _group_quantize_for_grouped_mlp(
split_sizes: Optional[torch.Tensor],
*,
tensor_offsets: Optional[torch.Tensor] = None,
use_dense_single_group: bool = False,
) -> GroupedTensor:
"""Quantize into grouped storage."""

if num_groups != 1 or not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)):
if (
num_groups != 1
or not use_dense_single_group
or not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer))
):
return tex.group_quantize(
tensor,
quantizer,
Expand Down Expand Up @@ -226,6 +231,7 @@ def _group_quantize_with_amax_for_grouped_mlp(
columnwise_amax: torch.Tensor,
*,
tensor_offsets: Optional[torch.Tensor] = None,
use_dense_single_group: bool = False,
) -> GroupedTensor:
"""Quantize with precomputed NVFP4 amaxes into grouped storage."""
if not isinstance(quantizer, NVFP4Quantizer):
Expand All @@ -235,9 +241,10 @@ def _group_quantize_with_amax_for_grouped_mlp(
num_groups,
split_sizes,
tensor_offsets=tensor_offsets,
use_dense_single_group=use_dense_single_group,
)

if num_groups != 1:
if num_groups != 1 or not use_dense_single_group:
return tex.nvfp4_group_quantize_with_amax(
tensor,
quantizer,
Expand Down Expand Up @@ -646,6 +653,7 @@ def _compute_grad_params(
scale_view_dtype,
sf_vec_size,
offsets,
use_dense_single_group,
):
"""Compute weight gradients and build grad_params for a GroupedLinear layer.
Returns the grad_params list in parameter registration order.
Expand Down Expand Up @@ -713,7 +721,7 @@ def _compute_grad_params(
"distributed-weight fused grouped-MLP requires delay_wgrad_compute=False."
)
if (
num_groups == 1
use_dense_single_group
and isinstance(grouped_x, (GroupedTensor, GroupedTensorStorage))
and isinstance(grouped_dy, (GroupedTensor, GroupedTensorStorage))
and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer))
Expand Down Expand Up @@ -1055,6 +1063,9 @@ def fuser_forward(
raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).")

num_groups = fc1_op.num_groups
use_dense_single_group = num_groups == 1 and (
_cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op))
)
Comment on lines +1066 to +1068

@greptile-apps greptile-apps Bot Sep 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P2 Regression coverage is missing

This capability gate fixes a specific numerical failure, but no automated test combines a padded single-group buffer, ScaledSReLU, MXFP8 or NVFP4, and backward gradient validation. The existing single-group MXFP8 test uses unpadded ScaledSwiGLU, while padded fused-MLP coverage uses GLU activations and four groups. Without a focused regression test, the dense-versus-packed layout mismatch and resulting NaN gradients could return without failing CI.

Knowledge Base Used:

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0
fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0

Expand Down Expand Up @@ -1125,7 +1136,7 @@ def fuser_forward(
# Older cuDNN frontends do not expose this specialization, so use the
# live generic offset calculation rather than caching CUDA metadata.
use_offsetless_metadata = (
num_groups == 1
use_dense_single_group
and unit_activation_scale
and isinstance(fc1_input_quantizer, MXFP8Quantizer)
and supports_single_group_runtime_offsets
Expand Down Expand Up @@ -1206,6 +1217,7 @@ def fuser_forward(
fc1_weight_quantizer,
num_groups,
None,
use_dense_single_group=use_dense_single_group,
)
else:
fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)]
Expand Down Expand Up @@ -1243,6 +1255,7 @@ def fuser_forward(
fc2_weight_quantizer,
num_groups,
None,
use_dense_single_group=use_dense_single_group,
)
else:
fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)]
Expand Down Expand Up @@ -1290,6 +1303,7 @@ def fuser_forward(
num_groups,
split_sizes,
tensor_offsets=fc1_x_tensor_offsets,
use_dense_single_group=use_dense_single_group,
)

use_nvfp4 = isinstance(fc1_input_quantizer, NVFP4Quantizer) or isinstance(
Expand Down Expand Up @@ -1421,7 +1435,7 @@ def fuser_forward(
fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor
fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4
if supports_single_group_runtime_offsets:
fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
fc1_activation_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
if self._pass_geglu_runtime_params:
fc1_activation_kwargs.update(
linear_offset=self._cudnn_linear_offset,
Expand All @@ -1438,7 +1452,7 @@ def fuser_forward(
if fc1_op.single_grouped_weight:
# Clone and swizzle scales for GEMM.
fc1_weight_for_gemm = grouped_fc1_weight.copy()
use_single_group_weight_swizzle = num_groups == 1
use_single_group_weight_swizzle = use_dense_single_group
if use_single_group_weight_swizzle:
fc1_weight_single = _single_quantized_tensor_from_grouped(fc1_weight_for_gemm)
fc1_weight_single._columnwise_data = None
Expand Down Expand Up @@ -1475,7 +1489,7 @@ def fuser_forward(
fc1_activation_kwargs["b_tensor"] = fc1_w_data
fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales
else:
use_single_discrete_weight = num_groups == 1
use_single_discrete_weight = use_dense_single_group
if use_single_discrete_weight:
fc1_weight_single = grouped_fc1_weight[0]
original_rowwise_scale = fc1_weight_single._rowwise_scale_inv
Expand Down Expand Up @@ -1577,6 +1591,7 @@ def fuser_forward(
fc1_kernel_out["amax_tensor"].view(-1),
fc1_kernel_out["post_rht_amax_tensor"].view(-1),
tensor_offsets=fc2_x_tensor_offsets,
use_dense_single_group=use_dense_single_group,
)
else:
grouped_fc2_x = _group_quantize_for_grouped_mlp(
Expand All @@ -1585,11 +1600,12 @@ def fuser_forward(
num_groups,
split_sizes,
tensor_offsets=fc2_x_tensor_offsets,
use_dense_single_group=use_dense_single_group,
)

fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device)
if (
num_groups == 1
use_dense_single_group
and grouped_fc2_x.columnwise_data is not None
and grouped_fc2_x.columnwise_scale_inv is not None
):
Expand Down Expand Up @@ -1647,7 +1663,7 @@ def fuser_forward(
with_gemm_swizzled_scales=True,
)

use_single_group_dense_fc2 = num_groups == 1
use_single_group_dense_fc2 = use_dense_single_group
fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device)
if use_single_group_dense_fc2:
fc2_out = _single_group_fc2_gemm(
Expand Down Expand Up @@ -1688,7 +1704,7 @@ def fuser_forward(
}
fc2_quant_kernel = self.grouped_gemm_quant_kernel()
if supports_single_group_runtime_offsets:
fc2_quant_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
fc2_quant_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group

if fc2_op.single_grouped_weight:
# Clone and swizzle scales for GEMM
Expand Down Expand Up @@ -1853,6 +1869,9 @@ def fuser_backward(
grad_output = grad_output.reshape(-1, fc2_weight_shape[0])
out_shape = list(grad_output.size())
num_groups = fc1_op.num_groups
use_dense_single_group = num_groups == 1 and (
_cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op))
)
fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0
fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0
device = fc1_weight_param.device
Expand Down Expand Up @@ -1959,6 +1978,7 @@ def fuser_backward(
num_groups,
split_sizes,
tensor_offsets=fc2_out_tensor_offsets,
use_dense_single_group=use_dense_single_group,
)

use_nvfp4 = (
Expand Down Expand Up @@ -2080,7 +2100,7 @@ def fuser_backward(
}
dactivation_kernel = self.grouped_gemm_dactivation_kernel()
if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
if self._cudnn_dact_func is not None:
fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor
fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func
Expand Down Expand Up @@ -2132,7 +2152,7 @@ def fuser_backward(
fc2_dactivation_kwargs["b_tensor"] = fc2_w_data
fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales
else:
use_single_discrete_weight = num_groups == 1
use_single_discrete_weight = use_dense_single_group
if use_single_discrete_weight:
fc2_weight_single = grouped_fc2_weight[0]
original_rowwise_data = fc2_weight_single._rowwise_data
Expand Down Expand Up @@ -2237,6 +2257,7 @@ def fuser_backward(
num_groups,
split_sizes,
tensor_offsets=fc2_x_tensor_offsets,
use_dense_single_group=use_dense_single_group,
)
else:
sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor")
Expand Down Expand Up @@ -2314,6 +2335,7 @@ def fuser_backward(
num_groups,
split_sizes,
tensor_offsets=fc1_dy_tensor_offsets,
use_dense_single_group=use_dense_single_group,
)
else:
grouped_fc1_dy = GroupedTensor(
Expand Down Expand Up @@ -2350,6 +2372,7 @@ def fuser_backward(
scale_view_dtype=scale_view_dtype,
sf_vec_size=sf_vec_size,
offsets=split_points,
use_dense_single_group=use_dense_single_group,
)

# Clear FC2 input tensor if possible
Expand All @@ -2375,7 +2398,7 @@ def fuser_backward(
if is_distributed_weight(fc1_leader):
grouped_fc1_weight = materialize_weight_for_backward(fc1_leader)

use_single_group_dense_dgrad = num_groups == 1
use_single_group_dense_dgrad = use_dense_single_group
if use_single_group_dense_dgrad:
grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device)
_single_group_dgrad_gemm(
Expand Down Expand Up @@ -2429,7 +2452,7 @@ def fuser_backward(
}
fc1_dgrad_kernel = self.grouped_gemm_quant_kernel()
if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group

if fc1_op.single_grouped_weight:
# Clone and swizzle scales for GEMM
Expand Down Expand Up @@ -2506,6 +2529,7 @@ def fuser_backward(
scale_view_dtype=scale_view_dtype,
sf_vec_size=sf_vec_size,
offsets=split_points,
use_dense_single_group=use_dense_single_group,
)

# Clear FC1 input tensor if possible
Expand Down