Single grouped weight fixes

[CarlosGomes98]

Description

A few fixes to the single grouped weight implementation

12406c9 — Preserve grouped weight initialization metadata
99b3f37 — Mark late grouped weights for delayed wgrad
aa4e233 — Preserve grouped MXFP8 columnwise usage

See #3178 for the original PR.
PR #3224 is very relevant to this one, both are required fixes.

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

[greptile-apps]

Greptile Summary

This PR delivers three targeted correctness fixes to the single_grouped_weight feature introduced in #3178, each paired with a dedicated test.

• module/grouped_linear.py: make_grouped_weights now stacks and transfers per-GEMM _high_precision_init_val CPU tensors to the new grouped parameter (so distributed optimizers build lossless FP32 masters from the original BF16/FP16 init rather than dequantized MXFP8); partially-preserved state raises RuntimeError instead of silently mixing sources.
• ops/basic/grouped_linear.py: A new register_parameter override marks the grouped weight with skip_backward_post_hook at attachment time, fixing the meta-device shell scenario where _apply_delay_wgrad_param_hooks ran before the weight existed; the same method now guards against a None weight for that case.
• ops/fused/grouped_mlp.py: columnwise_usage is made sticky (or fc1_weight_quantizer.columnwise_usage) for pre-quantized GroupedTensor weights so that an eager eval no-grad pass cannot drop columnwise buffers that a captured CUDA-graph training forward still requires.

Confidence Score: 5/5

Safe to merge — three targeted bug fixes each with an accompanying test; no existing API surfaces are changed.

All three fixes are narrow and well-scoped. The high-precision init-val transfer is gated behind an all(value is not None) guard and raises an explicit error on inconsistency. The register_parameter override only fires for name == weight with both single_grouped_weight and delay_wgrad_compute set. The columnwise_usage stickiness is a monotonic flag that corrects a latent crash under CUDA graph replay. Tests cover all three scenarios including edge cases like meta-device shells and partial init-val state.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Refactors inline closures for high-precision init-val get/clear/attach into module-level helpers so they can be shared. Semantics are identical to the original closures.
transformer_engine/pytorch/module/grouped_linear.py	Adds high-precision init-val transfer inside make_grouped_weights: per-GEMM CPU tensors are stacked (dim=0) and attached to the new grouped parameter before individual params are cleared.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Overrides register_parameter to mark a late-attached grouped weight with skip_backward_post_hook when delay_wgrad_compute=True; guards _apply_delay_wgrad_param_hooks against a missing weight in the meta-shell scenario.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Makes columnwise_usage sticky for pre-quantized GroupedTensor weights so eval no-grad passes cannot drop columnwise buffers still needed by captured CUDA-graph training forwards.
tests/pytorch/test_grouped_mlp.py	Adds meta-device delayed-wgrad and columnwise-usage persistence tests.
tests/pytorch/test_sanity.py	Adds high-precision init-val packing correctness and partial-state rejection tests.

_{Reviews (1): Last reviewed commit: "[PyTorch] Preserve grouped MXFP8 columnw..." | Re-trigger Greptile}

[zhongbozhu]

/te-ci pytorch L1

[zhongbozhu]

LGTM

[phu0ngng]

/te-ci pytorch L1