Rename MLM teacher arg

[AAnoosheh]

What does this PR do?

Type of change: Refactor

Overview: MLM arg changed from --teacher-model-config to --export-kd-teacher-model-config for consistency

Usage

# Add a code snippet demonstrating how to use this

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

Summary by CodeRabbit

• Chores
  • Updated export flag naming for knowledge distillation teacher model configuration.
  • Adjusted default top-k parameter for logits selection to 1024.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Two minor parameter updates: one replaces a command-line flag name for teacher model configuration export in a shell script, and another updates the default top-k value for logits selection in a distillation plugin class.

Changes

Cohort / File(s)	Summary
Shell Script Configuration examples/llm_qad/qad.sh	Renamed export flag from --teacher-model-config to --export-kd-teacher-model-config in CHECKPOINT_ARGS for knowledge distillation teacher model configuration export.
Distillation Plugin modelopt/torch/distill/plugins/megatron.py	Updated default top_k parameter value in TopKLogitsKLLoss.__init__ from 1000 to 1024 for top-k logits selection in teacher distribution.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~5 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The title 'Rename MLM teacher arg' accurately describes the main change: renaming --teacher-model-config to --export-kd-teacher-model-config. However, the other file change (updating top_k from 1000 to 1024) is unrelated to argument renaming and not reflected in the title.	The PR mixes two unrelated changes: renaming an argument and updating a parameter default value. Either split into two PRs or update the title to cover both changes (e.g., 'Update MLM teacher arg naming and top_k default').

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

modelopt/torch/distill/plugins/megatron.py (1)

380-386: Confirm the new default top_k=1024 is safe for all supported vocabs.

This default change can alter loss behavior or assert for small vocab sizes if any code relies on the default. Please verify usages and update any configs/tests/docs that assumed 1000.

#!/bin/bash
# Find instantiations relying on the default and config-driven top_k settings.
rg -n -C2 -- 'TopKLogitsKLLoss\(' -g'*.py'
rg -n -C2 -- 'logit_kl_topk' -g'*.py' -g'*.yaml' -g'*.yml' -g'*.conf'

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.82%. Comparing base (9857e0a) to head (43a5d0e).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #829   +/-   ##
=======================================
  Coverage   73.82%   73.82%           
=======================================
  Files         193      193           
  Lines       19745    19745           
=======================================
  Hits        14577    14577           
  Misses       5168     5168           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.