Update TE GroupedLinear to use single parameter mode

[hemildesai]

NVIDIA/TransformerEngine#2731

[SwekeR-463]

Is this up for grabs? I’d love to work on this.

From my read of TE PR #2731 (NVIDIA/TransformerEngine#2731), the main change in AutoModel is to make GroupedExpertsTE compatible with TE GroupedLinear single-parameter mode while preserving the existing stacked MoE checkpoint format.

Planned changes=>

• nemo_automodel/components/moe/experts.py to handle both TE layouts: legacy weight0 parameters and the new grouped weight parameter
• nemo_automodel/components/training/utils.py so EP grad scaling still matches expert parameters after the TE naming change
• Add unit tests for both layouts

[hemildesai]

Is this up for grabs? I’d love to work on this.

From my read of TE PR #2731 (NVIDIA/TransformerEngine#2731), the main change in AutoModel is to make GroupedExpertsTE compatible with TE GroupedLinear single-parameter mode while preserving the existing stacked MoE checkpoint format.

Planned changes=>

• nemo_automodel/components/moe/experts.py to handle both TE layouts: legacy weight0 parameters and the new grouped weight parameter
• nemo_automodel/components/training/utils.py so EP grad scaling still matches expert parameters after the TE naming change
• Add unit tests for both layouts

Hi, yes this is up for grabs, feel free to work on it. Also, maybe we can just drop the legacy format? That will keep the code simpler as we can remove all the hacks added to make that work.

[SwekeR-463]

Hi, yes this is up for grabs, feel free to work on it. Also, maybe we can just drop the legacy format? That will keep the code simpler as we can remove all the hacks added to make that work.

I have made the pr, waiting for review :).