From 43a5d0e35de719613ca15d94438cc77cca4612b6 Mon Sep 17 00:00:00 2001
From: Asha Anoosheh <aanoosheh@nvidia.com>
Date: Thu, 29 Jan 2026 05:43:31 -0800
Subject: [PATCH] Rename MLM arg

Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
---
 examples/llm_qad/qad.sh                    | 2 +-
 modelopt/torch/distill/plugins/megatron.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/llm_qad/qad.sh b/examples/llm_qad/qad.sh
index 52ec2bd6a..ac416ad35 100644
--- a/examples/llm_qad/qad.sh
+++ b/examples/llm_qad/qad.sh
@@ -181,7 +181,7 @@ CHECKPOINT_ARGS=" \
     ${LOAD_OPTIM_ARGS} \
     --load ${LOAD_CHECKPOINT_DIR} \
     --export-kd-teacher-load ${TEACHER_CKPT} \
-    --teacher-model-config ${TEACHER_MODEL_CONFIG}"
+    --export-kd-teacher-model-config ${TEACHER_MODEL_CONFIG}"
 
 # KD config (optional)
 if [[ -n "$KD_CFG_PATH" && -f "$KD_CFG_PATH" ]]; then
diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py
index b0eeeab74..dbfad6fb6 100644
--- a/modelopt/torch/distill/plugins/megatron.py
+++ b/modelopt/torch/distill/plugins/megatron.py
@@ -382,7 +382,7 @@ def __init__(
         model_config: "TransformerConfig",
         temperature: float = 1.0,
         reverse: bool = False,
-        top_k: int = 1000,
+        top_k: int = 1024,
     ):
         """Constructor.