fix(training): avoid double-counting MTP MoE metric layers

[odedovadia]

What does this PR do ?

Fix aggregate MoE metric normalization for unified hybrid patterns containing MTP layers.

The change separates the main decoder from the MTP suffix, counts their MoE layers independently, and passes those counts to MCore. For Nemotron 3.5 Lightning, the denominator changes from the incorrect 25 + 2 = 27 to 23 + 2 = 25.

Changelog

• Count main-decoder MoE layers before the MTP separator.
• Derive the MTP MoE count from the unified pattern when present.
• Add regression coverage for the Nemotron 3.5 Lightning pattern.

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

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?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

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

Additional Information

• Related to [bug] MoE metrics double-count MTP layers in unified hybrid patterns #5794

[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.

[Phlip79]

This logic has been changed on the Megatron-Core side in NVIDIA/Megatron-LM#6748.

[yaoyu-33]

/ok to test a3f1f25

[odedovadia]

This logic has been changed on the Megatron-Core side in NVIDIA/Megatron-LM#6748.

I see. I re-wrote this PR accordingly, as as far as I can tell a fix is still needed (jsut a slightly different one). I created a new function to properly count the real number of MoE/MTP layers according to the new MCore logic for both hybrid and non-hybrid models.

For example, Nemotron 3.5 Lightning counts num_moe_layers = num_layers + mtp_num_layers = 25 + 2 = 27 in main, which is wrong. After this new fix, we get num_moe_layers=24 (23 main MoE layers + 1 MTP, as MCore correctly scales the repeated MTPs).

I added some tests for different architectures too. While the issue was with hybrid models, I added a few tests for other architectures.

[Phlip79]

This logic has been changed on the Megatron-Core side in NVIDIA/Megatron-LM#6748.

I see. I re-wrote this PR accordingly, as as far as I can tell a fix is still needed (jsut a slightly different one). I created a new function to properly count the real number of MoE/MTP layers according to the new MCore logic for both hybrid and non-hybrid models.

For example, Nemotron 3.5 Lightning counts num_moe_layers = num_layers + mtp_num_layers = 25 + 2 = 27 in main, which is wrong. After this new fix, we get num_moe_layers=24 (23 main MoE layers + 1 MTP, as MCore correctly scales the repeated MTPs).

I added some tests for different architectures too. While the issue was with hybrid models, I added a few tests for other architectures.

Thanks for making the changes. LGTM.

[gautham-kollu]

/ok to test 7e1cd6e

[gautham-kollu]

Requesting review from @balasaajay/ @ko3n1g @malay-nagda