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"), 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 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) diff --git a/tests/examples/speculative_decoding/test_eagle.py b/tests/examples/speculative_decoding/test_eagle.py index 8d284d4ba1b..df5b00291f0 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 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}", 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",