Support auto_quantize for Megatron expert parallelism

[realAsma]

What does this PR do?

Type of change: Bug fix

This PR enables auto_quantize for Megatron expert parallel MoE flows by including the expert model parallel group when aggregating scores and costs and when synchronizing selected recipes. It also derives the search budget from the no-quant candidate costs in candidate_stats, so sharded expert layers use global candidate costs instead of local module weights.

Usage

model, search_state = mtq.auto_quantize(
    model,
    constraints={"effective_bits": 8.0},
    quantization_formats=[mtq.NVFP4_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG],
    data_loader=data_loader,
    forward_step=forward_step,
)

Testing

• Focused Megatron EP test from local log: python -m pytest tests/gpu_megatron/torch/quantization/plugins/test_megatron.py::test_auto_quantize_moe_ep -xvs in NGC PyTorch 26.01 (1 passed in 134.37s).
• Added unit coverage for deriving the auto_quantize budget from no-quant candidate costs.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: N/A

Additional Information

Base branch: jennifchen/super_nvfp4_recipe.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (3)

• main
• release/.*
• feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 41f783ac-34dc-4bde-9506-f0eb27da2467

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch asma/mcore-aq

---

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

[codecov]

Codecov Report

❌ Patch coverage is 45.45455% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.38%. Comparing base (cc4a570) to head (bea86d4).

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/plugins/megatron.py	8.33%	11 Missing ⚠️
modelopt/torch/quantization/algorithms.py	90.00%	1 Missing ⚠️

Additional details and impacted files

@@                        Coverage Diff                        @@
##           jennifchen/super_nvfp4_recipe    #1513      +/-   ##
=================================================================
- Coverage                          72.74%   69.38%   -3.37%     
=================================================================
  Files                                473      473              
  Lines                              51574    51590      +16     
=================================================================
- Hits                               37520    35794    -1726     
- Misses                             14054    15796    +1742     

Flag	Coverage Δ
examples	33.47% <45.45%> (-8.08%)	⬇️
gpu	51.01% <13.63%> (-0.02%)	⬇️
regression	14.97% <13.63%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.