[sync] DeepSeek V3 FSDP Fix for Precision-Aware Optimizer

[ananthsub]

What does this PR do ?

Sync with changes from NVIDIA/Megatron-LM#2466

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

Release Notes

• Changes

  • Modified precision-aware optimizer configuration handling to preserve default weight settings instead of automatically disabling them.

• Tests

  • Removed unit tests for precision-aware optimizer behavior with Megatron FSDP configuration.

_{✏️ 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 1147def

[ananthsub]

/ok to test c8bec24

[coderabbitai]

📝 Walkthrough

Walkthrough

Removed the code that sets preserve_fp32_weights to False when use_precision_aware_optimizer is enabled in the config validation logic, and removed the corresponding test methods that validated this behavior.

Changes

Cohort / File(s)	Summary
Config Simplification src/megatron/bridge/training/config.py	Removed the assignment that disables preserve_fp32_weights when use_precision_aware_optimizer is enabled, leaving the setting at its prior value.
Test Cleanup tests/unit_tests/training/test_config.py	Removed test methods test_megatron_fsdp_with_precision_aware_optimizer and test_megatron_fsdp_without_precision_aware_optimizer that validated the removed config behavior.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~5 minutes

Suggested reviewers

• maanug-nv
• yaoyu-33

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR description lacks test results or testing information despite changes affecting numerical precision and convergence behavior of preserve_fp32_weights with precision-aware optimizer.	Add explicit test results, convergence validation, regression testing data, or reference upstream testing to demonstrate no numerical regression before merge.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main change: removing the assignment that disables preserve_fp32_weights for precision-aware optimizer in Megatron FSDP, which is a specific fix for DeepSeek V3.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

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

✨ Finishing touches

• 📝 Generate docstrings

---

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