fix(launcher): preserve native SFT recipes in CI

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

scripts/performance/argument_parser.py

@@ -330,8 +330,8 @@ def parse_cli_args():
         "--data",
         type=str,
         default="mock",
-        choices=["mock", "rp2", "squad", "squad_packed", "c4"],
-        help="Dataset type to use",
+        choices=["mock", "rp2", "squad", "squad_packed", "c4", "recipe"],
+        help="Dataset type to use; recipe preserves the library recipe dataset",
     )
     data_args.add_argument(
         "--c4_root",

scripts/performance/run_recipe.py

@@ -69,7 +69,7 @@ def _apply_training_argparse_overrides(config, args):
         config.optimizer.min_lr = args.min_lr
 
     # Scheduler configuration
-    if args.warmup_iters:
+    if args.warmup_iters is not None:
         config.scheduler.lr_warmup_iters = args.warmup_iters
 
     # Checkpoint configuration
@@ -95,7 +95,9 @@ def _apply_training_argparse_overrides(config, args):
         pad_seq_to_mult = cp_size * 2 if cp_size > 1 else 1
 
         # Create dataset configuration based on type
-        if args.data == "mock":
+        if args.data == "recipe":
+            pass  # Preserve dataset type, native packing, and pinned source metadata.
+        elif args.data == "mock":
             config.dataset = create_mock_dataset_config(seq_length=args.seq_length or 8192)
         elif args.data == "rp2":
             if not args.dataset_paths or not args.index_mapping_dir:
@@ -196,6 +198,8 @@ def _apply_training_argparse_overrides(config, args):
         config.checkpoint.save_interval = args.save_interval or 1000
 
     # Validation configuration for convergence
+    if args.data == "recipe":
+        return config  # Keep the recipe validation cadence and LR horizon.
     if args.max_steps <= 100:
         config.train.eval_interval = args.max_steps
         config.train.eval_iters = 0  # Disable evaluation for short convergence runs
@@ -312,6 +316,8 @@ def _run_training(args, cli_overrides: list[str]) -> None:
 
         if args.model_family_name in DIFFUSION_FAMILIES:
             forward_step = _get_diffusion_step(args.model_family_name)
+        elif args.domain == "vlm":
+            from megatron.bridge.training.vlm_step import forward_step
         else:
             from megatron.bridge.training.gpt_step import forward_step
 
@@ -322,6 +328,8 @@ def _run_training(args, cli_overrides: list[str]) -> None:
 
         if args.model_family_name in DIFFUSION_FAMILIES:
             forward_step = _get_diffusion_step(args.model_family_name)
+        elif args.domain == "vlm":
+            from megatron.bridge.training.vlm_step import forward_step
         else:
             from megatron.bridge.training.gpt_step import forward_step
 

scripts/performance/utils/utils.py

@@ -890,7 +890,9 @@ def build_recipe_config(model_family_name: str, model_recipe_name: str, train_ta
     family_pkg_path = f"megatron.bridge.recipes.{model_family_name}"
     family_pkg = importlib.import_module(family_pkg_path)
 
-    if model_recipe_name == "deepseek_v3_32nodes" and train_task == "pretrain":
+    if model_recipe_name.endswith("_config"):
+        pass  # Complete exported recipe names already identify task, hardware, and precision.
+    elif model_recipe_name == "deepseek_v3_32nodes" and train_task == "pretrain":
         model_recipe_name = "deepseek_v3_pretrain_config_32nodes"
     elif train_task == "peft":
         model_recipe_name = f"{model_recipe_name}_peft_config"

tests/unit_tests/scripts/performance/test_recipe_environment.py

@@ -1586,3 +1586,73 @@ def exec_runner(executable, argv, env):
     assert calls[0] == ("environment", effective_recipe)
     assert calls[1][0] == "exec"
     assert calls[1][2][1].endswith("run_recipe.py")
+
+
+def test_complete_library_recipe_name_is_not_suffixed(monkeypatch):
+    """CI can select a hardware-specific library recipe without renaming it."""
+    from megatron.bridge.recipes.common import _sft_common
+
+    recipe = _sft_common()
+    name = "qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_fp8mx_config"
+    module = SimpleNamespace(**{name: lambda: recipe})
+    monkeypatch.setattr(utils.importlib, "import_module", lambda _: module)
+    result = utils.build_recipe_config("qwen", name, "sft", "ci-test")
+    assert result is recipe
+    assert result.logger.wandb_exp_name == "ci-test"
+
+
+@pytest.mark.parametrize("dataset_kind", ["text", "energon"])
+def test_recipe_data_keeps_native_packing_tokenizer_and_schedule(dataset_kind):
+    """A 100-step CI run must not replace SFT inputs or compress the LR schedule."""
+    from argument_parser import parse_cli_args
+
+    from megatron.bridge.data.builders import EnergonDatasetConfig, QwenVLEnergonTaskEncoderConfig
+    from megatron.bridge.recipes.common import _sft_common
+    from megatron.bridge.recipes.utils.dataset_utils import default_coderforge_config
+
+    args = parse_cli_args().parse_args(
+        [
+            "--model_family_name",
+            "qwen",
+            "--model_recipe_name",
+            "test_config",
+            "--num_gpus",
+            "16",
+            "--gpu",
+            "gb200",
+            "--task",
+            "sft",
+            "--use_recipes",
+            "--data",
+            "recipe",
+            "--max_steps",
+            "100",
+            "--warmup_iters",
+            "200",
+        ]
+    )
+    cfg = _sft_common()
+    cfg.dataset = (
+        EnergonDatasetConfig(
+            seq_length=131072,
+            micro_batch_size=1,
+            packing_buffer_size=64,
+            task_encoder=QwenVLEnergonTaskEncoderConfig(hf_processor_path="test-model"),
+        )
+        if dataset_kind == "energon"
+        else default_coderforge_config(seq_length=131072, enable_offline_packing=True)
+    )
+    dataset, tokenizer = cfg.dataset, cfg.tokenizer
+    cfg.scheduler.lr_decay_iters = 300000
+    cfg.validation.eval_interval = 50
+    cfg.validation.eval_iters = 8
+    result = run_recipe._apply_training_argparse_overrides(cfg, args)
+    assert result.dataset is dataset
+    assert result.tokenizer is tokenizer
+    assert result.train.train_iters == 100
+    assert result.scheduler.lr_warmup_iters == 200
+    assert result.scheduler.lr_decay_iters == 300000
+    assert result.validation.eval_interval == 50
+    assert result.validation.eval_iters == 8
+    assert result.train.eval_iters is None
+    assert result.train.eval_interval is None

tests/unit_tests/scripts/performance/test_run_script_entrypoint.py

@@ -286,3 +286,23 @@ def test_run_recipe_defers_training_framework_imports():
 
     assert "torch" not in top_level_imports
     assert not any(module.startswith("megatron.bridge") for module in top_level_imports)
+
+
+@pytest.mark.parametrize("domain", ["llm", "vlm"])
+def test_library_sft_dispatches_the_selected_domain(monkeypatch, domain):
+    """VL CI recipes must reach the image-aware training step."""
+    import importlib
+
+    from megatron.bridge.training.gpt_step import forward_step as gpt_step
+    from megatron.bridge.training.vlm_step import forward_step as vlm_step
+
+    cfg = object()
+    calls = []
+    monkeypatch.setattr(run_recipe, "_prepare_recipe", lambda *_args, **_kwargs: cfg)
+    common = importlib.import_module("megatron.bridge.utils.common_utils")
+    finetune_module = importlib.import_module("megatron.bridge.training.finetune")
+    monkeypatch.setattr(common, "get_rank_safe", lambda: 1)
+    monkeypatch.setattr(finetune_module, "finetune", lambda **kwargs: calls.append(kwargs))
+    args = SimpleNamespace(dryrun=False, task="sft", model_family_name="qwen_vl", domain=domain)
+    run_recipe._run_training(args, [])
+    assert calls == [{"config": cfg, "forward_step_func": vlm_step if domain == "vlm" else gpt_step}]