grouped linear single param fixes

[CarlosGomes98]

Description

These changes fix convergence for grouped linear single param.
Issues encountered:

Without the TE fused ops path: causes the experts to get requires_grad=0 and the parent to not receive updates from the individual experts.

With TE fused ops: discards the high precision metadata, so the master weights were starting from dequantized mxfp8 weights, rather than the original high precision weights

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

These changes resolve the immediate issue, but cleaner medium- to long-term solutions would require broader changes.
In particular, the bug could likely have been avoided if Transformer Engine supported packing expert weights before quantization. The current flow quantizes each expert separately and then packs the resulting parameters, requiring the packing code to manually preserve and transfer FP8 initialization metadata. This stretches the metadata lifecycle and makes it easy to lose information needed to construct the optimizer’s high-precision master weights. Packing the high-precision weights first and quantizing the grouped parameter once would provide a cleaner and more robust ownership model.

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

[timmoon10]

If I understand, this PR fixes 3 bugs:

• The grouped linear module with a single grouped weight does not actually populate the param grad. It computes discrete weight grads, but it never ended up in the single weight param. The fix is to wrap the conversion from single grouped weights to discrete weights within an autograd function.
• The grouped linear module with a single grouped weight was not handling high-precision weight caching (used in Mcore DDP/FSDP). We were caching high-precision weights for the discrete weights, but they were not registered in the single weight param. The fix is to concat the discrete high-precision values and register it in the single weight param.
• When constructing a grouped linear op with device="meta", it might incorrectly set some Mcore-specific param attrs before the params are initialized. Mcore does this when it constructs an internal fused to replace the grouped linear module implementation (see here).

This is all hairy and messy, but the fixes seem reasonable. It seems that if you enable single grouped weights or any Mcore-specific optimizations, then grouped linear becomes extremely delicate. Users should avoid unless they know what they are doing.

[greptile-apps]

Greptile Summary

This PR fixes grouped linear behavior for the single grouped-parameter path. The main changes are:

• Preserves high-precision initialization metadata when packing expert weights.
• Connects per-expert weight views back to the grouped parent for autograd.
• Marks late-attached grouped weights for delayed wgrad handling.
• Keeps fused grouped MLP columnwise quantizer usage after training forwards.
• Adds tests for gradients, metadata preservation, delayed wgrad attachment, and fused MXFP8 usage.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/basic/grouped_linear.py	Applies delayed-wgrad metadata when a single grouped weight is registered after construction.
transformer_engine/pytorch/module/grouped_linear.py	Adds grouped-weight autograd bridging and transfers high-precision init metadata during grouped weight packing.
transformer_engine/pytorch/module/base.py	Moves high-precision init metadata handling into shared helper functions.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Keeps columnwise quantizer usage enabled once a grad-enabled forward requests it.
tests/pytorch/test_grouped_mlp.py	Adds tests for late grouped-weight attachment and fused grouped MLP columnwise usage.
tests/pytorch/test_sanity.py	Adds tests for single grouped-parameter gradients and high-precision initialization preservation.

_{Reviews (4): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[zhongbozhu]

without this this, will there be runtime errors in actual training? how do we trigger it?

[CarlosGomes98]

Yes. The real underlying issue is lack of explicit API support on the TE side for the MCore pattern of initializing ops with device=meta and then injecting the module weights later. Replicating the issue requires

USE_TE_OPS=True
MOE_SINGLE_GROUPED_WEIGHT=True
NVTE_GROUPED_LINEAR_SINGLE_PARAM=1
DELAY_WGRAD_COMPUTE=True
NVTE_CUTEDSL_FUSED_GROUPED_MLP=1

The order is:

1. TE GroupedLinear OP created by MCore with device meta
2. It created placeholder parameters, but weight is not created yet
3. This part of the code before the changes calls self.weight.skip_backward_post_hook = True, but weight does not exist yet.

In the current code, self._parameters.get("weight") will return None in this case, self.weight wont be accessed. In the "regular" case where device=meta is not used, it will work as normal.

[zhongbozhu]

why?

[CarlosGomes98]

Same as below:

This non-fused path will operate on the individual expert weights, whereas MCore is attaching distributed training flags to the grouped parent. So we need to pass these flags down to the expert weights.

[zhongbozhu]

why do we need this for loop to setattr?

[CarlosGomes98]

This non-fused path will operate on the individual expert weights, whereas MCore is attaching distributed training flags to the grouped parent. So we need to pass these flags down to the expert weights.

[zhongbozhu]

how does it avoid the dequantize exactly?

[CarlosGomes98]

This is how TE's initialization contract for high precision works. If the parameter has a method get_high_precision_init_val, then MCore uses that to initialize the optimizer master weights. If its not present, then it dequantizes the model weights.

https://github.com/NVIDIA/Megatron-LM/blob/eadbaa6189fbd272843b4e2bd75ae20984f891d4/megatron/core/optimizer/distrib_optimizer.py#L412

[zhongbozhu]

There needs to be have a numerical test for both fuse_wgrad_accumulation and not fuse_wgrad_accumulation. Compared with discrete weights, with or without the fusion, the wgrad output should be extremely close to discrete.

[CarlosGomes98]

Agreed, added a numerical test

[zhongbozhu]

We need a less hacky implementation.

[CarlosGomes98]

We can hide this away by directly using the _ methods I introduced in base.py, but its effectively doing the same thing.

I dont disagree with your general idea that a lot of this is hacky, but as far as I'm aware this is all we can do given the current TE API. Making it less hacky would require a non-trivial redesign and implementation work which I wouldnt want to add to this PR