Replace ModelOpt generation server

examples/post_training/modelopt/generation_server.sh

@@ -14,11 +14,9 @@ if [ -z ${MLM_MODEL_CKPT} ]; then
     exit 1
 fi
 
-if [ -z ${DRAFT_LEN} ]; then
-    DRAFT_LEN=0
-fi
+TOOLS_DIR="$(realpath ${SCRIPT_DIR}/../../../tools)"
 
-${LAUNCH_SCRIPT} ${SCRIPT_DIR}/generation_server.py \
+${LAUNCH_SCRIPT} ${TOOLS_DIR}/run_text_generation_server.py \
     ${MODEL_ARGS} \
     --tensor-model-parallel-size ${TP} \
     --expert-tensor-parallel-size ${ETP} \

megatron/training/training.py

@@ -855,6 +855,18 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap
     args = get_args()
     args.model_type = model_type
 
+    if has_nvidia_modelopt:
+        from megatron.post_training.checkpointing import has_modelopt_state
+        # [ModelOpt]: Check if the checkpoint is a ModelOpt checkpoint and
+        # set a flag to use our model provider if so.
+        if args.load is not None and has_modelopt_state(args.load):
+            print_rank_0(f'ModelOpt checkpoint detected')
+            args.modelopt_enabled = True
+        elif getattr(args, "export_kd_teacher_load", None):
+            # For distillation ckpts without ModelOpt state
+            args.modelopt_enabled = True
+
+
     # Build model.
     def build_model():
         if (
@@ -1081,17 +1093,6 @@ def setup_model_and_optimizer(
     timers = get_timers()
     one_logger = get_one_logger()
 
-    if has_nvidia_modelopt:
-        from megatron.post_training.checkpointing import has_modelopt_state
-        # [ModelOpt]: Check if the checkpoint is a ModelOpt checkpoint and
-        # set a flag to use our model provider if so.
-        if args.load is not None and has_modelopt_state(args.load):
-            print_rank_0(f'ModelOpt checkpoint detected')
-            args.modelopt_enabled = True
-        elif getattr(args, "export_kd_teacher_load", None):
-            # For distillation ckpts without ModelOpt state
-            args.modelopt_enabled = True
-
     model = get_model(model_provider_func, model_type)
     unwrapped_model = unwrap_model(model)
 

tools/run_text_generation_server.py

@@ -32,6 +32,7 @@
 from megatron.core.inference.text_generation_server import MegatronServer
 from megatron.core.inference.text_generation_server.run_mcore_engine import run_mcore_engine
 from megatron.core.transformer.module import MegatronModule
+from megatron.post_training.arguments import add_modelopt_args
 from megatron.training import get_model, print_rank_0
 from model_provider import model_provider
 
@@ -120,6 +121,7 @@ def add_text_generate_args(parser):
         default=None,
         help='Deprecated in favor of `--inference-max-batch-size`',
     )
+    add_modelopt_args(parser)
     return parser