fix: handle list-format quant_cfg from ModelOpt PR #1094 by ChenhanYu · Pull Request #4187 · NVIDIA/Megatron-LM

ChenhanYu · 2026-04-07T21:43:41Z

ModelOpt PR #1094 changed quant_cfg from a flat dict keyed by quantizer name to an ordered list of QuantizerCfgEntry dicts with an explicit 'quantizer_name' field. quantize.py was directly indexing the dict (e.g. quant_cfg['medusa_heads*'] = ...) which crashes with a TypeError when quant_cfg is a list.

Add _set_quant_cfg_entry() and _get_quant_cfg_entry() helpers that write to and read from either format, and update all call sites in get_modelopt_torch_quantization_config(), get_first_layers_disabled_config(), and get_last_layers_disabled_config() to use them.

What does this PR do ?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

copy-pr-bot · 2026-04-07T21:43:45Z

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.

github-actions · 2026-04-07T21:43:52Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

copy-pr-bot · 2026-04-15T03:36:53Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

ModelOpt PR NVIDIA#1094 changed quant_cfg from a flat dict keyed by quantizer name to an ordered list of QuantizerCfgEntry dicts with an explicit 'quantizer_name' field. quantize.py was directly indexing the dict (e.g. quant_cfg['*medusa_heads**'] = ...) which crashes with a TypeError when quant_cfg is a list. Add _set_quant_cfg_entry() and _get_quant_cfg_entry() helpers that write to and read from either format, and update all call sites in get_modelopt_torch_quantization_config(), get_first_layers_disabled_config(), and get_last_layers_disabled_config() to use them. Signed-off-by: chenhany <chenhany@nvidia.com>

ModelOpt PR NVIDIA#1094 changed QuantizerCfgEntry to require quantizer attributes under a 'cfg' sub-key rather than at the top level. Update fp8_config and fp4_config in quantize.py accordingly. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

initialize_megatron() and pretrain() no longer accept extra_args_provider and args_defaults after commit 97aca2f. Use parse_and_validate_args() from megatron.training.arguments to parse args before calling initialize_megatron()/pretrain(). Affected files: - quantize.py, validate.py, prune.py, mmlu.py, export.py - convert_model.py, offline_feature_extract.py, generate.py - finetune.py (uses pretrain() instead of initialize_megatron()) Signed-off-by: Chenhan Yu <chenhany@nvidia.com>

…ency-ext Some container versions renamed get_write_results_queue to _get_write_results_queue. Add fallback import to handle both.

Set DISABLE_NVRX=1 to force HAVE_NVRX=False and use mcore's own checkpoint implementation. Workaround for containers where nvrx is incompatible with the installed PyTorch version.

nvidia-resiliency-ext's FileSystemWriterAsync is incompatible with some container PyTorch versions. DISABLE_NVRX forces mcore fallback.

jenchen13 · 2026-04-22T18:51:54Z

 #!/bin/bash

+# Disable nvidia-resiliency-ext (incompatible with some container PyTorch versions)
+export DISABLE_NVRX=1


can you add this in arguments.sh instead of in multiple places?

jenchen13 · 2026-04-22T18:54:48Z

+
+    ModelOpt PR #1094 changed quant_cfg from a flat dict keyed by quantizer name
+    to an ordered list of QuantizerCfgEntry dicts with an explicit "quantizer_name"
+    field.  This helper writes to either format so callers don't need to branch.


could this be moved to ModelOpt quant cfg class to add get/set functions so that it is backward compatible?

The old format is problematic and deprecated, but backward compatibility is kept:

The ModelOpt code automatically convert the old formatted configs into new format configs at the time parsing the config

The issues we were hitting are from code manipulating the old formatted configs, such as updating quantization format of one quantizer. These code is not in consideration for backward compatibility.

We do not want to bring those manipulation code back into modelopt library.

I've updated the code to convert from old dict format to list format if it detects an old modelopt config, and to only manipulate list configs

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>

jenchen13 · 2026-04-27T21:01:17Z

/ok to test a62125b

svcnvidia-nemo-ci · 2026-04-28T13:40:17Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25056285782

svcnvidia-nemo-ci · 2026-04-28T13:51:32Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25056874744

svcnvidia-nemo-ci · 2026-04-28T17:51:32Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25068932293

…A#4187) Signed-off-by: chenhany <chenhany@nvidia.com> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Jennifer Chen <jennifchen@nvidia.com>

…A#4187) Signed-off-by: chenhany <chenhany@nvidia.com> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Jennifer Chen <jennifchen@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

…A#4187) Signed-off-by: chenhany <chenhany@nvidia.com> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Jennifer Chen <jennifchen@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

…A#4187) Signed-off-by: chenhany <chenhany@nvidia.com> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Jennifer Chen <jennifchen@nvidia.com>

svcnvidia-nemo-ci marked this pull request as draft April 7, 2026 21:43

ChenhanYu force-pushed the chenhany/fix-quant-cfg-list-format branch 2 times, most recently from cf20883 to 7d8d36f Compare April 8, 2026 04:42

ChenhanYu force-pushed the chenhany/fix-quant-cfg-list-format branch from 1550efc to 8bc5455 Compare April 15, 2026 03:36

ChenhanYu and others added 3 commits April 20, 2026 14:38

ChenhanYu force-pushed the chenhany/fix-quant-cfg-list-format branch from 8bc5455 to c11adb0 Compare April 20, 2026 21:51

ChenhanYu added 3 commits April 20, 2026 16:04

fix: fallback for privatized get_write_results_queue in nvidia-resili…

bda6752

…ency-ext Some container versions renamed get_write_results_queue to _get_write_results_queue. Add fallback import to handle both.

fix: add DISABLE_NVRX env var to bypass nvidia-resiliency-ext

084cd5e

Set DISABLE_NVRX=1 to force HAVE_NVRX=False and use mcore's own checkpoint implementation. Workaround for containers where nvrx is incompatible with the installed PyTorch version.

fix: set DISABLE_NVRX=1 in all modelopt example scripts

fc6a931

nvidia-resiliency-ext's FileSystemWriterAsync is incompatible with some container PyTorch versions. DISABLE_NVRX forces mcore fallback.

jenchen13 self-requested a review April 21, 2026 16:08

jenchen13 reviewed Apr 22, 2026

View reviewed changes

jenchen13 and others added 4 commits April 27, 2026 11:43

revert: remove DISABLE_NVRX=1 from modelopt example scripts

7e08b1b

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

revert: restore torch.py get_write_results_queue fallback to main

9056598

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

fix quantize.py quant config modification

0348040

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>

only support list-type modelopt config in quantize.py

4560fed

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>

jenchen13 requested a review from shengliangxu April 27, 2026 19:38

jenchen13 marked this pull request as ready for review April 27, 2026 19:38

svcnvidia-nemo-ci requested a review from a team April 27, 2026 19:38

svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 27, 2026

copy-pr-bot Bot had a problem deploying to test April 27, 2026 19:39 Error

svcnvidia-nemo-ci added the complexity: medium label Apr 27, 2026

jenchen13 requested review from AAnoosheh, kevalmorabia97 and shengliangxu and removed request for shengliangxu April 27, 2026 19:41

use existing find_quant_cfg_entry_by_path from mtq

0fc2c08

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>

copy-pr-bot Bot temporarily deployed to test April 27, 2026 19:48 Inactive

kevalmorabia97 approved these changes Apr 27, 2026

View reviewed changes

Merge branch 'main' into chenhany/fix-quant-cfg-list-format

a62125b

copy-pr-bot Bot temporarily deployed to test April 27, 2026 20:20 Inactive

shengliangxu approved these changes Apr 27, 2026

View reviewed changes

jenchen13 added this pull request to the merge queue Apr 28, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Apr 28, 2026

jenchen13 added this pull request to the merge queue Apr 28, 2026

Merged via the queue into NVIDIA:main with commit 9816140 Apr 28, 2026
69 checks passed

Navigation Menu

Uh oh!

fix: handle list-format quant_cfg from ModelOpt PR #1094 - #4187

fix: handle list-format quant_cfg from ModelOpt PR #1094#4187
jenchen13 merged 12 commits into
NVIDIA:mainfrom
ChenhanYu:chenhany/fix-quant-cfg-list-format

ChenhanYu commented Apr 7, 2026

Uh oh!

copy-pr-bot Bot commented Apr 7, 2026

Uh oh!

github-actions Bot commented Apr 7, 2026

Uh oh!

copy-pr-bot Bot commented Apr 15, 2026

Uh oh!

jenchen13 Apr 22, 2026

Uh oh!

jenchen13 Apr 22, 2026

Uh oh!

shengliangxu Apr 27, 2026

Uh oh!

jenchen13 Apr 27, 2026

Uh oh!

jenchen13 commented Apr 27, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 28, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 28, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Apr 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

ChenhanYu commented Apr 7, 2026

What does this PR do ?

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Apr 7, 2026

Uh oh!

github-actions Bot commented Apr 7, 2026

Uh oh!

copy-pr-bot Bot commented Apr 15, 2026

Uh oh!

jenchen13 Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

jenchen13 Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

shengliangxu Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

jenchen13 Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

jenchen13 commented Apr 27, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 28, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 28, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Apr 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants