[OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data

tests/_test_utils/torch/calib_data.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Local calibration text for tests that only need some text, not a particular Hub dataset."""
+
+import json
+import random
+from pathlib import Path
+
+__all__ = ["write_calib_jsonl"]
+
+_TEXT = (
+    "the of and to in is that for it as was with be by on not he this are or his from at which "
+    "but have an had they you were their one all we can her has there been if more when will "
+    "would who so no model layer weight scale range value sample batch token cache quantize"
+)
+_WORDS = _TEXT.split()
+
+
+def write_calib_jsonl(path: Path | str, num_samples: int = 64, words_per_sample: int = 150) -> str:
+    """Write ``num_samples`` random-word texts to ``path`` as ``{"text": ...}`` JSON lines.
+
+    Every text is long enough to fill the default 512-token calibration window with a character
+    level tokenizer. The texts differ from each other and are the same on every call.
+    """
+    rng = random.Random(0)
+    texts = [" ".join(rng.choices(_WORDS, k=words_per_sample)) for _ in range(num_samples)]
+    Path(path).write_text("".join(json.dumps({"text": t}) + "\n" for t in texts), encoding="utf-8")
+    return str(path)

tests/_test_utils/torch/puzzletron/utils.py

@@ -270,9 +270,9 @@ def save_dummy_dataset(dataset_path: Path | str):
         },
     ]
 
-    # Prepare train and val splits with sample repeated, 2500 samples are for
-    # 128 samples with block-size 8192 and LLama3 tokenizer
-    data = [{"conversation": sample}] * 2500
+    # Prepare train and val splits with the sample repeated. With the tiny tokenizer a sample is
+    # 1424 tokens, so 400 samples give 278 blocks of 2048 tokens; the pruning scores read at most 100.
+    data = [{"conversation": sample}] * 400
 
     # For train-val splits
     data_dict = DatasetDict({"train": Dataset.from_list(data), "valid": Dataset.from_list(data)})

tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py

@@ -17,6 +17,7 @@
 import pytest
 import torch
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
+from _test_utils.torch.calib_data import write_calib_jsonl
 from _test_utils.torch.transformers_models import (
     create_tiny_gpt_oss_dir,
     create_tiny_llama_dir,
@@ -26,6 +27,12 @@
 from safetensors import safe_open
 
 
+@pytest.fixture(scope="module")
+def calib_jsonl(tmp_path_factory):
+    """Local calibration texts; the test checks the exported structure, so it needs no Hub dataset."""
+    return write_calib_jsonl(tmp_path_factory.mktemp("calib") / "calib.jsonl")
+
+
 # Here we map each qformat -> the suffix we expect in the generated safetensors directory
 @pytest.mark.parametrize(
     (
@@ -55,6 +62,7 @@
 )
 def test_unified_hf_export_and_check_safetensors(
     tmp_path,
+    calib_jsonl,
     qformat,
     expected_suffix,
     fuse_input_scale,
@@ -103,11 +111,15 @@ def test_unified_hf_export_and_check_safetensors(
         pyt_ckpt_path=tiny_model_dir,
         qformat=qformat,
         export_path=output_dir,
-        dataset="cnn_dailymail",
+        dataset=calib_jsonl,
         # This test only checks the exported safetensors structure (not accuracy), so a
         # small calibration set is enough. Avoids the 1024-sample default on a toy model.
         calib_size=64,
+        # One batch holds the whole set, so hf_ptq does not probe for the largest batch.
+        batch_size=64,
     )
+    # The generate() previews before and after quantization do not affect the export.
+    cmd_parts.append("--skip_generate")
 
     # Run the command
     # current transformers T5 model seem to be incompatible with multiple GPUs

tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct.yaml

@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths

tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-8B/Qwen3-8B.yaml

@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths

tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-VL-30B-A3B-Instruct/Qwen3-VL-30B-A3B-Instruct.yaml

@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths

tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml

@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override: # Optional override for sequence lengths

tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.2-3B-Instruct/Llama-3.2-3B-Instruct.yaml

@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override: # Optional override for sequence lengths

tests/gpu/torch/puzzletron/resources/configs/mistralai/Mistral-Small-24B-Instruct-2501/Mistral-Small-24B-Instruct-2501.yaml

@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths

tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml

@@ -20,7 +20,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths

tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-Nano-12B-v2/NVIDIA-Nemotron-Nano-12B-v2.yaml

@@ -19,7 +19,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override:       # Optional override for sequence lengths

tests/gpu/torch/puzzletron/resources/configs/openai/gpt-oss-20b/gpt-oss-20b.yaml

@@ -20,7 +20,7 @@ build_replacement_library:
   add_attention_no_ops: true
 
 calc_subblock_stats:
-  batch_sizes: [64, 96, 128]
+  batch_sizes: [96]  # must include the batch_size in mip.subblock_stats_args
   prefill_seq_len: 4096
   generation_seq_len: 4096
   num_active_tokens_override: # Optional override for sequence lengths

tests/gpu/torch/puzzletron/test_nas_convert.py

@@ -27,62 +27,6 @@
 import modelopt.torch.utils.distributed as dist
 
 
-def test_nas_convert_ffn_pruning(project_root_path: Path, tmp_path: Path):
-    spawn_multiprocess_job(
-        size=torch.cuda.device_count(),
-        job=partial(_test_nas_convert_ffn_pruning_multiprocess_job, project_root_path, tmp_path),
-        backend="nccl",
-    )
-
-
-def _test_nas_convert_ffn_pruning_multiprocess_job(
-    project_root_path: Path, tmp_path: Path, rank: int, size: int
-):
-    dist.setup(timeout=timedelta(minutes=10))
-    # Setup the test model and data.
-    puzzle_dir, llama_checkpoint_path, dataset_path = setup_test_model_and_data(
-        tmp_path, rank, "meta-llama/Llama-3.1-8B-Instruct"
-    )
-    hydra_config_dir = project_root_path / "tests/gpu/torch/puzzletron/resources/configs"
-    hydra_config_name = "meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct"
-
-    #
-    # Run the mnt.convert() step
-    #
-    input_model = mtpz.puzzletron_nas_plugin.PuzzletronModel()
-    mtn.convert(
-        input_model,
-        mode=[
-            (
-                "puzzletron",
-                {
-                    "puzzle_dir": str(puzzle_dir),
-                    "input_model_path": str(llama_checkpoint_path),
-                    "hydra_config_dir": str(hydra_config_dir),
-                    "hydra_config_name": hydra_config_name,
-                    "dataset_path": str(dataset_path),
-                },
-            )
-        ],
-    )
-
-    #
-    # Check assertions
-    #
-    if rank == 0:
-        # assertions for the score_pruning_activations step
-        rank = int(os.environ["RANK"])
-        rank_filepath = (
-            f"pruning/pruning_scores/ffn_iterative/100samples_diverse_mini/rank_{rank}.pth"
-        )
-        assert (puzzle_dir / rank_filepath).is_file()
-
-        # assertions for the pruning_ckpts step
-        assert (puzzle_dir / "ckpts/ffn_256_attn_no_op").exists()
-
-    dist.cleanup()
-
-
 def test_nas_convert_attn_pruning(project_root_path: Path, tmp_path: Path):
     spawn_multiprocess_job(
         size=torch.cuda.device_count(),

tests/gpu/torch/puzzletron/test_nas_search.py

@@ -76,9 +76,16 @@ def _test_nas_search_multiprocess_job(
     )
 
     #
-    # Check assertions for mtn.search() step
+    # Check assertions for the mtn.convert() and mtn.search() steps
     #
     if rank == 0:
+        # assertions for the score_pruning_activations step of mtn.convert()
+        rank_filepath = "pruning/pruning_scores/ffn_iterative/100samples_diverse_mini/rank_0.pth"
+        assert (puzzle_dir / rank_filepath).is_file()
+
+        # assertions for the pruning_ckpts step of mtn.convert()
+        assert (puzzle_dir / "ckpts/ffn_256_attn_no_op").exists()
+
         # assertions for the build_library_and_stats step
         assert (puzzle_dir / "replacement_library.json").is_file()
         assert (puzzle_dir / "subblock_stats.json").is_file()

tests/gpu/torch/quantization/test_gptq.py

@@ -18,6 +18,7 @@
 
 import pytest
 import torch
+from _test_utils.torch.calib_data import write_calib_jsonl
 from _test_utils.torch.transformers_models import get_tiny_llama
 from conftest import requires_triton
 
@@ -130,14 +131,14 @@ def test_gptq_export_roundtrip():
 @pytest.mark.parametrize(
     "quant_cfg", [mtq.NVFP4_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG, mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG]
 )
-def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer):
+def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer, tmp_path):
     model = get_tiny_llama(vocab_size=tiny_tokenizer.vocab_size).to("cuda")
     model.eval()
 
     quant_cfg = copy.deepcopy(quant_cfg)
     quant_cfg["algorithm"] = {"method": "gptq", "layerwise": {"enable": True}}
     calib_dataloader = get_dataset_dataloader(
-        dataset_name="cnn_dailymail",
+        dataset_name=write_calib_jsonl(tmp_path / "calib.jsonl", num_samples=8),
         tokenizer=tiny_tokenizer,
         batch_size=2,
         num_samples=8,

tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py

@@ -49,7 +49,7 @@ def test_worker_context_falls_back_to_spawn_without_forkserver(monkeypatch):
     strict=False,
     reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
 )
-def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
+def test_megatron_preprocess_data_with_jsonl_path(tmp_path, tiny_tokenizer_path):
     input_jsonl = tmp_path / "minipile.jsonl"
     input_jsonl.write_text(
         "".join(json.dumps({"text": f"Sample document {index}."}) + "\n" for index in range(4)),
@@ -67,7 +67,7 @@ def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
     prefixes = megatron_preprocess_data(
         jsonl_paths=input_jsonl,
         output_dir=tmp_path,
-        tokenizer_name_or_path="gpt2",
+        tokenizer_name_or_path=tiny_tokenizer_path,
         json_keys=["text"],
         workers=1,
     )
@@ -81,7 +81,7 @@ def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
     assert Path(prefixes[0] + ".idx").stat().st_size > 0, "Index file should not be empty"
 
 
-def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path):
+def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path, tiny_tokenizer_path):
     input_path = tmp_path / "data.jsonl"
     input_path.write_text(
         "".join(json.dumps({"text": f"Document {index} " * 100}) + "\n" for index in range(10)),
@@ -91,7 +91,7 @@ def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path):
     common_args = {
         "jsonl_paths": input_path,
         "output_dir": tmp_path,
-        "tokenizer_name_or_path": "gpt2",
+        "tokenizer_name_or_path": tiny_tokenizer_path,
         "json_keys": "text",
         "workers": 2,
     }
@@ -115,13 +115,13 @@ def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path):
     strict=False,
     reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
 )
-def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path):
+def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path, tiny_tokenizer_path):
     common_args = {
         "hf_dataset": "nanotron/minipile_100_samples",
         "hf_split": "train",
         "hf_max_samples_per_split": 100,
         "hf_streaming": True,
-        "tokenizer_name_or_path": "gpt2",
+        "tokenizer_name_or_path": tiny_tokenizer_path,
         "json_keys": "text",
         "workers": 2,
     }
@@ -147,15 +147,17 @@ def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path):
     strict=False,
     reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
 )
-def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_path):
+def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(
+    tmp_path, tiny_tokenizer_path
+):
     args = {
         "hf_dataset": "nanotron/minipile_100_samples",
         "hf_split": "train",
         "hf_max_samples_per_split": 1,
         "hf_streaming": True,
         "max_tokens": 100,
         "output_dir": tmp_path,
-        "tokenizer_name_or_path": "gpt2",
+        "tokenizer_name_or_path": tiny_tokenizer_path,
         "json_keys": "text",
         "max_sequence_length": 16,
         "workers": 2,
@@ -253,7 +255,7 @@ def test_megatron_preprocess_data_reasoning_content(tmp_path):
     )
 
 
-def test_megatron_preprocess_data_with_gzip_input(tmp_path):
+def test_megatron_preprocess_data_with_gzip_input(tmp_path, tiny_tokenizer_path):
     gz_path = tmp_path / "data.jsonl.gz"
     with gzip.open(gz_path, "wt", encoding="utf-8") as f:
         for i in range(10):
@@ -262,7 +264,7 @@ def test_megatron_preprocess_data_with_gzip_input(tmp_path):
     prefixes = megatron_preprocess_data(
         jsonl_paths=gz_path,
         output_dir=tmp_path,
-        tokenizer_name_or_path="gpt2",
+        tokenizer_name_or_path=tiny_tokenizer_path,
         json_keys=["text"],
         workers=1,
         strip_newlines=True,
@@ -324,15 +326,15 @@ def test_megatron_preprocess_data_tool_calls_arguments_normalized(tmp_path):
     strict=False,
     reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
 )
-def test_megatron_preprocess_data_hf_streaming_warning(tmp_path):
+def test_megatron_preprocess_data_hf_streaming_warning(tmp_path, tiny_tokenizer_path):
     # hf_streaming without hf_max_samples_per_split should warn and fall back to non-streaming
     with pytest.warns(UserWarning, match="hf_streaming"):
         megatron_preprocess_data(
             hf_dataset="nanotron/minipile_100_samples",
             hf_split="train",
             hf_streaming=True,
             output_dir=tmp_path,
-            tokenizer_name_or_path="gpt2",
+            tokenizer_name_or_path=tiny_tokenizer_path,
             json_keys=["text"],
             workers=1,
         )