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 fixes to the single grouped weight feature introduced in #3178. All three fixes address correctness gaps that surface when a grouped weight is used across training and eval steps, constructed on the meta device, or packed with MXFP8.

• High-precision init-val transfer (module/grouped_linear.py, module/base.py): when discrete per-GEMM weights are packed into a single GroupedTensor parameter, the BF16/FP16 master-weight seed values are now stacked and forwarded to the grouped parameter so distributed optimizers see identical initialization to the discrete layout.
• Late-attachment delayed-wgrad marking (ops/basic/grouped_linear.py): a register_parameter override ensures skip_backward_post_hook is set at the moment the grouped parent is attached to a meta-device shell, before DDP/FSDP installs backward hooks.
• Columnwise quantizer preservation (ops/fused/grouped_mlp.py): the MXFP8 path now updates the weight's own quantizer in-place and passes None (not False) for columnwise during eval, so a usage flag set during training is never silently dropped.

Confidence Score: 5/5

Safe to merge — all three fixes are narrowly scoped, well-tested, and address real correctness gaps without changing any shared code paths.

The changes are precise bug fixes with direct test coverage for each scenario. The set_usage(columnwise=None) change correctly leverages the existing None-means-no-op contract in Quantizer.set_usage. The register_parameter override is guarded by name == 'weight' and single_grouped_weight, so it cannot affect non-grouped paths. The high-precision init-val stacking is consistent with the shapes expected by the test assertions. No silent data-loss or backward-hook ordering issues were found.

Files Needing Attention: No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Refactors inline get/clear lambdas into module-level _get_high_precision_init_val, _clear_high_precision_init_val, and _attach_high_precision_init_val helpers; functionally equivalent to the old code.
transformer_engine/pytorch/module/grouped_linear.py	Adds high-precision init-val transfer when packing discrete weights into a GroupedTensor: collects, consistency-checks, stacks, attaches to the grouped parameter, and clears from originals.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Adds a register_parameter override to mark skip_backward_post_hook at grouped-weight attachment time (late-attach case); fixes _apply_delay_wgrad_param_hooks to tolerate a None weight on meta-device shells.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Preserves MXFP8 columnwise quantizer usage across training→eval transitions by updating the weight's own quantizer in-place and using None instead of False for columnwise when grad is not required.
tests/pytorch/test_grouped_mlp.py	Adds two new tests: meta-device delayed-wgrad hook attachment, and MXFP8 columnwise usage preservation across training→eval forward passes.
tests/pytorch/test_sanity.py	Adds two sanity tests: one verifying stacked init-val matches discrete-weight layout, one verifying partial-clearing raises RuntimeError.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant GroupedLinear_ops as GroupedLinear (ops)
    participant GroupedLinear_mod as GroupedLinear (module)
    participant GroupedTensor
    participant Quantizer

    Note over Caller,GroupedTensor: Fix 1 — high-precision init-val preservation
    Caller->>GroupedLinear_mod: make_grouped_weights()
    GroupedLinear_mod->>GroupedLinear_mod: collect high_precision_init_vals from weight0..N
    GroupedLinear_mod->>GroupedLinear_mod: consistency check (all or none must have init val)
    GroupedLinear_mod->>GroupedTensor: make_grouped_tensor_with_shapes()
    GroupedLinear_mod->>GroupedLinear_mod: torch.stack(init_vals) → attach to grouped_parameter
    GroupedLinear_mod->>GroupedLinear_mod: clear init_vals from individual weights

    Note over Caller,GroupedLinear_ops: Fix 2 — late-attach delayed-wgrad marking
    Caller->>GroupedLinear_ops: "GroupedLinear(device=meta, delay_wgrad_compute=True)"
    GroupedLinear_ops-->>Caller: "shell (weight=None)"
    Caller->>GroupedLinear_ops: register_parameter(weight, grouped_param)
    GroupedLinear_ops->>GroupedLinear_ops: "override: set skip_backward_post_hook=True"
    GroupedLinear_ops-->>Caller: parameter marked before DDP hooks install

    Note over Caller,Quantizer: Fix 3 — columnwise usage preservation (MXFP8)
    Caller->>GroupedLinear_ops: "forward() training, requires_grad=True"
    GroupedLinear_ops->>Quantizer: "set_usage(rowwise=True, columnwise=True)"
    Note right of Quantizer: columnwise_usage=True stored on weight's quantizer
    Caller->>GroupedLinear_ops: "forward() eval, requires_grad=False"
    GroupedLinear_ops->>Quantizer: "set_usage(rowwise=True, columnwise=None)"
    Note right of Quantizer: columnwise_usage preserved (None = no-op)

Loading

_{Reviews (5): Last reviewed commit: "Fix inconsistent line endings" | 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