[OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests

tests/examples/llm_eval/test_llm_eval.py

@@ -32,9 +32,10 @@ def test_lm_eval_hf(tmp_path):
         ["python", "lm_eval_hf.py"],
         model="hf",
         model_args=f"pretrained={model_dir}",
-        tasks="mmlu",
+        # One subtask takes the same harness and gate path as the 57-subtask mmlu group, which
+        # spends minutes loading datasets on a model whose context is only 32 tokens
+        tasks="mmlu_abstract_algebra",
         num_fewshot=5,
-        limit=0.1,
         batch_size=8,
         # Exercise the accuracy gate: reads the results file and enforces the bound
         output_path=str(tmp_path / "results"),

tests/examples/llm_qat/test_llm_qat.py

@@ -19,6 +19,7 @@
 import pytest
 import torch
 from _test_utils.examples.run_command import run_example_command
+from _test_utils.fs_utils import assert_unmodified_tree
 from safetensors.torch import load_file
 
 # Mapping from backend name to accelerate config file
@@ -102,6 +103,38 @@ def _run_export(ckpt_dir: str, export_dir: str):
     )
 
 
+@pytest.fixture(scope="module")
+def dataset_cache_dir(tmp_path_factory):
+    """Tokenized-dataset cache for every quantize/train run below.
+
+    The cache key covers the tokenizer, max length and blend (not the command), so the dataset is
+    built once for quantize.py (8192) and once for train.py (128) instead of once per command.
+    """
+    return str(tmp_path_factory.mktemp("dataset_cache"))
+
+
+@pytest.fixture(scope="module")
+def nvfp4_ptq_dir(tiny_qwen3_path, tmp_path_factory, dataset_cache_dir):
+    """NVFP4 PTQ checkpoint shared by the QAT and QAD tests, which only read it.
+
+    quantize.py ignores the QAD-only config keys, so qat_nvfp4.yaml and qad_nvfp4.yaml give the
+    same arguments and one run serves both.
+    """
+    ptq_dir = tmp_path_factory.mktemp("ptq")
+    _run_quantize(
+        "configs/train/qat_nvfp4.yaml",
+        [
+            "--model_name_or_path", tiny_qwen3_path,
+            "--recipe", "general/ptq/nvfp4_default-kv_fp8",
+            "--calib_size", "64",
+            "--output_dir", str(ptq_dir),
+        ],
+        cache_dir=dataset_cache_dir,
+    )
+    with assert_unmodified_tree(ptq_dir) as path:
+        yield str(path)
+
+
 def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path):
     """Test dataset_utils.py standalone CLI pre-tokenization."""
     cache_dir = tmp_path / "dataset_cache"
@@ -122,33 +155,18 @@ def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path):
     "deepspeed",
     "ddp",
 ])
-def test_qwen3_qat_nvfp4(tiny_qwen3_path, tmp_path, backend):
-    ptq_output_dir = tmp_path / "ptq"
+def test_qwen3_qat_nvfp4(nvfp4_ptq_dir, dataset_cache_dir, tmp_path, backend):
     qat_output_dir = tmp_path / "qat"
-    cache_dir = str(tmp_path / "dataset_cache")
-
-    # Step 1: Quantize
-    _run_quantize(
-        "configs/train/qat_nvfp4.yaml",
-        [
-            "--model_name_or_path", tiny_qwen3_path,
-            "--recipe", "general/ptq/nvfp4_default-kv_fp8",
-            "--calib_size", "64",
-            "--output_dir", str(ptq_output_dir),
-        ],
-        cache_dir=cache_dir,
-    )
 
-    # Step 2: QAT
     _run_train(
         "configs/train/qat_nvfp4.yaml",
         [
-            "--model_name_or_path", str(ptq_output_dir),
+            "--model_name_or_path", nvfp4_ptq_dir,
             "--do_train", "True",
             "--output_dir", str(qat_output_dir),
         ],
         backend=backend,
-        cache_dir=cache_dir,
+        cache_dir=dataset_cache_dir,
     )
 
 @pytest.mark.skip(reason="FSDP2 LoRA checkpoint save omits adapter_model.safetensors")
@@ -211,43 +229,28 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
     "fsdp2",
     "deepspeed",
 ])
-def test_qwen3_qad_nvfp4(tiny_qwen3_path, tmp_path, backend):
-    ptq_output_dir = tmp_path / "ptq"
+def test_qwen3_qad_nvfp4(tiny_qwen3_path, nvfp4_ptq_dir, dataset_cache_dir, tmp_path, backend):
     qad_output_dir = tmp_path / "qad"
-    cache_dir = str(tmp_path / "dataset_cache")
-
-    # Step 1: Quantize student
-    _run_quantize(
-        "configs/train/qad_nvfp4.yaml",
-        [
-            "--model_name_or_path", tiny_qwen3_path,
-            "--recipe", "general/ptq/nvfp4_default-kv_fp8",
-            "--calib_size", "64",
-            "--output_dir", str(ptq_output_dir),
-        ],
-        cache_dir=cache_dir,
-    )
 
-    # Step 2: QAD (quantization-aware distillation)
+    # QAD (quantization-aware distillation) of the shared quantized student
     _run_train(
         "configs/train/qad_nvfp4.yaml",
         [
-            "--model_name_or_path", str(ptq_output_dir),
+            "--model_name_or_path", nvfp4_ptq_dir,
             "--do_train", "True",
             "--output_dir", str(qad_output_dir),
             "--distill", "True",
             "--teacher_model", tiny_qwen3_path,
         ],
         backend=backend,
-        cache_dir=cache_dir,
+        cache_dir=dataset_cache_dir,
     )
 
 
-def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path):
+def test_qwen3_qlora_nvfp4(tiny_qwen3_path, dataset_cache_dir, tmp_path):
     ptq_output_dir = tmp_path / "ptq"
-    cache_dir = str(tmp_path / "dataset_cache")
 
-    # Step 1: Quantize with compression for QLoRA
+    # Step 1: Quantize with compression for QLoRA (its own checkpoint: the others are not compressed)
     _run_quantize(
         "configs/train/qlora_nvfp4.yaml",
         [
@@ -257,7 +260,7 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path):
             "--compress", "True",
             "--output_dir", str(ptq_output_dir),
         ],
-        cache_dir=cache_dir,
+        cache_dir=dataset_cache_dir,
     )
 
     # Step 2: QLoRA training
@@ -271,7 +274,7 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path):
             "--output_dir", str(qlora_output_dir),
         ],
         backend="ddp",
-        cache_dir=cache_dir,
+        cache_dir=dataset_cache_dir,
     )
 
     # Step 3: Export the QLoRA checkpoint for deployment

tests/examples/megatron_bridge/conftest.py

@@ -44,19 +44,23 @@ def _fast_example_runner():
 
 
 @pytest.fixture(autouse=True)
-def _isolate_megatron_global_state():
+def _isolate_megatron_global_state(request):
     """Reset shared state around every test so a failure cannot cascade into the next one.
 
     In-process steps share the interpreter. Besides Megatron's singletons, Transformer-Engine
     records its chosen attention backend in ``NVTE_*``, which failed a Mamba hybrid that ran after
     an attention model -- so the environment is restored wholesale rather than by naming variables.
     """
+    # Only tests that take ``num_gpus`` run example steps; the gc-heavy reset is wasted on the rest.
+    reset_state = "num_gpus" in request.fixturenames
     env_before = os.environ.copy()
-    reset_megatron_global_state()
+    if reset_state:
+        reset_megatron_global_state()
     try:
         yield
     finally:
-        reset_megatron_global_state()
+        if reset_state:
+            reset_megatron_global_state()
         os.environ.clear()
         os.environ.update(env_before)
 

tests/examples/speculative_decoding/test_eagle.py

@@ -147,9 +147,17 @@ def test_llama_eagle3(tiny_llama_path,
         *_TINY_EAGLE_ARCH,
     ]
 
+    env = None
+    if num_gpus == 1 and not mix_hidden_states:
+        # The 1-GPU PR smoke compiles with mode="max-autotune"; benchmarking Triton GEMM templates
+        # is a large share of that on a tiny model and the test only needs the compiled path to
+        # run. Runs on more GPUs (the nightly) keep the full autotune.
+        env = {**os.environ, "TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS": "ATEN"}
+
     run_example_command(
         ["./launch_train.sh", "--config", EAGLE3_YAML, *overrides],
         "speculative_decoding",
+        env=env,
     )
 
 
@@ -222,7 +230,6 @@ def test_convert_to_vllm_ckpt(tiny_llama_path, eagle_output_dir):
 @pytest.mark.parametrize(
     ("model_source", "use_fake_base"),
     [
-        (None, False),  # tiny_llama (from fixture), no FakeBase
         pytest.param("moonshotai/Kimi-K2.5", True),  # remote HF repo, FakeBaseModel
         pytest.param(  # remote HF repo, no FakeBaseModel
             "moonshotai/Kimi-K2-Thinking",
@@ -235,22 +242,23 @@ def test_convert_to_vllm_ckpt(tiny_llama_path, eagle_output_dir):
             marks=pytest.mark.manual(reason="skip redundand test, too slow"),
         ),
     ],
-    ids=["tinyllama", "kimi-k2.5", "kimi-k2-thinking", "minimax-m2.5"],
+    ids=["kimi-k2.5", "kimi-k2-thinking", "minimax-m2.5"],
 )
 def test_offline_eagle3_training(
-    tiny_llama_path,
     tiny_daring_anteater_path,
     tmp_path,
     eagle_output_dir,
     model_source,
     use_fake_base,
 ):
-    """Test Eagle3 training with pre-computed hidden states (offline mode / FakeBaseModel)."""
-    model_path = tiny_llama_path if model_source is None else model_source
-    model_id = "tinyllama" if model_source is None else model_source.split("/")[-1]
+    """Test Eagle3 training with pre-computed hidden states on a FakeBaseModel.
+
+    Offline training on a real base model is covered by test_eagle_offline_ptq.py.
+    """
+    model_id = model_source.split("/")[-1]
     output_subdir = eagle_output_dir / f"eagle-{model_id}-offline"
 
-    cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    cfg = AutoConfig.from_pretrained(model_source, trust_remote_code=True)
     if hasattr(cfg, "text_config"):  # vlm: get text_config
         cfg = cfg.text_config
 
@@ -262,7 +270,7 @@ def test_offline_eagle3_training(
     )
 
     overrides = [
-        f"model.model_name_or_path={model_path}",
+        f"model.model_name_or_path={model_source}",
         "model.trust_remote_code=true",
         f"model.use_fake_base_for_offline={str(use_fake_base).lower()}",
         f"data.data_path={tiny_daring_anteater_path}",

tests/examples/speculative_decoding/test_eagle_offline_ptq.py

@@ -97,7 +97,8 @@ def test_offline_eagle_training(tiny_llama_path, tiny_daring_anteater_path, offl
         f"training.output_dir={output_dir}",
         "training.num_train_epochs=1",
         "training.learning_rate=1e-5",
-        "training.training_seq_len=64",
+        # Shorter than the dumped sequences, so the offline collator truncates them.
+        "training.training_seq_len=16",
         "training.save_steps=1",
         # torch.compile is smoke-tested once by test_llama_eagle3[1-False]; skip its warmup here.
         "eagle.eagle_use_torch_compile=false",