Skip to content

Move KD teacher loading to after Float16Module#4394

Merged
ko3n1g merged 7 commits intoNVIDIA:mainfrom
AAnoosheh:aanoosheh/kd-teacher-load-fix
May 4, 2026
Merged

Move KD teacher loading to after Float16Module#4394
ko3n1g merged 7 commits intoNVIDIA:mainfrom
AAnoosheh:aanoosheh/kd-teacher-load-fix

Conversation

@AAnoosheh
Copy link
Copy Markdown
Contributor

@AAnoosheh AAnoosheh commented Apr 20, 2026 •
edited
Loading

What does this PR do ?

Fix bug where KD teacher loading before Float16Module would undesirably convert router bias to 16-bit instead of being always held in 32-bit.

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"

  1. When your PR is ready, click Ready for Review.
  2. 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
Move KD teacher loading to after Float16Module
487cec0 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented Apr 20, 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
Prevent recursion
8f1fa32 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
@AAnoosheh AAnoosheh marked this pull request as ready for review April 28, 2026 17:43
@AAnoosheh AAnoosheh requested review from a team as code owners April 28, 2026 17:43
@AAnoosheh AAnoosheh self-assigned this Apr 28, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci requested a review from a team April 28, 2026 17:43
jenchen13
jenchen13 reviewed Apr 28, 2026
View reviewed changes
Comment thread megatron/post_training/model_builder.py
@AAnoosheh
Add comment
38a5330 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Apr 30, 2026
@AAnoosheh
Copy link
Copy Markdown
Contributor Author

AAnoosheh commented Apr 30, 2026

/ok to test 38a5330

@AAnoosheh AAnoosheh enabled auto-merge April 30, 2026 14:08
@AAnoosheh
Safeguard args
a061477 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
@AAnoosheh
Copy link
Copy Markdown
Contributor Author

AAnoosheh commented Apr 30, 2026

/ok to test a061477

@AAnoosheh AAnoosheh added this pull request to the merge queue Apr 30, 2026
@svcnvidia-nemo-ci
Copy link
Copy Markdown

svcnvidia-nemo-ci commented Apr 30, 2026

🔄 Merge queue validation started!

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

@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Apr 30, 2026
@AAnoosheh AAnoosheh added this pull request to the merge queue May 1, 2026
@svcnvidia-nemo-ci
Copy link
Copy Markdown

svcnvidia-nemo-ci commented May 1, 2026

🔄 Merge queue validation started!

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

@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks May 1, 2026
@AAnoosheh
Copy link
Copy Markdown
Contributor Author

AAnoosheh commented May 1, 2026

/ok to test d9bd2eb

@AAnoosheh AAnoosheh enabled auto-merge May 1, 2026 20:17
@AAnoosheh AAnoosheh added this pull request to the merge queue May 1, 2026
@svcnvidia-nemo-ci
Copy link
Copy Markdown

svcnvidia-nemo-ci commented May 1, 2026

🔄 Merge queue validation started!

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

@svcnvidia-nemo-ci
Copy link
Copy Markdown

svcnvidia-nemo-ci commented May 1, 2026

🔄 Merge queue validation started!

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

@svcnvidia-nemo-ci
Copy link
Copy Markdown

svcnvidia-nemo-ci commented May 2, 2026

🔄 Merge queue validation started!

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

@svcnvidia-nemo-ci
Copy link
Copy Markdown

svcnvidia-nemo-ci commented May 2, 2026

🔄 Merge queue validation started!

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

@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks May 2, 2026
@ko3n1g ko3n1g enabled auto-merge May 2, 2026 21:42
@ko3n1g
Copy link
Copy Markdown
Contributor

ko3n1g commented May 2, 2026

/ok to test cde70b0

@chtruong814
Copy link
Copy Markdown
Contributor

chtruong814 commented May 3, 2026

/ok to test

@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented May 3, 2026

/ok to test

@chtruong814, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

@chtruong814
Copy link
Copy Markdown
Contributor

chtruong814 commented May 3, 2026

/ok to test fd75d6f

@ko3n1g ko3n1g added this pull request to the merge queue May 4, 2026
@svcnvidia-nemo-ci
Copy link
Copy Markdown

svcnvidia-nemo-ci commented May 4, 2026

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit 99abdc8 May 4, 2026
545 of 556 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@deepakn94 deepakn94 deepakn94 approved these changes

@jenchen13 jenchen13 jenchen13 approved these changes

@kevalmorabia97 kevalmorabia97 kevalmorabia97 approved these changes

@yaoyu-33 yaoyu-33 yaoyu-33 approved these changes

Assignees

@AAnoosheh AAnoosheh

Labels

Approved All necessary approvals have been made complexity: low Run functional tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

8 participants

@AAnoosheh @svcnvidia-nemo-ci @ko3n1g @chtruong814 @deepakn94 @jenchen13 @kevalmorabia97 @yaoyu-33