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/_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/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 b45943e6e5e..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 @@ -55,6 +55,7 @@ ) def test_unified_hf_export_and_check_safetensors( tmp_path, + calib_jsonl, qformat, expected_suffix, fuse_input_scale, @@ -103,11 +104,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 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() diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 496db08578a..12b387450d3 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -130,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): +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="cnn_dailymail", + dataset_name=calib_jsonl, tokenizer=tiny_tokenizer, batch_size=2, num_samples=8, 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, )