fix: update TE GroupedLinear integration for single-parameter mode

[SwekeR-463]

What does this PR do ?

Update GroupedExpertsTE to use TE GroupedLinear single-parameter mode and keep AutoModel’s MoE state dict format unchanged.

Changelog

• Switched TE MoE expert construction to single_grouped_parameter=True.
• Updated GroupedExpertsTE weight handling to read and write the grouped weight parameter directly.
• Kept AutoModel MoE checkpoint serialization in stacked tensor format.
• Updated EP grad-scaling name matching for the grouped TE parameter layout.
• Added and updated unit tests for the new TE path.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to Update TE GroupedLinear to use single parameter mode #1538

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test 5346557

[hemildesai]

Hi @SwekeR-463, thanks a lot for your contribution. Do you have an example wandb run verifying the convergence before and after this change?

[SwekeR-463]

Hi @SwekeR-463, thanks a lot for your contribution. Do you have an example wandb run verifying the convergence before and after this change?

I haven't done any runs, will do and update.

[akoumpa]

/ok to test bf85bc9

[SwekeR-463]

Hi @SwekeR-463, thanks a lot for your contribution. Do you have an example wandb run verifying the convergence before and after this change?

Hello @hemildesai, I attempted to run experiments to verify convergence before and after the change, but ran into repeated setup issues on my end. Rather than delaying further, I wanted to inform. 🙂

[akoumpa]

/ok to test f6f3fcd

[akoumpa]

Hi @SwekeR-463 I restarted CI, I apologize for the long delay.