From 12ab37f574bc1f00caa55d4add60863b22fa8177 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 18:20:02 +0000 Subject: [PATCH 1/5] Use the tiny local tokenizer in six preprocess_data tests Six tests in test_megatron_preprocess_data.py tokenize with gpt2 from the Hugging Face Hub. megatron_preprocess_data loads the tokenizer once for the vocabulary size and again in every worker, so each test makes several hub round trips, and none of them depends on gpt2's vocabulary: they check output naming, non-empty .bin and .idx files, and that max_tokens shortens the output. Switch them to the tiny_tokenizer_path fixture, the checked-in byte-level tokenizer that the blend test already uses. Three of the six read only local data (the jsonl path, the jsonl max_tokens and the gzip tests), so the tokenizer download was their only network access. The PR history lists the jsonl max_tokens test, with the HF split max_tokens test, among the lane's 120 s timeouts (3 failures each); it no longer has a network dependency to time out on. The other three (HF split max_tokens, HF resume and the HF streaming warning) still read the real HF dataset, which is the code path they exist to cover; only their tokenizer changes. The Qwen3 tests keep their real tokenizer: they need its chat template and EOS handling. The blend test already uses the tiny tokenizer, and its HF split source is the path it covers, so it is unchanged. No test ids, assertions or markers change. Measured on a CPU host with a cold Hugging Face cache, the module's 11 tests take 42.8 and 44.0 s before and 35.0 and 36.1 s after (two runs each, with the conftest's forkserver prestart). Per test: jsonl path 2.7-3.4 s to 0.01 s, gzip 1.4 s to 0.05 s, HF split max_tokens 5.8-5.9 s to 4.3 s, HF resume 4.4 s to 2.8-3.0 s, HF streaming warning 3.6-3.8 s to 2.3-2.7 s. The jsonl max_tokens test takes about 14 s in both, because here it is the first test to wait for the forkserver's cold import; in the lane the session fixture starts that import earlier. CI jobs start without a warm hub cache, so a similar saving is expected. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. Ran the module on CPU, next to a copy of the tiny_tokenizer_path fixture and the forkserver prestart from tests/gpu_megatron/conftest.py, with the Hub reachable: 11 passed before and after, and the collected test ids are identical. Not run: the gpu_megatron lane itself, which needs GPUs and prebuilds CUDA extensions in its conftest. Signed-off-by: Shengliang Xu --- .../plugins/test_megatron_preprocess_data.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py index b676e1e3f85..c70f4fa2631 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py @@ -40,7 +40,7 @@ def test_worker_context_falls_back_to_spawn_without_forkserver(monkeypatch): _prestart_forkserver() -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)), @@ -58,7 +58,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, ) @@ -72,7 +72,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)), @@ -82,7 +82,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, } @@ -100,13 +100,13 @@ def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path): assert limited_size < full_size -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, } @@ -126,7 +126,9 @@ def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path): assert limited_size < full_size -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", @@ -134,7 +136,7 @@ def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_pa "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, @@ -226,7 +228,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): @@ -235,7 +237,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, @@ -291,7 +293,7 @@ def test_megatron_preprocess_data_tool_calls_arguments_normalized(tmp_path): ) -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( @@ -299,7 +301,7 @@ def test_megatron_preprocess_data_hf_streaming_warning(tmp_path): 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, ) From 038232cbd734ff1c343ac94d22a234829b8e85c6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:02:41 +0000 Subject: [PATCH 2/5] Shrink the puzzletron test data and fold a duplicate NAS test into another - save_dummy_dataset wrote 2500 copies of one conversation into each of the train and valid splits. The tiny test tokenizer turns a conversation into 1424 tokens, so that is 1738 blocks of 2048 tokens, while the pruning scores read 100 blocks and the scoring and realize steps 2 each. Every dataloader build tokenizes the whole split, and a run builds three of them. The 2500 came from an older 8192-token Llama-3 setup. Use 400 conversations (278 blocks). On a CPU the three builds plus the dataset save took 2.8 s per test with 2500 conversations and 0.36 s with 400. The scoring step checks lm_loss with a tolerance of 0.15; on random-init toy Llamas the validation loss moved by at most 0.02 between the two sizes, because the same text repeats and only the choice of blocks changes. - calc_subblock_stats ran 12 passes (batch sizes 64, 96 and 128 times four dtype sets), but the MIP and the teacher memory and parameter checks only read the entry for the batch_size in mip.subblock_stats_args, which is 96 in all nine configs. batch_sizes: [96] gives 4 passes (0.56 s to 0.18 s for a toy Llama on a CPU; MoE configs have more subblocks per pass). - test_nas_convert_ffn_pruning ran the same mtn.convert call on the same toy Llama and config as test_nas_search (21.6 s in the CI logs, with its own 2-rank spawn). Its two assertions move into test_nas_search, which checks them with the other rank-0 assertions after the search. Left alone: the Nemotron-3 toy keeps its 128 experts, because shrinking it changes EXPECTED_TEACHER_MEMORY_MIB and EXPECTED_TEACHER_NUM_PARAMS, which need a GPU runner with mamba_ssm to regenerate. Testing: timed the dataloader builds, the subblock-stats passes and the loss sensitivity on CPU with the repository's own helpers and dataloader code, collected tests/gpu/torch/puzzletron (16 tests, one fewer than before), and ran ruff. Not run: the puzzletron tests themselves, which need GPUs and mamba_ssm. Expect the gpu job to lose roughly 50 s (the removed test is 22 s of it); the CI run will show the real number. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/puzzletron/utils.py | 6 +- .../Qwen2.5-7B-Instruct.yaml | 2 +- .../configs/Qwen/Qwen3-8B/Qwen3-8B.yaml | 2 +- .../Qwen3-VL-30B-A3B-Instruct.yaml | 2 +- .../Llama-3.1-8B-Instruct.yaml | 2 +- .../Llama-3.2-3B-Instruct.yaml | 2 +- .../Mistral-Small-24B-Instruct-2501.yaml | 2 +- ...DIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml | 2 +- .../NVIDIA-Nemotron-Nano-12B-v2.yaml | 2 +- .../openai/gpt-oss-20b/gpt-oss-20b.yaml | 2 +- .../gpu/torch/puzzletron/test_nas_convert.py | 56 ------------------- tests/gpu/torch/puzzletron/test_nas_search.py | 9 ++- 12 files changed, 20 insertions(+), 69 deletions(-) diff --git a/tests/_test_utils/torch/puzzletron/utils.py b/tests/_test_utils/torch/puzzletron/utils.py index 2ce44858ef5..525b597518d 100644 --- a/tests/_test_utils/torch/puzzletron/utils.py +++ b/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)}) diff --git a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct.yaml b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct.yaml index 2843f0b97ac..6691a65d9d7 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen2.5-7B-Instruct/Qwen2.5-7B-Instruct.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-8B/Qwen3-8B.yaml b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-8B/Qwen3-8B.yaml index cd82a47271a..59e69520a4b 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-8B/Qwen3-8B.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-VL-30B-A3B-Instruct/Qwen3-VL-30B-A3B-Instruct.yaml b/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-VL-30B-A3B-Instruct/Qwen3-VL-30B-A3B-Instruct.yaml index 00b21ea9798..77567a777f9 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/Qwen/Qwen3-VL-30B-A3B-Instruct/Qwen3-VL-30B-A3B-Instruct.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml b/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml index 8e2e0786b33..880f630d83b 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.2-3B-Instruct/Llama-3.2-3B-Instruct.yaml b/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.2-3B-Instruct/Llama-3.2-3B-Instruct.yaml index 78cb6bd73c4..ab45abe41f9 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.2-3B-Instruct/Llama-3.2-3B-Instruct.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/resources/configs/mistralai/Mistral-Small-24B-Instruct-2501/Mistral-Small-24B-Instruct-2501.yaml b/tests/gpu/torch/puzzletron/resources/configs/mistralai/Mistral-Small-24B-Instruct-2501/Mistral-Small-24B-Instruct-2501.yaml index e042c4bb627..a6593200354 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/mistralai/Mistral-Small-24B-Instruct-2501/Mistral-Small-24B-Instruct-2501.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml b/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml index ab2b09e6794..33a8318215d 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-Nano-12B-v2/NVIDIA-Nemotron-Nano-12B-v2.yaml b/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-Nano-12B-v2/NVIDIA-Nemotron-Nano-12B-v2.yaml index 906b7338d8c..05b22051392 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/nvidia/NVIDIA-Nemotron-Nano-12B-v2/NVIDIA-Nemotron-Nano-12B-v2.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/resources/configs/openai/gpt-oss-20b/gpt-oss-20b.yaml b/tests/gpu/torch/puzzletron/resources/configs/openai/gpt-oss-20b/gpt-oss-20b.yaml index 60d6568df60..d57ff295796 100644 --- a/tests/gpu/torch/puzzletron/resources/configs/openai/gpt-oss-20b/gpt-oss-20b.yaml +++ b/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 diff --git a/tests/gpu/torch/puzzletron/test_nas_convert.py b/tests/gpu/torch/puzzletron/test_nas_convert.py index 57239f1c0e4..7ddb0584fc7 100644 --- a/tests/gpu/torch/puzzletron/test_nas_convert.py +++ b/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(), diff --git a/tests/gpu/torch/puzzletron/test_nas_search.py b/tests/gpu/torch/puzzletron/test_nas_search.py index 2cdc3e26d44..cd3e3f967b2 100644 --- a/tests/gpu/torch/puzzletron/test_nas_search.py +++ b/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() From 2254e5b7ea3bca0c39eaa4c41762bbe0091233cb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:03:30 +0000 Subject: [PATCH 3/5] Calibrate the unified HF export test on local text, with no previews or batch probing Each of the 11 test_unified_hf_export_and_check_safetensors params starts a fresh hf_ptq process, 22 s per param in the CI logs. The test only checks the structure of the exported safetensors, so three parts of that run are not needed: - Calibration text came from cnn_dailymail on the Hub, a network round trip per param. Use a local jsonl of 64 random-word texts instead, written by a new helper, write_calib_jsonl, in _test_utils/torch/calib_data.py. Each text fills the 512-token calibration window and differs from the others, so the activation statistics stay varied. - hf_ptq probed for the largest calibration batch. --batch_size 64 puts the whole calibration set (calib_size 64) in one batch. - hf_ptq ran model.generate() before and after quantization. --skip_generate skips those previews; the export is not affected. Testing: no GPU here, so I ran the test on a CPU host through a python shim that lets hf_ptq pass its GPU guard and passes --device cpu. 8 of the 11 params pass with the new command line, scale-equality assertions included: fp8, nvfp4, int4_awq, w4a8_awq_beta, int8_weight_only, w4a16_nvfp4, the Qwen3 MoE nvfp4 and the GPT-OSS fp8. The other three (nvfp4 AWQ-lite, SVDQuant and the MSE sweep) need CUDA or Triton and did not run. The dataloader returns one (64, 512) batch with a full attention mask from the helper's file. Not run: the GPU job itself. From the per-param CI timings this should save about 38 s per `gpu` job. Signed-off-by: Shengliang Xu --- tests/_test_utils/torch/calib_data.py | 41 +++++++++++++++++++ ...unified_hf_export_and_check_safetensors.py | 14 ++++++- 2 files changed, 54 insertions(+), 1 deletion(-) create mode 100644 tests/_test_utils/torch/calib_data.py diff --git a/tests/_test_utils/torch/calib_data.py b/tests/_test_utils/torch/calib_data.py new file mode 100644 index 00000000000..03622986d82 --- /dev/null +++ b/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) diff --git a/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py b/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py index b45943e6e5e..7a67c026b3b 100644 --- a/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py +++ b/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 From 612c1609aedff4c5892a866f1d3b4eeedbb98e65 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 00:03:50 +0000 Subject: [PATCH 4/5] Calibrate the GPTQ end-to-end test on local text test_gptq_e2e_flow streamed cnn_dailymail from the Hub for each of its three params, a network round trip per param (about 6 s per job in the CI logs) and a test that can fail on a Hub hiccup. The GPTQ calibration only needs some text, and the test has no accuracy assertion, so it now reads 8 texts from the local jsonl that write_calib_jsonl (added with the unified-export change) creates. Testing: ruff check and format with the pinned ruff 0.15.20. On a CPU host the GPTQ flow completes on that jsonl for the FP8 and INT4 blockwise configs; the NVFP4 config needs CUDA. Not run: the GPU test itself. Expected saving is a few seconds per gpu job. Signed-off-by: Shengliang Xu --- tests/gpu/torch/quantization/test_gptq.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 496db08578a..abe3386064a 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/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, From 55e1c27450585b5ddd3ccb62aab8b69063ee38fb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 19:04:47 +0000 Subject: [PATCH 5/5] Share one read-only calibration file across the gpu tests Review feedback on this PR: the calibration texts are never modified by a test, so one session-wide fixture can serve every test that needs some text instead of each test or module writing its own file. `calib_jsonl` in tests/gpu/conftest.py writes the 64 texts once per session and wraps the directory in assert_unmodified_tree, the check the shared tiny-tokenizer fixture already uses (chmod would not stop a root CI container). The unified HF export test drops its module-scoped fixture of the same name, and the GPTQ end-to-end test takes the fixture instead of writing 8 texts into its tmp_path. It still asks the dataloader for 8 samples, and the file starts with the same 8 texts the old 8-text file held, since the generator is seeded and sequential. tests/gpu_megatron has no test that calibrates on text (its tests use model.get_input()), so it gets no copy of the fixture. Testing: pre-commit passes. On a CPU host, with the real fixture and the repo's get_dataset_dataloader: the 8 samples read from the shared file equal those of the old 8-text file, and two reads leave the directory unmodified (the fixture's teardown check passes). `--setup-plan` shows the unified export test using the conftest fixture. Not run: the GPTQ and export tests themselves, which need a GPU. Signed-off-by: Shengliang Xu --- tests/gpu/conftest.py | 14 ++++++++++++++ ...test_unified_hf_export_and_check_safetensors.py | 7 ------- tests/gpu/torch/quantization/test_gptq.py | 5 ++--- 3 files changed, 16 insertions(+), 10 deletions(-) diff --git a/tests/gpu/conftest.py b/tests/gpu/conftest.py index 40a272e8838..408743727c1 100644 --- a/tests/gpu/conftest.py +++ b/tests/gpu/conftest.py @@ -15,6 +15,8 @@ import pytest import torch +from _test_utils.fs_utils import assert_unmodified_tree +from _test_utils.torch.calib_data import write_calib_jsonl from _test_utils.torch.distributed.utils import ( DistributedWorkerPool, default_worker_teardown, @@ -41,3 +43,15 @@ def dist_workers(_dist_pools): world_size=world_size, backend="nccl", teardown_fn=default_worker_teardown ), ) + + +@pytest.fixture(scope="session") +def calib_jsonl(tmp_path_factory): + """Local calibration texts for tests that only need some text, not a particular Hub dataset. + + Shared by the whole session: a test that writes into its directory fails at teardown. + """ + calib_dir = tmp_path_factory.mktemp("calib") + path = write_calib_jsonl(calib_dir / "calib.jsonl") + with assert_unmodified_tree(calib_dir): + yield path diff --git a/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py b/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py index 7a67c026b3b..c7c3fb743f0 100644 --- a/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py +++ b/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py @@ -17,7 +17,6 @@ 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, @@ -27,12 +26,6 @@ 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( ( diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index abe3386064a..12b387450d3 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -18,7 +18,6 @@ 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 @@ -131,14 +130,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, tmp_path): +def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer, calib_jsonl): 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=write_calib_jsonl(tmp_path / "calib.jsonl", num_samples=8), + dataset_name=calib_jsonl, tokenizer=tiny_tokenizer, batch_size=2, num_samples=8,