From 0787a380a8c811a18babcef5fd62850b4d6ab451 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:04:46 +0000 Subject: [PATCH 1/6] Evaluate one MMLU subtask in test_lm_eval_hf test_lm_eval_hf ran `--tasks mmlu --limit 0.1`: the 57-subtask group, about 5.6k loglikelihood requests, on a tiny Qwen3 whose max_position_embeddings is 32, so every context is truncated to 32 tokens. Its runtime is dominated by loading one HF dataset config per subtask and does not depend on the host (over 331 CI runs: min 198.9 s, median 212.5 s, mean 221.0 s). Run `mmlu_abstract_algebra` (100 test documents, 5-shot) instead. It goes through the same lm_eval_hf.py code path and the same accuracy gate (which requires exactly one task and reads results[task]['acc,none']); only the number of datasets and documents changes. The gate's lower bound of 0.1 is unchanged. The label distribution of that split is 22/26/31/21, so even a model that always predicts one letter scores at least 0.21. Testing: ran the test's command on CPU (CUDA hidden) with the tiny Qwen3 the test builds. Old configuration: 223.4 s wall, mmlu/acc 0.2343, gate PASS. New configuration: 22.4 to 24.0 s wall, mmlu_abstract_algebra/acc 0.27, gate PASS. Also ran ruff check, ruff format --check and pytest --collect-only on the file. A GPU CI run is still needed to confirm the accuracy with CUDA bf16 numerics. Signed-off-by: Shengliang Xu --- tests/examples/llm_eval/test_llm_eval.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/examples/llm_eval/test_llm_eval.py b/tests/examples/llm_eval/test_llm_eval.py index d65f497e320..ec28ab28aa9 100644 --- a/tests/examples/llm_eval/test_llm_eval.py +++ b/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"), From d06a644b30762a33794747e3a64b14700284fc4b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:05:20 +0000 Subject: [PATCH 2/6] Quantize once per module in the llm_qat QAT and QAD tests test_qwen3_qat_nvfp4 (fsdp2, deepspeed, ddp) and test_qwen3_qad_nvfp4 (fsdp2, deepspeed) each ran quantize.py on the same tiny Qwen3 with the same recipe, and every quantize.py and train.py call built its own tokenized dataset cache (the cache key contains max_length, 8192 in quantize.py and 128 in train.py, so even the two steps of one test could not share it). With test_qwen3_qlora_nvfp4 that is 6 quantize.py runs and 12 cold dataset builds per module. These tests take 104 to 121 s each on CI (qlora 169 s). Run quantize.py once in a module-scoped fixture and let the five tests read the checkpoint, guarded by assert_unmodified_tree like the shared tiny-model fixtures. qat_nvfp4.yaml and qad_nvfp4.yaml give quantize.py identical arguments (distill and teacher_model are not quantize arguments), so one checkpoint serves both. test_qwen3_qlora_nvfp4 keeps its own PTQ because it passes --compress True. All quantize.py and train.py calls now share one dataset cache directory, so the 8192 and 128 variants are each built once; calibration keeps its 8192 tokens. test_dataset_utils_pretokenize keeps its own cache directory so it still checks that the CLI writes one. Fixture time is outside pytest-timeout (timeout_func_only = true), as for the other shared fixtures in tests/examples. Testing: node IDs are unchanged. With run_example_command faked, quantize.py runs 6 -> 2 times and cold dataset builds go 12 -> 2 (plus the pretokenize test's own); a train.py that writes into the shared checkpoint is reported by the fixture teardown. Real quantize.py and train.py (QAT, then QAD from the checkpoint made with the QAT config) ran on CPU with the int4 and fp8 recipes, since NVFP4 dynamic-block quantization needs CUDA: train.py left the checkpoint unmodified and the QAD run reused the dataset cache the QAT run built. The saving, about 3 minutes of the 21.7 minute lane, is an estimate from the CI timings; a GPU CI run is needed to measure it. Signed-off-by: Shengliang Xu --- tests/examples/llm_qat/test_llm_qat.py | 85 +++++++++++++------------- 1 file changed, 44 insertions(+), 41 deletions(-) diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py index a7b610a807c..bec89bd1b97 100644 --- a/tests/examples/llm_qat/test_llm_qat.py +++ b/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 From 27559eeff244cb2cbe9145b997d5c8653486dd1c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:05:31 +0000 Subject: [PATCH 3/6] Skip Triton GEMM autotune in the 1-GPU EAGLE3 compile smoke test test_llama_eagle3[1-False] is the one compiled EAGLE3 training run on the 1-GPU PR lane: hf_eagle.py compiles _eagle_forward with mode="max-autotune". It takes 175.5 s per job against 40.6 s for [1-True], the same flow without compile (call time 158.5 s vs 34.7 s on PRs, 147.8 s vs 34.4 s nightly), so about 135 s is torch.compile. That includes benchmarking Triton GEMM templates for a draft model that is one tiny decoder layer. Launch it with TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=ATEN, the mechanism test_diffusers.py already uses for its compile tests. The mode stays max-autotune, so cudagraphs, coordinate-descent tuning, Dynamo tracing and Inductor lowering still run; only the Triton GEMM template search is dropped. There is no conv variable here because the draft has no convolutions. [2-False] (2 GPUs, nightly) keeps the full autotune. Testing: with run_example_command faked, only [1-False] gets the variable (on top of the full os.environ) and the other parametrizations pass no env. On torch 2.14 torch._inductor.config.max_autotune_gemm_backends reads the variable and the max-autotune mode options do not override it. CPU cannot reproduce the Triton autotune, so the saving is not measured; a GPU CI run is needed to see how much of the 135 s it removes. Signed-off-by: Shengliang Xu --- tests/examples/speculative_decoding/test_eagle.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/examples/speculative_decoding/test_eagle.py b/tests/examples/speculative_decoding/test_eagle.py index 8d284d4ba1b..30f88ec985b 100644 --- a/tests/examples/speculative_decoding/test_eagle.py +++ b/tests/examples/speculative_decoding/test_eagle.py @@ -147,9 +147,17 @@ def test_llama_eagle3(tiny_llama_path, *_TINY_EAGLE_ARCH, ] + env = None + if cp_size == 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. [2-False] (2 GPUs, nightly) keeps 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, ) From 8155b36ce0aba83278ebd8e96716fb8fd98039ef Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:06:01 +0000 Subject: [PATCH 4/6] Drop the duplicate tiny-llama offline EAGLE3 training test test_offline_eagle3_training[tinyllama] trains one step (num_train_epochs=0.1) of offline EAGLE3 on the tiny llama with its real base model. It takes 34.9 s per job, almost all launcher start-up. test_eagle_offline_ptq.py's test_offline_eagle_training (35.2 s) already runs that mode on the same model (use_fake_base_for_offline defaults to false) with hidden states dumped by compute_hidden_states_hf.py, and the FakeBaseModel case keeps the synthetic dump path. The only thing the dropped case covered that stage 2 did not is the offline collator truncating dumps: its 128-token dumps were cut to training_seq_len=64, while stage 2's dumps are 28 tokens against 64, so they were only padded. Set stage 2 to training_seq_len=16 so the collator truncates there. The collator is also unit tested on CPU (test_collator_truncates, test_collator_pads). Stage 3 (hf_ptq.py --specdec_offline_dataset) takes its collator length from --calib_seq and does not read the tokenizer saved with the checkpoint, which now records model_max_length 16. Testing: ran stage 1 (compute_hidden_states_hf.py, two 28-token dumps) and stage 2 (main.py, offline EAGLE3 on the tiny llama) on CPU with training_seq_len=64 and 16 (use_cpu, bf16 and tf32 off are the only changes from the test): both train and write the checkpoint, and the collator pads 28 -> 64 and truncates 28 -> 16. With run_example_command faked, the launch commands of the remaining parametrizations are unchanged and the only removed test ID is test_offline_eagle3_training[tinyllama]; nothing else reads its eagle-tinyllama-offline output. The ~30 s saving per job is not measured. Signed-off-by: Shengliang Xu --- .../examples/speculative_decoding/test_eagle.py | 16 ++++++++-------- .../test_eagle_offline_ptq.py | 3 ++- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/tests/examples/speculative_decoding/test_eagle.py b/tests/examples/speculative_decoding/test_eagle.py index 30f88ec985b..7325fe72e84 100644 --- a/tests/examples/speculative_decoding/test_eagle.py +++ b/tests/examples/speculative_decoding/test_eagle.py @@ -230,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", @@ -243,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 @@ -270,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}", diff --git a/tests/examples/speculative_decoding/test_eagle_offline_ptq.py b/tests/examples/speculative_decoding/test_eagle_offline_ptq.py index 5fca5a3d71c..3bd1be71dba 100644 --- a/tests/examples/speculative_decoding/test_eagle_offline_ptq.py +++ b/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", From 7e4c1b8fbba8dac39fe378fdf1973a232cf5d190 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 9 Oct 2026 23:44:44 +0000 Subject: [PATCH 5/6] Skip the Megatron state reset for megatron_bridge tests that use no GPU The megatron_bridge conftest resets Megatron's global state before and after every test: destroy_model_parallel, destroy_rerun_state_machine, then gc.collect() and torch.cuda.empty_cache(). With torch, Megatron and modelopt imported the collection walks a large heap and costs about 1.3 s per test in the recorded durations. 77 tests in this directory are pure Python (the mlflow wiring tests and the checkpoint-path tests), never run an example step and leave nothing to reset, so the lane pays about 100 s per job for nothing. The megatron_bridge job runs at 68-75 minutes against its 75-minute nightly limit. Reset the Megatron state only for tests that request `num_gpus`, which is how every test that runs an example step reaches a GPU in this directory. The environment restore stays for every test: it is cheap, and the NVTE_* leak it guards against does not depend on a GPU. Tests that run steps keep both resets, and run_example_in_process and run_torchrun_in_process already reset at the start of each step. A future test that runs a step without taking `num_gpus` would lose the reset around it, though its steps would still reset at their start. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. Ran a throwaway pytest module (not part of this change) on CPU against the patched conftest with the real fixture, counting calls to reset_megatron_global_state: plain tests trigger none, tests that request `num_gpus` directly or through another fixture trigger one before and one after, and the environment is restored in both cases. The same module fails against the old conftest. Not run: the megatron_bridge suite itself, because megatron.bridge and the GPUs are not available on this machine. The next GPU CI run should show the setup and teardown time of the mlflow tests drop to near zero. Signed-off-by: Shengliang Xu --- tests/examples/megatron_bridge/conftest.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/examples/megatron_bridge/conftest.py b/tests/examples/megatron_bridge/conftest.py index d8da852f11a..31481fbbdb6 100644 --- a/tests/examples/megatron_bridge/conftest.py +++ b/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) From 5b4b6eab3a554ddcaa58cc4f69217ecf8341af52 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 19:16:33 +0000 Subject: [PATCH 6/6] Limit the EAGLE3 autotune shortcut to single-GPU runs Review feedback on this PR: `cp_size == 1` does not mean one GPU. launch_train.sh starts one accelerate process per visible GPU, so on the 2-GPU nightly runner `[1-False]` runs with two processes and still got TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=ATEN, although the comment says the nightly keeps the full autotune. Key the shortcut on the GPU count instead. Testing: pre-commit passes. Calling test_llama_eagle3 on a CPU host with run_example_command captured: only cp_size=1, mix_hidden_states=False and num_gpus=1 gets the ATEN variable; that case with 2 GPUs, and every other case, keep the full autotune; cp_size=2 on one GPU is still skipped. Not run: the example job itself, which needs GPUs. Signed-off-by: Shengliang Xu --- tests/examples/speculative_decoding/test_eagle.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/examples/speculative_decoding/test_eagle.py b/tests/examples/speculative_decoding/test_eagle.py index 7325fe72e84..df5b00291f0 100644 --- a/tests/examples/speculative_decoding/test_eagle.py +++ b/tests/examples/speculative_decoding/test_eagle.py @@ -148,10 +148,10 @@ def test_llama_eagle3(tiny_llama_path, ] env = None - if cp_size == 1 and not mix_hidden_states: + 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. [2-False] (2 GPUs, nightly) keeps the full autotune. + # run. Runs on more GPUs (the nightly) keep the full autotune. env = {**os.environ, "TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS": "ATEN"} run_example_command(