Improve device-init grouped linear module with single grouped weight support

[zhongbozhu]

Description

Fixes numerical issues when using single weight for TE module grouped linear. Limit the single weight feature to the grouped tensor API instead of the legacy path.

TODO: test E2E convergence, unit test directly from Mcore.

Note: needs to pay extra attention to whether bias grad and weight grad are generated properly.

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

[zhongbozhu]

Note: this code block was deleted because it was duplicated

[zhongbozhu]

/te-ci pytorch

[greptile-apps]

Greptile Summary

This PR adds a single_grouped_weight / single_grouped_bias experimental feature to both the module-level and op-fuser GroupedLinear paths. Instead of N separate per-GEMM weight{i} parameters, a single GroupedTensor parameter is registered, reducing optimizer state and checkpoint size. The PR also refactors the grouped-tensor path predicate into exported is_op_fuser_grouped_tensor_path_supported / is_module_grouped_tensor_path_supported helpers and extends group_quantize with output-reuse and noop_flag support.

• Single-param layout: New make_grouped_weights() / _make_grouped_biases() methods build a GroupedTensor from a packed storage tensor; checkpoint compatibility is maintained via _remap_grouped_weight_state_dict_keys.
• C++ extension: group_quantize gains an output parameter for in-place workspace reuse across iterations, plus noop_flag threading into the FP8 block-scaling kernel.
• CUDA kernel fix: group_quantize_mxfp8.cuh corrects per-member scale-base computation for SAME_BOTH_DIMS shape representation.

Confidence Score: 3/5

Do not merge until the two P1 issues are resolved: meta-device AttributeError on construction and the missing skip guard in test_single_grouped_weight_rejects_host_m_splits.

Two confirmed P1 defects — one causes an AttributeError at module construction time on meta device with specific flags, and one causes a test to fail with the wrong exception type on pre-Hopper hardware. Both are straightforward one-line fixes but block correctness on non-Hopper CI runners.

Files Needing Attention: transformer_engine/pytorch/ops/basic/grouped_linear.py (lines 289-297) and tests/pytorch/test_grouped_linear.py (lines 48-63)

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds single_grouped_weight/bias support, make_grouped_weights(), state-dict remapping, and grouped-tensor forward/backward paths; logic is correct but early return in make_grouped_weights() for float8_current_scaling silently skips weight registration
transformer_engine/pytorch/ops/basic/grouped_linear.py	Refactors grouped-tensor predicate and adds single-param support; P1 bug: _apply_delay_wgrad_param_hooks accesses self.weight before registration on meta device
transformer_engine/pytorch/csrc/extensions/cast.cpp	Extends group_quantize with output reuse and noop_flag; identity/shape pre-checks are correct, noop_flag threading into FP8 block-scaling looks sound
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Fixes per-member scale-base for SAME_BOTH_DIMS; integer division is used without remainder assertion — benign if callers enforce divisibility
tests/pytorch/test_grouped_linear.py	538-line test suite for new single-param feature; test_single_grouped_weight_rejects_host_m_splits is missing _require_native_grouped_tensor_gemm() skip guard (P1)
tests/pytorch/test_grouped_mlp.py	Replaces manual capability checks with is_op_fuser_grouped_tensor_path_supported; adds single_grouped_bias storage test; looks correct
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Minor fix: replaces _get_bias_tensors+stack with _get_packed_bias_tensor to support single_grouped_bias; correct

_{Reviews (13): Last reviewed commit: "relax atol for bf16 sum errors" | Re-trigger Greptile}

[timmoon10]

The biggest change in this PR is that TE is abandoning any attempt to make single_grouped_weight=True a general feature. Things must be exactly right, or we crash. Given how delicate and experimental this feature has been, I'm not opposed.

The second change is that users must opt-in to access the grouped GEMM kernel. This is also reasonable, since it has alignment requirements for m_splits and it's helpful having a way for users to accept that stricter contract.

We are experiencing many test failures. Given that single_grouped_weight is no longer a general feature, I think it's reasonable we move the corresponding tests to test_grouped_linear.py and test_grouped_mlp.py.

[timmoon10]

I get that the h2d memcpy is suboptimal, but it's trivially easy to handle. Erroring out seems excessively rigid.

Suggested change

	raise ValueError(
	"The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 "
	"tensor, or select grouped_gemm_backend='legacy'."
	)
	m_splits = m_splits.to(device=device)

We need to handle the d2h case anyways when the user has specified grouped_gemm_backend="grouped_tensor", but it's not supported and we fallback to split-quantize.

[zhongbozhu]

But if the alignment is not provided in the first place, converting it to a device tensor also wouldn't work right, I am okay with another alignment check before adding this H2D copy.

[timmoon10]

These are misleading names.

• Legacy implies we are deprecating it, but in fact it is the default backend and it is the fallback when the grouped tensor impl is not supported.
• As mentioned in https://github.com/NVIDIA/TransformerEngine/pull/3224/changes#r3618050729, it's confusing that we might not actually perform the backend specified by the user. The user might put grouped_gemm_backend="grouped_tensor" and we'll actually do split-quantize because the cuBLAS version is old or something.

Really, this is not a backend config. It's a hint where the user promises to provide m_splits with a certain alignment. I'd suggest changing this to a bool like enable_grouped_tensor_backend. Alternatively if we want to generalize in case we add more backends in the future, we could have a string like enabled_backends (enabled_backends="grouped_tensor,future_backend").

[zhongbozhu]

and we'll actually do split-quantize because the cuBLAS version is old or something

Should we just raise error instead of silently doing this fallback? As a user, it's actually annoying to specify some optimization but it's not working because a fallback get triggered.

I am okay with changing the TE interface into a bool flag like enable_grouped_tensor_backend instead of a str.

[zhongbozhu]

renamed

[zhongbozhu]

TODO: I don't quite understand why this is the case, but without it, cublas will complain cannot find supported gemm. Unclear if it's related with TE integration or a cublas issue or if upgrading to the next cublas version can fix it.
cc @denera @timmoon10

[zhongbozhu]

/te-ci pytorch