Add MTP acceptance rate metrics by yfw · Pull Request #3458 · NVIDIA/Megatron-LM

yfw · 2026-02-17T19:50:11Z

What does this PR do ?

Refactor MTPLossLoggingHelper to track acceptance rate (correct/total predictions) alongside loss metrics. Computes distributed argmax across tensor parallel ranks to determine predictions without gathering full vocab logits (conditionally). Logs per-step and cumulative acceptance rates.

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

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
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

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

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!

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

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.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

copy-pr-bot · 2026-02-17T19:50:15Z

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.

macandro96 · 2026-06-02T20:17:38Z

/ok to test 2393096

macandro96 · 2026-06-02T21:41:51Z

/ok to test d6d1eea

svcnvidia-nemo-ci · 2026-06-03T21:51:34Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-06-03T23:35:41Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-06-04T00:44:46Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-06-04T04:26:39Z

🔄 Merge queue validation started!

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

Refactor MTPLossLoggingHelper to track acceptance rate (correct/total predictions) alongside loss metrics. Computes distributed argmax across tensor parallel ranks to determine predictions without gathering full vocab logits. Logs per-step and cumulative acceptance rates. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

…ate tests Update existing tests to use the new save_metrics_to_tracker API and clean_metrics_in_tracker, replacing the old save_loss_to_tracker and clean_loss_in_tracker calls. Add test coverage for acceptance rate metrics tracking. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

macandro96 · 2026-06-04T15:58:12Z

/ok to test 0402398

svcnvidia-nemo-ci · 2026-06-04T17:28:05Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-06-04T18:21:11Z

🔄 Merge queue validation started!

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

yfw force-pushed the yifu/mtp_acceptance_rate branch 2 times, most recently from 398f25a to 159431c Compare April 2, 2026 21:35

yfw mentioned this pull request Apr 2, 2026

Clip mtp grads separately when mtp_detach_heads=True #4116

Merged

5 tasks

yfw force-pushed the yifu/mtp_acceptance_rate branch 4 times, most recently from 7463948 to a13d976 Compare April 10, 2026 17:09

macandro96 force-pushed the yifu/mtp_acceptance_rate branch from a13d976 to 2393096 Compare June 2, 2026 20:17

copy-pr-bot Bot temporarily deployed to public June 2, 2026 20:18 Inactive

copy-pr-bot Bot temporarily deployed to test June 2, 2026 20:18 Inactive

copy-pr-bot Bot temporarily deployed to public June 2, 2026 20:21 Inactive

copy-pr-bot Bot temporarily deployed to public June 2, 2026 20:22 Inactive

macandro96 force-pushed the yifu/mtp_acceptance_rate branch from 2393096 to d6d1eea Compare June 2, 2026 21:40

macandro96 marked this pull request as ready for review June 2, 2026 21:41

macandro96 requested review from a team as code owners June 2, 2026 21:41

copy-pr-bot Bot temporarily deployed to public June 2, 2026 21:42 Inactive

svcnvidia-nemo-ci added the complexity: medium label Jun 2, 2026

copy-pr-bot Bot temporarily deployed to test June 2, 2026 21:43 Inactive

copy-pr-bot Bot temporarily deployed to public June 2, 2026 21:45 Inactive

copy-pr-bot Bot temporarily deployed to public June 2, 2026 21:46 Inactive

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Jun 3, 2026

macandro96 enabled auto-merge June 3, 2026 20:51

macandro96 added this pull request to the merge queue Jun 3, 2026

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

macandro96 added this pull request to the merge queue Jun 3, 2026

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

macandro96 added this pull request to the merge queue Jun 4, 2026

macandro96 removed this pull request from the merge queue due to a manual request Jun 4, 2026

yfw and others added 2 commits June 4, 2026 08:55

macandro96 force-pushed the yifu/mtp_acceptance_rate branch from 603e175 to 0402398 Compare June 4, 2026 15:56

copy-pr-bot Bot temporarily deployed to public June 4, 2026 15:58 Inactive

copy-pr-bot Bot temporarily deployed to test June 4, 2026 15:59 Inactive

copy-pr-bot Bot temporarily deployed to public June 4, 2026 16:02 Inactive

copy-pr-bot Bot temporarily deployed to public June 4, 2026 16:11 Inactive

macandro96 enabled auto-merge June 4, 2026 16:20

macandro96 added this pull request to the merge queue Jun 4, 2026

Jintao-Huang mentioned this pull request Jun 17, 2026

[compat] MTP compat megatron-core main branch / dev branch modelscope/mcore-bridge#127

Merged

Navigation Menu

Uh oh!

Add MTP acceptance rate metrics - #3458

Add MTP acceptance rate metrics#3458
macandro96 merged 2 commits into
NVIDIA:mainfrom
yfw:yifu/mtp_acceptance_rate

yfw commented Feb 17, 2026 •

edited by macandro96

Loading

Uh oh!

copy-pr-bot Bot commented Feb 17, 2026

Uh oh!

macandro96 commented Jun 2, 2026

Uh oh!

macandro96 commented Jun 2, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 3, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Jun 3, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 4, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Jun 4, 2026

Uh oh!

Uh oh!

macandro96 commented Jun 4, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 4, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Uh oh!

Conversation

yfw commented Feb 17, 2026 • edited by macandro96 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Feb 17, 2026

Uh oh!

macandro96 commented Jun 2, 2026

Uh oh!

macandro96 commented Jun 2, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 3, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Jun 3, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 4, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Jun 4, 2026

Uh oh!

Uh oh!

macandro96 commented Jun 4, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 4, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

yfw commented Feb 17, 2026 •

edited by macandro96

Loading

(Step 1): Add PR label `Expert Review`