Skip to content

Move KD teacher loading to after Float16Module - #4394

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

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

Conversation

@AAnoosheh

@AAnoosheh AAnoosheh commented Apr 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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-pr-bot Bot commented Apr 20, 2026

Copy link
Copy Markdown

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

AAnoosheh commented Apr 30, 2026

Copy link
Copy Markdown
Contributor Author

/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

AAnoosheh commented Apr 30, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test a061477

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

svcnvidia-nemo-ci commented Apr 30, 2026

Copy link
Copy Markdown
Contributor

🔄 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

svcnvidia-nemo-ci commented May 1, 2026

Copy link
Copy Markdown
Contributor

🔄 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

AAnoosheh commented May 1, 2026

Copy link
Copy Markdown
Contributor Author

/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

svcnvidia-nemo-ci commented May 1, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented May 1, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented May 2, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented May 2, 2026

Copy link
Copy Markdown
Contributor

🔄 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

ko3n1g commented May 2, 2026

Copy link
Copy Markdown
Contributor

/ok to test cde70b0

@chtruong814

chtruong814 commented May 3, 2026

Copy link
Copy Markdown
Contributor

/ok to test

@copy-pr-bot

copy-pr-bot Bot commented May 3, 2026

Copy link
Copy Markdown

/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

chtruong814 commented May 3, 2026

Copy link
Copy Markdown
Contributor

/ok to test fd75d6f

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

svcnvidia-nemo-ci commented May 4, 2026

Copy link
Copy Markdown
Contributor

🔄 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
@AAnoosheh
AAnoosheh deleted the aanoosheh/kd-teacher-load-fix branch June 11, 2026 20:24
yhgalaxy pushed a commit to yhgalaxy/Megatron-LM that referenced this pull request Jun 17, 2026
@AAnoosheh @ko3n1g
@chtruong814 @yhgalaxy
Move KD teacher loading to after Float16Module (NVIDIA#4394)
1c15846 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
mathemakitten pushed a commit to mathemakitten/Megatron-LM that referenced this pull request Jul 1, 2026
@AAnoosheh @ko3n1g
@chtruong814 @mathemakitten
Move KD teacher loading to after Float16Module (NVIDIA#4394)
79d400e 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
jon-barker pushed a commit to jon-barker/Megatron-LM that referenced this pull request Jul 10, 2026
@AAnoosheh @ko3n1g
@chtruong814 Jon Barker
Move KD teacher loading to after Float16Module (NVIDIA#4394)
ebdf6e4 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026
@AAnoosheh @ko3n1g
@chtruong814 @terminator123
Move KD teacher loading to after Float16Module (NVIDIA#4394)
74994ad 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
svcnvidia-nemo-ci pushed a commit to dimapihtar/Megatron-LM that referenced this pull request Aug 4, 2026
@AAnoosheh @ko3n1g
@chtruong814 @dimapihtar
Move KD teacher loading to after Float16Module (NVIDIA#4394)
9047fed 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
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