Various ModelOpt fixes: QAD test for CICD, use model builder config instead of model provider, allow loading teacher ckpt independently of student - #4520

Merged

AAnoosheh merged 26 commits into

NVIDIA:mainfrom

AAnoosheh:aanoosheh/qad-cicd-test

Jul 14, 2026

AAnoosheh commented Apr 29, 2026 •

edited

Loading

Contributor

What does this PR do ?

Adds a new functional test for testing PTQ + KD in succession using Model Optimizer.

⚠️ 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.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

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.

AAnoosheh self-assigned this

copy-pr-bot Bot commented Apr 29, 2026

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.

AAnoosheh force-pushed the aanoosheh/qad-cicd-test branch from cb52062 to d06c767 Compare

May 11, 2026 13:04

copy-pr-bot Bot commented May 11, 2026

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.

AAnoosheh marked this pull request as ready for review

June 18, 2026 15:20

AAnoosheh requested review from a team as code owners

June 18, 2026 15:20

copy-pr-bot Bot temporarily deployed to test

June 18, 2026 15:21

Inactive

svcnvidia-nemo-ci added the complexity: medium label

copy-pr-bot Bot temporarily deployed to test

June 19, 2026 12:05

Inactive

kevalmorabia97 approved these changes

View reviewed changes

AAnoosheh force-pushed the aanoosheh/qad-cicd-test branch from b856371 to 9c75554 Compare

June 23, 2026 18:57

copy-pr-bot Bot temporarily deployed to public

June 23, 2026 18:57

Inactive

copy-pr-bot Bot temporarily deployed to test

June 23, 2026 18:58

Inactive

copy-pr-bot Bot temporarily deployed to public

June 23, 2026 19:01

Inactive

copy-pr-bot Bot temporarily deployed to public

June 23, 2026 19:01

Inactive

copy-pr-bot Bot temporarily deployed to public

June 23, 2026 19:01

Inactive

copy-pr-bot Bot temporarily deployed to public

June 23, 2026 19:01

Inactive

copy-pr-bot Bot temporarily deployed to public

June 23, 2026 19:02

Inactive

copy-pr-bot Bot temporarily deployed to public

June 23, 2026 19:11

Inactive

jenchen13 approved these changes

View reviewed changes

megatron/post_training/model_builder.py

copy-pr-bot Bot temporarily deployed to public

July 1, 2026 15:58

Inactive

copy-pr-bot Bot temporarily deployed to test

July 1, 2026 15:59

Inactive

copy-pr-bot Bot temporarily deployed to public

July 1, 2026 16:02

Inactive

copy-pr-bot Bot temporarily deployed to public

July 1, 2026 16:02

Inactive

copy-pr-bot Bot temporarily deployed to public

July 1, 2026 16:02

Inactive

copy-pr-bot Bot temporarily deployed to public

July 1, 2026 16:02

Inactive

copy-pr-bot Bot temporarily deployed to public

July 1, 2026 16:02

Inactive

copy-pr-bot Bot temporarily deployed to public

July 1, 2026 16:11

Inactive

chtruong814 approved these changes

View reviewed changes


          Re-enable total loss golden vals

99f386d

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

svcnemo-autobot mentioned this pull request

AUT-681 Fix QAD functional test: use local jsonl for PTQ calibration to avoid offline HF Hub fetch #5651

Closed

AAnoosheh commented Jul 6, 2026

Contributor Author

/ok to test 99f386d

copy-pr-bot Bot temporarily deployed to public

July 6, 2026 13:12

Inactive

copy-pr-bot Bot temporarily deployed to test

July 6, 2026 13:12

Inactive

copy-pr-bot Bot temporarily deployed to public

July 6, 2026 13:15

Inactive

copy-pr-bot Bot temporarily deployed to public

July 6, 2026 13:15

Inactive

copy-pr-bot Bot temporarily deployed to public

July 6, 2026 13:15

Inactive

copy-pr-bot Bot temporarily deployed to public

July 6, 2026 13:16

Inactive

yaoyu-33 approved these changes

View reviewed changes

jaredcasper reviewed

View reviewed changes

megatron/training/training.py Outdated

AAnoosheh added 2 commits

July 7, 2026 18:49


          Fix more issues

a137f30

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>


          Refactor Modelopt builder

5d57b6c

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

AAnoosheh commented Jul 7, 2026

Contributor Author

/ok to test 5d57b6c

AAnoosheh added 5 commits

July 8, 2026 14:17


          Fix more goldens

6dac7b2

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>


          Move ModelOpt builder selection out of ModelConfig, into pretrain_gpt…

4d3af48

….py/pretrain_hybrid.py

Per Jared/Maanu feedback in review: ModelOpt-specific branching doesn't
belong in the shared ModelConfig base class (base.py's get_builder_cls()
was checking args.modelopt_enabled and swapping in ModelOptModelBuilder
for every config type). Revert base.py entirely - it is now byte-for-byte
identical to NVIDIA/Megatron-LM main again.

Instead:
- Add ModelOptModelConfig(GPTModelConfig) and ModelOptHybridModelConfig
  (HybridModelConfig) in post_training/model_builder.py, each overriding
  only the `builder` ClassVar - construction still goes through
  gpt_config_from_args/hybrid_config_from_args, which now take an
  optional model_config_cls to build a compatible subclass without
  duplicating arg-derivation logic.
- Their builders (ModelOptModelBuilder/ModelOptHybridModelBuilder) share
  a `build_model()` override via `_ModelOptBuildModelMixin`, since
  modelopt_gpt_hybrid_builder already dispatches on args.export_model_type
  internally - but each is tied to its own proper parent
  (GPTModelBuilder/HybridModelBuilder respectively) for
  build_distributed_models(), rather than both piggybacking on
  GPTModelBuilder's.
- pretrain_gpt.py's and pretrain_hybrid.py's __main__ blocks now select
  the ModelOpt config vs the plain one based on args.modelopt_enabled,
  right after parse_and_validate_args() - this works now that args is
  available in __main__ before pretrain() is called.
- maybe_enable_modelopt (the args.modelopt_enabled detection, previously
  private to training.py) now lives in megatron.post_training.utils,
  alongside print_distributed_quant_summary - both need modelopt
  installed at import time, so it's a natural fit. training.py imports
  it once at module level (in the same try/except as
  get_tensor_shapes_adjust_fn_for_distillation) and both call sites just
  guard the call with `if has_nvidia_modelopt:`. pretrain_gpt.py and
  pretrain_hybrid.py pick it up via the same try/except block they
  already use for add_modelopt_args.

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>


          Move load_kd_teacher_checkpoint to post_training module

bda26ce

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>


          Merge remote-tracking branch 'github/main' into aanoosheh/qad-cicd-test

728be54

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

# Conflicts:
#	megatron/post_training/checkpointing.py


          Fix note

ad86d92

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

AAnoosheh commented Jul 8, 2026

Contributor Author

/ok to test ad86d92

maanug-nv reviewed

View reviewed changes

megatron/training/argument_utils.py

maanug-nv approved these changes

View reviewed changes

maanug-nv left a comment

Contributor

lgtm overall now. left one minor suggestion

AAnoosheh added 3 commits

July 8, 2026 22:25


          Address nit comment

860606c

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>


          Add some things that help but don't fix determinism yet

61ee9dc

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>


          Remove moe-permute-fusion

0f680d9

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>

AAnoosheh commented Jul 13, 2026

Contributor Author

/ok to test 0f680d9

deepakn94 approved these changes

View reviewed changes

svcnvidia-nemo-ci commented Jul 14, 2026

Contributor

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci mentioned this pull request

chore: nightly sync main into dev (30_07_2026) #6159

Closed

lhb8125 mentioned this pull request

2026年7月 Megatron-LM 工作汇总 lhb8125/Megatron-LM#56

Open

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Labels

Approved complexity: medium Run functional tests