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 bug fixes to the single grouped weight path introduced in #3178: preserving high-precision initialization metadata when packing discrete weights into a GroupedTensor, correctly deferring skip_backward_post_hook assignment when a meta-device op receives its grouped weight late, and preventing MXFP8 columnwise buffer deallocation during eval when a CUDA-graph-captured training forward still holds references.

• High-precision init preservation (module/grouped_linear.py, module/base.py): Three module-level helpers (_get/clear/attach_high_precision_init_val) replace inline closures. make_grouped_weights stacks and transfers the per-weight init tensors to the packed parameter, and raises RuntimeError when the init state is inconsistent across weights.
• Delayed wgrad on meta-device ops (ops/basic/grouped_linear.py): A register_parameter override marks skip_backward_post_hook at attachment time when the grouped parent arrives after construction; _apply_delay_wgrad_param_hooks now safely skips a missing weight rather than AttributeError-ing.
• MXFP8 columnwise usage stickiness (ops/fused/grouped_mlp.py): set_usage now ORs input_requires_grad with fc1/fc2_weight_quantizer.columnwise_usage so an eval forward cannot reset to columnwise=False when a captured training graph still needs the columnwise buffers.

Confidence Score: 5/5

All three fixes are narrow and well-targeted, each addressing a distinct silent failure mode without touching shared infrastructure beyond the refactored helpers.

Each bug is exercised by a dedicated new test that directly asserts the previously-broken invariant. The refactored _attach_high_precision_init_val helper is a pure extraction of existing inline closure logic with identical semantics, and the defensive getattr guards in the register_parameter override prevent regressions during super().init() ordering.

No files require special attention. The MXFP8 columnwise stickiness fix in ops/fused/grouped_mlp.py is the most subtle change, but the accompanying test validates the end-to-end behavior.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Extracts three helper functions for high-precision init val lifecycle; original call-site updated to use them. Logic is equivalent to the inline closures they replace.
transformer_engine/pytorch/module/grouped_linear.py	Transfers high-precision init vals to the packed grouped parameter in make_grouped_weights, with consistency guard. Clean and correct.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Adds register_parameter override to mark skip_backward_post_hook on late-attached grouped weights; _apply_delay_wgrad_param_hooks now null-guards for meta-device shells.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	ORs columnwise_usage with input_requires_grad when calling set_usage for FC1 and FC2 single-grouped-weight paths; prevents eval from releasing MXFP8 columnwise buffers still referenced by a captured training graph.
tests/pytorch/test_grouped_mlp.py	Adds tests for delayed-wgrad meta-device late attachment and MXFP8 columnwise-usage preservation across train/eval transitions. Coverage is accurate and targeted.
tests/pytorch/test_sanity.py	Adds tests verifying high-precision init val packing (round-trip equality) and partial-init rejection. Well-scoped and matches the production invariants.

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

[zhongbozhu]

/te-ci pytorch L1

[zhongbozhu]

LGTM

[phu0ngng]

/te-ci pytorch L1

[timmoon10]

LGTM

[timmoon10]

The op and the param purposely have different quantizers so that you can change them independently. However, we are incorrectly reusing the same quantizer in two places, and getting unexpected interactions. The op's quantizer is set in each forward pass:

TransformerEngine/transformer_engine/pytorch/ops/basic/grouped_linear.py

Line 697 in aa4e233

weight_quantizer.set_usage(rowwise=True, columnwise=requires_grad)

The proper fix is keep the op and param quantizers distinct.

[timmoon10]

/te-ci pytorch