[sync] clear aux losses tracker after eval, add cuda graph validation

[ananthsub]

What does this PR do ?

Sync training changes from NVIDIA/Megatron-LM#3056

Changelog

• Add specific line by line info of high level changes in this PR.

GitHub Actions CI

See the CI sectionin 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 # (issue)

Summary by CodeRabbit

• Bug Fixes

  • Fixed MoE auxiliary losses incorrectly accumulating in training metrics across evaluation boundaries, ensuring more accurate performance monitoring for mixture-of-experts models.
  • Added configuration validation for CUDA graphs with Transformer Engine to prevent incompatible settings and guide proper setup.

• Tests

  • Expanded test coverage for CUDA graph and Transformer Engine configuration validation scenarios.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

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

[ananthsub]

/ok to test e3c39a8

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds runtime validation to enforce NCCL_GRAPH_REGISTER=0 when using CUDA graphs with Transformer Engine and expandable segments configured, and introduces MoE auxiliary loss tracker cleanup at evaluation boundaries to prevent metric leakage.

Changes

Cohort / File(s)	Summary
CUDA Graph Validation src/megatron/bridge/training/config.py	Added validation in TrainingConfig.validate() to assert that when expandable_segments is enabled, NCCL_GRAPH_REGISTER must be set to "0"; raises error with guidance if constraint is violated.
MoE Auxiliary Loss Tracking src/megatron/bridge/training/train.py	Imported clear_aux_losses_tracker and added conditional calls to reset MoE auxiliary losses at evaluation boundaries (after energy monitor resume and after evaluation completion) when num_moe_experts is configured.
Validation Test Coverage tests/unit_tests/training/test_config.py	Added unit test test_cuda_graph_transformer_engine_expandable_segments_validation to verify NCCL_GRAPH_REGISTER validation behavior with and without expandable_segments, including environment variable patching.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• [sync] DeepSeek V3 FSDP Fix for Precision-Aware Optimizer #2093: Modifies the same TrainingConfig.validate function with different validation logic for CUDA graph configurations.

Suggested reviewers

• gautham-kollu
• yaoyu-33

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly and clearly summarizes the two main changes: clearing the auxiliary losses tracker after evaluation and adding CUDA graph validation.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	PR contains minor bug fixes and validation improvements (~47 lines) with comprehensive unit test validating CUDA graph validation logic.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@tests/unit_tests/training/test_config.py`:
- Around line 1143-1174: The new test function
test_cuda_graph_transformer_engine_expandable_segments_validation is missing a
pytest mark; add a class/function decorator such as `@pytest.mark.unit` (or the
appropriate category per guidelines) immediately above the def line and ensure
pytest is imported at the top of the file (add "import pytest" if absent); keep
the decorator placement consistent with other tests in the file so the test is
properly categorized by the test runner.

[gautham-kollu]

Can we do this in the post_init in TransformerConfig