diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml index 8adadbac7af..c7a04e17ef9 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.github/workflows/_example_tests_runner.yml @@ -47,8 +47,12 @@ jobs: echo "PATH=${PATH}:/usr/local/tensorrt/targets/x86_64-linux-gnu/bin" >> $GITHUB_ENV - name: Install dependencies run: | - # use `python -m pip` instead of `pip` to avoid conflicts with system pip for nemo containers - python -m pip install ".${{ inputs.pip_install_extras }}" + # Uninstall conflicting system-wide installed modelopt in nemo containers + pip uninstall -y nvidia-modelopt || true + + # Use `python -m pip` instead of `pip` to avoid conflicts with system pip for nemo containers + # Editable install so example scripts launched as subprocesses resolve modelopt to the same source path as the test process + python -m pip install -e ".${{ inputs.pip_install_extras }}" if [[ "${{ inputs.example }}" == *"diffusers"* ]]; then echo "Uninstalling apex for diffusers: T5 Int8 (PixArt) + Apex is not supported as per https://github.com/huggingface/transformers/issues/21391" diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 4bcbf3502e8..ef968dacbbe 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -35,14 +35,14 @@ jobs: strategy: fail-fast: false matrix: - example: [llm_distill, llm_qat, llm_sparsity, diffusers_sparsity, specdec_bench] + example: [diffusers_sparsity, gpt-oss, llm_distill, llm_qat, llm_sparsity, specdec_bench] include: - example: speculative_decoding docker_image: "26.01" uses: ./.github/workflows/_example_tests_runner.yml secrets: inherit with: - docker_image: "nvcr.io/nvidia/pytorch:${{ matrix.docker_image || '26.04' }}-py3" + docker_image: "nvcr.io/nvidia/pytorch:${{ matrix.docker_image || '26.05' }}-py3" example: ${{ matrix.example }} timeout_minutes: 30 pip_install_extras: "[hf,dev-test]" @@ -59,7 +59,7 @@ jobs: uses: ./.github/workflows/_example_tests_runner.yml secrets: inherit with: - docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc16" + docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc17" example: ${{ matrix.example }} pip_install_extras: "[hf,dev-test]" runner: linux-amd64-gpu-rtxpro6000-latest-1 @@ -73,7 +73,7 @@ jobs: uses: ./.github/workflows/_example_tests_runner.yml secrets: inherit with: - docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc16" + docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc17" example: ${{ matrix.example }} pip_install_extras: "[hf,dev-test]" runner: linux-amd64-gpu-rtxpro6000-latest-2 @@ -102,8 +102,9 @@ jobs: uses: ./.github/workflows/_example_tests_runner.yml secrets: inherit with: - docker_image: "nvcr.io/nvidia/tensorrt:26.04-py3" + docker_image: "nvcr.io/nvidia/tensorrt:26.05-py3" example: ${{ matrix.example }} + timeout_minutes: 45 pip_install_extras: "[onnx,hf,dev-test]" runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index ca8d7e6d0b6..ab23036e0f4 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -39,16 +39,16 @@ jobs: matrix: include: - example: gpu - timeout: 75 - container_image: nvcr.io/nvidia/pytorch:26.04-py3 + timeout: 60 + container_image: nvcr.io/nvidia/pytorch:26.05-py3 - example: gpu_megatron timeout: 60 container_image: nvcr.io/nvidia/nemo:26.04 - example: gpu_trtllm timeout: 30 - container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc16 + container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc17 - example: gpu_vllm - timeout: 30 + timeout: 15 container_image: docker.io/vllm/vllm-openai:v0.20.0 runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} timeout-minutes: ${{ matrix.timeout }} diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 26dc2108063..fc2ae364cba 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -58,7 +58,7 @@ jobs: linux: needs: [check-dco] runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 15 steps: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup @@ -78,7 +78,7 @@ jobs: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] runs-on: windows-latest - timeout-minutes: 30 + timeout-minutes: 15 steps: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 @@ -90,7 +90,7 @@ jobs: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 15 strategy: fail-fast: false matrix: @@ -115,7 +115,7 @@ jobs: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 15 strategy: fail-fast: false matrix: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 541ab6a51d4..66656ed5e22 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -163,14 +163,20 @@ nox -s "unit-3.12(torch_211, tf_latest)" ### Test design principles -- **Develop with focused tests.** During development, write as many focused - tests as needed, including lower-level unit tests or internal probes, to - understand and harden behavior. -- **Curate production tests and keep them lean.** Before staging or committing, - decide which tests should be checked in. Checked-in tests should document - expected behavior, protect against regressions, or flag backward-incompatible - behavior changes. Remove redundant lower-level tests when a higher-level test - already covers the same behavior, keeping CI/CD fast and lean. +- **Develop with focused tests.** During development, write as many focused tests as needed, including lower-level + unit tests or internal probes, to understand and harden behavior. +- **Curate production tests and keep them lean.** Before staging or committing, decide which tests should be checked + in. Checked-in tests should document expected behavior, protect against regressions, or flag backward-incompatible + behavior changes. Remove redundant lower-level tests when a higher-level test already covers the same behavior, + keeping CI/CD fast and lean. +- **Keep `tests/unit` offline — no HuggingFace Hub access.** Unit tests must be hermetic so they never flake on + network/timeout issues. Do not call `from_pretrained("/")`, `load_dataset("")`, + `snapshot_download(...)`, etc. with Hub IDs. Instead build dummy models, tokenizers, configs, and datasets locally — + e.g. the `create_tiny_*` helpers and `get_tiny_tokenizer()` in `tests/_test_utils/`, or a small on-disk dataset + directory written with `datasets.Dataset.from_dict(...).to_parquet(...)`. +- **Respect the per-test timeout.** `tests/conftest.py` applies a default per-test call timeout by directory; override a + single slow test with `@pytest.mark.timeout()`, and register any new top-level `tests//` in that + mapping (collection errors until you do). ## ✍️ Signing your work diff --git a/examples/diffusers/quantization/diffusion_trt.py b/examples/diffusers/quantization/diffusion_trt.py index 2ae32ea8ee7..125d285be95 100644 --- a/examples/diffusers/quantization/diffusion_trt.py +++ b/examples/diffusers/quantization/diffusion_trt.py @@ -65,14 +65,14 @@ @torch.inference_mode() -def generate_image(pipe, prompt, image_name, torch_autocast=False): +def generate_image(pipe, prompt, image_name, torch_autocast=False, num_inference_steps=30): context = torch.autocast("cuda") if torch_autocast else nullcontext() seed = 42 with context: image = pipe( prompt, output_type="pil", - num_inference_steps=30, + num_inference_steps=num_inference_steps, generator=torch.Generator("cuda").manual_seed(seed), ).images[0] image.save(image_name) @@ -186,6 +186,12 @@ def main(): help="Use torch.autocast() during inference or benchmarking", ) parser.add_argument("--skip-image", action="store_true", help="Skip image generation") + parser.add_argument( + "--num-inference-steps", + type=int, + default=30, + help="Number of denoising steps for image generation (lower is faster; tests use few).", + ) args = parser.parse_args() image_name = args.save_image_as if args.save_image_as else f"{args.model}.png" @@ -235,7 +241,9 @@ def main(): ) if not args.skip_image: - generate_image(pipe, args.prompt, image_name, args.torch_autocast) + generate_image( + pipe, args.prompt, image_name, args.torch_autocast, args.num_inference_steps + ) return backbone.to("cuda") @@ -322,7 +330,7 @@ def main(): pipe.to("cuda") if not args.skip_image: - generate_image(pipe, args.prompt, image_name, args.torch_autocast) + generate_image(pipe, args.prompt, image_name, args.torch_autocast, args.num_inference_steps) print(f"Image generated using {args.model} model saved as {image_name}") if args.benchmark: diff --git a/examples/llm_eval/run_simple_eval.sh b/examples/llm_eval/run_simple_eval.sh index 2240a5567d0..5f40b4ce8b3 100644 --- a/examples/llm_eval/run_simple_eval.sh +++ b/examples/llm_eval/run_simple_eval.sh @@ -22,6 +22,7 @@ MODEL_NAME=$1 EVALS=$2 BUILD_MAX_OUTPUT_LEN=${3:-2048} PORT=${4:-8000} +NUM_EXAMPLES=${5:-} # optional: limit examples per eval (default: full eval) if [ ! -d "human-eval" ]; then git clone https://github.com/openai/human-eval.git @@ -42,4 +43,9 @@ popd export OPENAI_API_KEY="local" export OPENAI_BASE_URL="http://localhost:$PORT/v1" -python -m simple-evals.simple_evals --model $MODEL_NAME --evals $EVALS --max_tokens $BUILD_MAX_OUTPUT_LEN +examples_flag="" +if [ -n "$NUM_EXAMPLES" ]; then + examples_flag="--examples $NUM_EXAMPLES" +fi + +python -m simple-evals.simple_evals --model $MODEL_NAME --evals $EVALS --max_tokens $BUILD_MAX_OUTPUT_LEN $examples_flag diff --git a/examples/llm_ptq/scripts/huggingface_example.sh b/examples/llm_ptq/scripts/huggingface_example.sh index 5c7889cc341..541c349da08 100755 --- a/examples/llm_ptq/scripts/huggingface_example.sh +++ b/examples/llm_ptq/scripts/huggingface_example.sh @@ -328,7 +328,7 @@ if [[ $TASKS =~ "livecodebench" || $TASKS =~ "simple_eval" ]]; then if [[ $TASKS =~ "simple_eval" ]]; then echo "Using the following config: max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE" - bash run_simple_eval.sh $MODEL_NAME $SIMPLE_EVAL_TASKS $BUILD_MAX_OUTPUT_LEN $PORT | tee $SAVE_PATH/simple_eval.txt + bash run_simple_eval.sh $MODEL_NAME $SIMPLE_EVAL_TASKS $BUILD_MAX_OUTPUT_LEN $PORT $SIMPLE_EVAL_LIMIT | tee $SAVE_PATH/simple_eval.txt echo "Simple eval results are saved under $SAVE_PATH/simple_eval.txt." fi diff --git a/examples/llm_ptq/scripts/parser.sh b/examples/llm_ptq/scripts/parser.sh index 2a9a28b3566..3be7706c4e6 100644 --- a/examples/llm_ptq/scripts/parser.sh +++ b/examples/llm_ptq/scripts/parser.sh @@ -38,7 +38,7 @@ parse_options() { CAST_MXFP4_TO_NVFP4=false # Parse command-line options - ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,auto_quantize_bits:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_checkpoint:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4" -n "$0" -- "$@") + ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,auto_quantize_bits:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_checkpoint:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4" -n "$0" -- "$@") eval set -- "$ARGS" while true; do @@ -60,6 +60,7 @@ parse_options() { --lm_eval_tasks ) LM_EVAL_TASKS="$2"; shift 2;; --lm_eval_limit ) LM_EVAL_LIMIT="$2"; shift 2;; --simple_eval_tasks ) SIMPLE_EVAL_TASKS="$2"; shift 2;; + --simple_eval_limit ) SIMPLE_EVAL_LIMIT="$2"; shift 2;; --trust_remote_code ) TRUST_REMOTE_CODE=true; shift;; --use_seq_device_map ) USE_SEQ_DEVICE_MAP=true; shift;; --gpu_max_mem_percentage ) GPU_MAX_MEM_PERCENTAGE="$2"; shift 2;; @@ -159,6 +160,7 @@ parse_options() { echo "lm_eval_tasks: $LM_EVAL_TASKS" echo "lm_eval_limit: $LM_EVAL_LIMIT" echo "simple_eval_tasks: $SIMPLE_EVAL_TASKS" + echo "simple_eval_limit: $SIMPLE_EVAL_LIMIT" echo "num_sample: $NUM_SAMPLES" echo "use_seq_device_map: $USE_SEQ_DEVICE_MAP" echo "gpu_max_mem_percentage: $GPU_MAX_MEM_PERCENTAGE" diff --git a/examples/llm_sparsity/weight_sparsity/data_prep.py b/examples/llm_sparsity/weight_sparsity/data_prep.py index d91caaba8b6..62be755eeca 100644 --- a/examples/llm_sparsity/weight_sparsity/data_prep.py +++ b/examples/llm_sparsity/weight_sparsity/data_prep.py @@ -39,6 +39,13 @@ def preprocess_function(sample): def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--save_path", type=str, default="data") + parser.add_argument( + "--max_samples", + type=int, + default=None, + help="If set, keep only the first N rows of each split before processing. Greatly " + "speeds up preparation for smoke tests (cnn_dailymail train is ~287k rows).", + ) return parser.parse_args() @@ -48,6 +55,14 @@ def main(): # Load dataset from the hub dataset = load_dataset(dataset_id, name=dataset_config) + if args.max_samples is not None: + dataset = type(dataset)( + { + split: ds.select(range(min(args.max_samples, len(ds)))) + for split, ds in dataset.items() + } + ) + # process dataset tokenized_dataset = dataset.map( preprocess_function, batched=True, remove_columns=list(dataset["train"].features) diff --git a/examples/llm_sparsity/weight_sparsity/hf_pts.py b/examples/llm_sparsity/weight_sparsity/hf_pts.py index 77574c1c2c4..f778ca32abf 100644 --- a/examples/llm_sparsity/weight_sparsity/hf_pts.py +++ b/examples/llm_sparsity/weight_sparsity/hf_pts.py @@ -20,38 +20,15 @@ import numpy as np import torch -from datasets import load_dataset -from torch.utils.data import DataLoader from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer import modelopt.torch.opt as mto import modelopt.torch.sparsity as mts +from modelopt.torch.utils import get_dataset_dataloader DEFAULT_PAD_TOKEN = "[PAD]" -def get_calib_dataloader( - data="cnn_dailymail", tokenizer=None, batch_size=1, calib_size=512, block_size=512, device=None -): - print("Loading calibration dataset") - if data == "cnn_dailymail": - dataset = load_dataset("abisee/cnn_dailymail", name="3.0.0", split="train") - dataset = dataset["article"][:calib_size] - else: - raise NotImplementedError - - batch_encoded = tokenizer( - dataset, return_tensors="pt", padding=True, truncation=True, max_length=block_size - ) - if device: - batch_encoded = batch_encoded.to(device) - batch_encoded = batch_encoded["input_ids"] - - calib_dataloader = DataLoader(batch_encoded, batch_size=batch_size, shuffle=False) - - return calib_dataloader - - def smart_tokenizer_and_embedding_resize( special_tokens_dict: dict, tokenizer: PreTrainedTokenizer, @@ -80,7 +57,8 @@ def get_tokenizer(ckpt_path: str, model_max_length: int, trust_remote_code: bool tokenizer = AutoTokenizer.from_pretrained( ckpt_path, model_max_length=model_max_length, - padding_side="right", + # Left padding is recommended for calibration (get_dataset_dataloader warns otherwise). + padding_side="left", use_fast=False, trust_remote_code=trust_remote_code, ) @@ -126,13 +104,13 @@ def main(args): model=model, ) - calib_size = args.calib_size - # Get calibration dataloader - calib_dataloader = get_calib_dataloader( + calib_dataloader = get_dataset_dataloader( + dataset_name=args.dataset, tokenizer=tokenizer, batch_size=args.batch_size, - calib_size=calib_size, + num_samples=args.calib_size, + max_sample_length=args.model_max_length, device=args.device, ) @@ -160,11 +138,18 @@ def main(args): "--model_name_or_path", help="Specify where the PyTorch checkpoint path is", required=True ) parser.add_argument("--device", default="cuda") + parser.add_argument( + "--dataset", + default="cnn_dailymail", + help="Calibration dataset: a ModelOpt-registered name (see get_supported_datasets()), " + "a HuggingFace dataset id, or a local .jsonl path.", + ) parser.add_argument("--dtype", help="Model data type.", default="fp16") parser.add_argument( "--model_max_length", + type=int, default=2048, - help="Maximum sequence length. Sequences will be right padded (and possibly truncated).", + help="Maximum sequence length used for both the tokenizer and calibration sequences.", ) parser.add_argument("--batch_size", help="Batch size for calibration.", type=int, default=1) parser.add_argument( diff --git a/pyproject.toml b/pyproject.toml index 07bc449e176..74927947215 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -305,6 +305,8 @@ skips = [ # print execution time for 50 slowest tests and generate coverage reports addopts = "-v -ra --instafail --cov-report=term-missing --cov-report=html --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers" pythonpath = ["tests/"] +# Apply per-test timeouts (see tests/conftest.py) to the test call only, not fixture setup/teardown +timeout_func_only = true markers = [ "integration: Tests that require external services or other non-hermetic dependencies", "manual: Only run when --run-manual is given", diff --git a/tests/_test_utils/deploy_utils.py b/tests/_test_utils/deploy_utils.py index c711ca060d5..a1a51de8a4b 100644 --- a/tests/_test_utils/deploy_utils.py +++ b/tests/_test_utils/deploy_utils.py @@ -17,6 +17,7 @@ import os import subprocess import sys +import traceback import pytest import torch @@ -84,8 +85,6 @@ def _run_trtllm_deploy( ) deployer._deploy_trtllm_impl() except Exception: - import traceback - traceback.print_exc() pytest.fail(traceback.format_exc()) @@ -111,8 +110,6 @@ def _run_vllm_deploy( ) deployer._deploy_vllm_impl() except Exception: - import traceback - traceback.print_exc() pytest.fail(traceback.format_exc()) @@ -138,8 +135,6 @@ def _run_sglang_deploy( ) deployer._deploy_sglang_impl() except Exception: - import traceback - traceback.print_exc() pytest.fail(traceback.format_exc()) diff --git a/tests/_test_utils/examples/models.py b/tests/_test_utils/examples/models.py index 8bf2b95a60c..dd1d8755f84 100644 --- a/tests/_test_utils/examples/models.py +++ b/tests/_test_utils/examples/models.py @@ -48,38 +48,28 @@ def _select_path(remote_id: str, local_id: str) -> str: local_id="TinyLlama-1.1B-Chat-v1.0", ) -SXDL_PATH = _select_path( - remote_id="stabilityai/stable-diffusion-xl-base-1.0", - local_id="stable-diffusion-xl-base-1.0", -) - -PIXART_PATH = _select_path( - remote_id="PixArt-alpha/PixArt-XL-2-1024-MS", - local_id="PixArt-XL-2-1024-MS", -) - -LLAVA_PATH = _select_path( - remote_id="llava-hf/llava-1.5-7b-hf", - local_id="llava-1.5-7b-hf", -) - QWEN_VL_PATH = _select_path( remote_id="Qwen/Qwen3-VL-2B-Instruct", local_id="Qwen3-VL-2B-Instruct", ) # Diffusers -FLUX_SCHNELL_PATH = _select_path( - remote_id="hf-internal-testing/tiny-flux-pipe", - local_id="black-forest-labs/FLUX.1-schnell", -) - -SDXL_1_0_PATH = _select_path( +SDXL_PATH = _select_path( remote_id="hf-internal-testing/tiny-sdxl-pipe", - local_id="stabilityai/stable-diffusion-xl-base-1.0", + local_id="stable-diffusion-xl-base-1.0", ) SD3_PATH = _select_path( remote_id="hf-internal-testing/tiny-sd3-pipe", - local_id="stabilityai/stable-diffusion-3-medium-diffusers", + local_id="stable-diffusion-3-medium-diffusers", +) + +FLUX_SCHNELL_PATH = _select_path( + remote_id="hf-internal-testing/tiny-flux-pipe", + local_id="FLUX.1-schnell", +) + +PIXART_PATH = _select_path( + remote_id="PixArt-alpha/PixArt-XL-2-1024-MS", + local_id="PixArt-XL-2-1024-MS", ) diff --git a/tests/_test_utils/onnx/lib_test_models.py b/tests/_test_utils/onnx/lib_test_models.py index 97bc22b7121..2c34a350852 100644 --- a/tests/_test_utils/onnx/lib_test_models.py +++ b/tests/_test_utils/onnx/lib_test_models.py @@ -679,11 +679,13 @@ def build_conv_batchnorm_sig_mul_model(): def build_conv_act_pool_model(include_reshape_node=False): - # Define your model inputs and outputs + # Define your model inputs and outputs. Kept tiny (this is a QDQ-placement + # structure test, independent of channel/spatial sizes) so calibration over the + # ONNX runtime stays fast. input_names = ["input_0"] output_names = ["output_0"] - input_shapes = [(32, 64, 256, 256)] - output_shapes = [(32, 128, 128, 128)] + input_shapes = [(1, 8, 32, 32)] + output_shapes = [(1, 16, 16, 16)] inputs = [ helper.make_tensor_value_info(input_name, onnx.TensorProto.FLOAT, input_shape) @@ -762,44 +764,44 @@ def build_conv_act_pool_model(include_reshape_node=False): helper.make_tensor( name="weights_1", data_type=onnx.TensorProto.FLOAT, - dims=(128, 64, 3, 3), - vals=np.random.uniform(low=0.5, high=1.0, size=128 * 64 * 3 * 3), + dims=(16, 8, 3, 3), + vals=np.random.uniform(low=0.5, high=1.0, size=16 * 8 * 3 * 3), ), helper.make_tensor( name="bias_1", data_type=onnx.TensorProto.FLOAT, - dims=(128,), - vals=np.random.uniform(low=0.5, high=1.0, size=128), + dims=(16,), + vals=np.random.uniform(low=0.5, high=1.0, size=16), ), helper.make_tensor( name="bn1_scale", data_type=onnx.TensorProto.FLOAT, - dims=(128,), - vals=np.random.uniform(low=0.5, high=1.0, size=128), + dims=(16,), + vals=np.random.uniform(low=0.5, high=1.0, size=16), ), helper.make_tensor( name="bn1_bias", data_type=onnx.TensorProto.FLOAT, - dims=(128,), - vals=np.random.uniform(low=0.5, high=1.0, size=128), + dims=(16,), + vals=np.random.uniform(low=0.5, high=1.0, size=16), ), helper.make_tensor( name="bn1_mean", data_type=onnx.TensorProto.FLOAT, - dims=(128,), - vals=np.random.uniform(low=0.5, high=1.0, size=128), + dims=(16,), + vals=np.random.uniform(low=0.5, high=1.0, size=16), ), helper.make_tensor( name="bn1_var", data_type=onnx.TensorProto.FLOAT, - dims=(128,), - vals=np.random.uniform(low=0.5, high=1.0, size=128), + dims=(16,), + vals=np.random.uniform(low=0.5, high=1.0, size=16), ), helper.make_tensor( name="weights_2", data_type=onnx.TensorProto.FLOAT, - dims=(128, 128, 3, 3), - vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3), + dims=(16, 16, 3, 3), + vals=np.random.uniform(low=0.5, high=1.0, size=16 * 16 * 3 * 3), ), ] if include_reshape_node: @@ -808,7 +810,7 @@ def build_conv_act_pool_model(include_reshape_node=False): name="shape_1", data_type=onnx.TensorProto.INT64, dims=(4,), - vals=(32, 128, 256, 256), + vals=(1, 16, 32, 32), ), ) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index dec0413883c..23e67821371 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -55,11 +55,10 @@ def init_process(rank, size, job=None, backend="gloo", port=None): def spawn_multiprocess_job(size, job, backend="gloo"): port = get_free_port() - + ctx = mp.get_context("spawn") processes = [] - mp.set_start_method("spawn", force=True) for rank in range(size): - p = mp.Process(target=init_process, args=(rank, size, job, backend, port)) + p = ctx.Process(target=init_process, args=(rank, size, job, backend, port)) p.start() processes.append(p) diff --git a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py index 8559192718f..c9de64e6c60 100644 --- a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py +++ b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py @@ -15,13 +15,18 @@ import pytest import torch +import torch.nn as nn import torch.nn.functional as F from _test_utils.torch.quantization.quant_utils import quant from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization import utils as quant_utils from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.model_calib import max_calibrate +from modelopt.torch.quantization.model_calib import ( + enable_stats_collection, + finish_stats_collection, + max_calibrate, +) from modelopt.torch.quantization.nn import QuantLinear, SequentialQuantizer, TensorQuantizer @@ -148,6 +153,10 @@ def test_entropy_and_percentile_calib(self): """Don't really have a good way to test it.""" quant_attr_cfg1 = QuantizerAttributeConfig(calibrator="histogram") quantizer1 = TensorQuantizer(quant_attr_cfg1, if_calib=True, if_quant=False).to(self.device) + # Reduce histogram bins (default 2048) before the first collect to keep the + # entropy KL-divergence search ~8x cheaper; the assertion compares two equal + # computes so the result stays valid. + quantizer1._calibrator._num_bins = 512 x_1 = torch.rand(3, 6, 7, 7).to(self.device) x_2 = torch.rand(3, 6, 7, 7).to(self.device) @@ -230,13 +239,6 @@ def test_use_constant_amax(self): def test_use_constant_amax_skips_calibration(self): """Test that use_constant_amax quantizers are disabled during calibration and re-enabled after.""" - import torch.nn as nn - - from modelopt.torch.quantization.model_calib import ( - enable_stats_collection, - finish_stats_collection, - ) - # Build a small model with one use_constant_amax quantizer and one normal quantizer model = nn.ModuleDict( { diff --git a/tests/_test_utils/torch/tokenizer/chat_template.jinja b/tests/_test_utils/torch/tokenizer/chat_template.jinja new file mode 100644 index 00000000000..faecb702a2c --- /dev/null +++ b/tests/_test_utils/torch/tokenizer/chat_template.jinja @@ -0,0 +1,21 @@ +{# Terse generation-tagged chat template for the tiny test tokenizer. + + Uses plain-text "Q:/A:" role prefixes rather than ChatML special tokens + (<|im_start|> etc.): the 128-vocab tiny tokenizer has no such tokens, so it + byte-encodes them into ~70 extra tokens, which blows past tight max-seq-len + limits in example tests (e.g. compute_hidden_states with max-seq-len=32 on a + 32-position tiny model). The {% generation %} blocks still let unit tests + exercise answer-only-loss assistant masking via return_assistant_tokens_mask. #} +{{- bos_token -}} +{%- for message in messages -%} + {%- if message['role'] == 'user' -%} + {{- 'Q: ' + message['content'] + '\n' -}} + {%- elif message['role'] == 'assistant' -%} + {{- 'A: ' -}} + {%- generation -%} + {{- message['content'] -}} + {%- endgeneration -%} + {{- '\n' -}} + {%- endif -%} +{%- endfor -%} +{{- eos_token -}} diff --git a/tests/_test_utils/torch/tokenizer/tokenizer_config.json b/tests/_test_utils/torch/tokenizer/tokenizer_config.json index bdd427826a5..7b6301834be 100644 --- a/tests/_test_utils/torch/tokenizer/tokenizer_config.json +++ b/tests/_test_utils/torch/tokenizer/tokenizer_config.json @@ -1,6 +1,5 @@ { "bos_token": "<|begin_of_text|>", - "chat_template": "{{ bos_token }}{% for message in messages %}{% if message['role'] == 'user' %}Q: {{ message['content'] }}{% elif message['role'] == 'assistant' %}A: {{ message['content'] }}{% endif %}{{ eos_token }}{% endfor %}", "clean_up_tokenization_spaces": true, "eos_token": "<|eot_id|>", "pad_token": "<|eot_id|>", diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index 8fc88964662..3aad8771c2b 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -45,8 +45,16 @@ TINY_TOKENIZER_PATH = Path(__file__).parent / "tokenizer" -def get_tiny_tokenizer() -> "transformers.PreTrainedTokenizerBase": - return AutoTokenizer.from_pretrained(TINY_TOKENIZER_PATH) +def get_tiny_tokenizer(*, pad_side: str = "left") -> "transformers.PreTrainedTokenizerBase": + """Returns a tiny tokenizer for tests. + + Default to left padding, which is what decoder-LM calibration/generation expects and what + ``get_dataset_dataloader`` warns about otherwise. Callers needing right padding can override + with ``pad_side="right"``. + """ + tokenizer = AutoTokenizer.from_pretrained(TINY_TOKENIZER_PATH) + tokenizer.padding_side = pad_side + return tokenizer ##### Qwen3 ##### diff --git a/tests/conftest.py b/tests/conftest.py index a4e65ff2ae3..16a9f3f2614 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,8 +44,22 @@ def pytest_addoption(parser): ) +# Default per-test `call` wall-clock cap (seconds) by top-level tests/ subdirectory +# Every collectible test group must be listed here else collection errors occur +# A test can override its cap by adding ``@pytest.mark.timeout(...)`` +_DEFAULT_TIMEOUT = { + "examples": 300, + "gpu": 120, + "gpu_megatron": 120, + "gpu_trtllm": 60, + "gpu_vllm": 60, + "regression": 180, + "unit": 60, +} + + def pytest_collection_modifyitems(config, items): - """Skip tests with specific markers unless their corresponding flag is provided.""" + """Skip flag-gated tests and apply a default per-test timeout based on the test directory.""" skip_marks = [ ("manual", "--run-manual"), ("release", "--run-release"), @@ -58,6 +72,33 @@ def pytest_collection_modifyitems(config, items): if mark_name in item.keywords: item.add_marker(skipper) + tests_root = Path(__file__).parent + for item in items: + if item.get_closest_marker("timeout") is not None or not item.path.is_relative_to( + tests_root + ): + continue + # First path component under tests/ is the group dir (unit, gpu, examples, ...). + # Crash loudly (rather than silently skip) if a group has no configured default, so a + # newly added tests// must be given an explicit timeout in the mapping above. + group = item.path.relative_to(tests_root).parts[0] + if group not in _DEFAULT_TIMEOUT: + raise pytest.UsageError( + f"tests/{group}/ has no default timeout; add '{group}' to " + "_DEFAULT_TIMEOUT in tests/conftest.py." + ) + item.add_marker(pytest.mark.timeout(_DEFAULT_TIMEOUT[group])) + + +@pytest.fixture +def tiny_tokenizer(): + """Real tiny HF tokenizer (vocab=128) shared across unit and gpu test lanes.""" + # Lazy import: transformers_models.py runs ``pytest.importorskip("transformers")`` + # at module load, which we don't want to trigger at conftest import time. + from _test_utils.torch.transformers_models import get_tiny_tokenizer + + return get_tiny_tokenizer() + @pytest.fixture def skip_on_windows(): diff --git a/tests/examples/diffusers/test_cache_diffusion.py b/tests/examples/diffusers/test_cache_diffusion.py index 18f49525781..7129f3163f2 100644 --- a/tests/examples/diffusers/test_cache_diffusion.py +++ b/tests/examples/diffusers/test_cache_diffusion.py @@ -18,7 +18,7 @@ import pytest import torch -from _test_utils.examples.models import PIXART_PATH, SXDL_PATH +from _test_utils.examples.models import PIXART_PATH, SDXL_PATH from _test_utils.examples.run_command import MODELOPT_ROOT from diffusers import DiffusionPipeline, PixArtAlphaPipeline @@ -29,9 +29,8 @@ def test_sdxl_cachify(): pipe = DiffusionPipeline.from_pretrained( - SXDL_PATH, + SDXL_PATH, torch_dtype=torch.float16, - variant="fp16", use_safetensors=True, ).to("cuda") cachify.prepare(pipe, SDXL_DEFAULT_CONFIG) @@ -39,7 +38,8 @@ def test_sdxl_cachify(): prompt = "A random person with a head that is made of flowers, photo by James C. Leyendecker, \ Afrofuturism, studio portrait, dynamic pose, national geographic photo, retrofuturism, biomorphicy" generator = torch.Generator(device="cuda").manual_seed(2946901) - pipe(prompt=prompt, generator=generator, num_inference_steps=30).images[0] + # 8 steps still exercises the step-modulo cache pattern; this is a runs-without-error smoke test. + pipe(prompt=prompt, generator=generator, num_inference_steps=8).images[0] # Clear cuda memory as pytest doesnt clear it between tests del pipe torch.cuda.empty_cache() @@ -55,7 +55,8 @@ def test_pixart_cachify(): prompt = "a small cactus with a happy face in the Sahara desert" generator = torch.Generator(device="cuda").manual_seed(2946901) - pipe(prompt=prompt, generator=generator, num_inference_steps=30).images[0] + # 8 steps still exercises the step-modulo cache pattern; this is a runs-without-error smoke test. + pipe(prompt=prompt, generator=generator, num_inference_steps=8).images[0] # Clear cuda memory as pytest doesnt clear it between tests del pipe torch.cuda.empty_cache() diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index 5b117b41b3f..15c5eb44934 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -17,7 +17,7 @@ from typing import NamedTuple import pytest -from _test_utils.examples.models import FLUX_SCHNELL_PATH, SD3_PATH, SDXL_1_0_PATH +from _test_utils.examples.models import FLUX_SCHNELL_PATH, SD3_PATH, SDXL_PATH from _test_utils.examples.run_command import run_example_command from _test_utils.torch.misc import minimum_sm @@ -45,13 +45,13 @@ def _run_cmd(self, script: str, *args: str) -> None: def _format_args(self) -> list[str]: return [ "--calib-size", - "8", + "4", "--percentile", "1.0", "--alpha", "0.8", "--n-steps", - "20", + "2", "--batch-size", "2", "--format", @@ -93,6 +93,8 @@ def inference(self, tmp_path: Path) -> None: str(tmp_path / f"{self.name}_{self.format_type}_onnx/model.onnx"), "--dq-only", "--torch-autocast", + "--num-inference-steps", + "2", ) @@ -118,7 +120,7 @@ def inference(self, tmp_path: Path) -> None: pytest.param( DiffuserModel( name="sdxl-1.0", - path=SDXL_1_0_PATH, + path=SDXL_PATH, dtype="Half", format_type="fp8", quant_algo="max", @@ -128,7 +130,7 @@ def inference(self, tmp_path: Path) -> None: ), DiffuserModel( name="sdxl-1.0", - path=SDXL_1_0_PATH, + path=SDXL_PATH, dtype="Half", format_type="int8", quant_algo="smoothquant", @@ -271,8 +273,8 @@ def test_wan22_quantization(wan_model: Wan22Model, tiny_wan22_path: str, tmp_pat ("flux-schnell", FLUX_SCHNELL_PATH, True), ("sd3-medium", SD3_PATH, False), ("sd3-medium", SD3_PATH, True), - ("sdxl-1.0", SDXL_1_0_PATH, False), - ("sdxl-1.0", SDXL_1_0_PATH, True), + ("sdxl-1.0", SDXL_PATH, False), + ("sdxl-1.0", SDXL_PATH, True), ], ids=[ "flux_schnell_torch", @@ -296,6 +298,8 @@ def test_diffusion_trt_torch( "--override-model-path", model_path, "--torch", + "--num-inference-steps", + "2", ] if torch_compile: cmd_args.append("--torch-compile") diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index a5c81d36937..88821bbf8f7 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -17,7 +17,7 @@ from typing import NamedTuple import pytest -from _test_utils.examples.models import FLUX_SCHNELL_PATH, SDXL_1_0_PATH +from _test_utils.examples.models import FLUX_SCHNELL_PATH, SDXL_PATH from _test_utils.examples.run_command import run_example_command from _test_utils.torch.misc import minimum_sm @@ -41,11 +41,11 @@ def quantize_and_export_hf(self, tmp_path: Path) -> Path: "--override-model-path", self.path, "--calib-size", - "8", + "4", "--batch-size", "2", "--n-steps", - "20", + "2", "--percentile", "1.0", "--alpha", @@ -72,7 +72,7 @@ def quantize_and_export_hf(self, tmp_path: Path) -> Path: [ DiffuserHfExportModel( name="sdxl-1.0", - path=SDXL_1_0_PATH, + path=SDXL_PATH, dtype="Half", format_type="int8", quant_algo="smoothquant", @@ -90,7 +90,7 @@ def quantize_and_export_hf(self, tmp_path: Path) -> Path: pytest.param( DiffuserHfExportModel( name="sdxl-1.0", - path=SDXL_1_0_PATH, + path=SDXL_PATH, dtype="Half", format_type="fp8", quant_algo="max", diff --git a/tests/examples/gpt_oss/test_gpt_oss_qat.py b/tests/examples/gpt-oss/test_gpt_oss_qat.py similarity index 96% rename from tests/examples/gpt_oss/test_gpt_oss_qat.py rename to tests/examples/gpt-oss/test_gpt_oss_qat.py index f36dd9e023e..f358a450f17 100644 --- a/tests/examples/gpt_oss/test_gpt_oss_qat.py +++ b/tests/examples/gpt-oss/test_gpt_oss_qat.py @@ -19,6 +19,7 @@ import pytest from _test_utils.examples.run_command import run_example_command from _test_utils.torch.misc import minimum_gpu +from _test_utils.torch.transformers_models import create_tiny_gpt_oss_dir from datasets import Dataset, DatasetDict @@ -116,24 +117,9 @@ def gpt_oss_qat_training(self, tmp_path, sft_dir=None): if sft_dir is None: sft_dir = tmp_path / f"{model_name}-sft" - # If SFT directory doesn't exist, create a mock one for standalone testing + # If SFT directory doesn't exist, create a tiny stand-in model for standalone testing if not sft_dir.exists(): - sft_dir.mkdir() - - # Create minimal config.json for the mock model - config_content = { - "model_type": "gpt_oss", - "hidden_size": 5120, - "num_attention_heads": 40, - "num_hidden_layers": 44, - "vocab_size": 100000, - "torch_dtype": "bfloat16", - } - - import json - - with open(sft_dir / "config.json", "w") as f: - json.dump(config_content, f) + sft_dir = create_tiny_gpt_oss_dir(tmp_path, with_tokenizer=True) qat_output_dir = tmp_path / f"{model_name}-qat" diff --git a/tests/examples/llm_eval/test_llm_eval.py b/tests/examples/llm_eval/test_llm_eval.py index 356430ea6f6..465ba06f4a5 100644 --- a/tests/examples/llm_eval/test_llm_eval.py +++ b/tests/examples/llm_eval/test_llm_eval.py @@ -15,6 +15,7 @@ import subprocess +import pytest from _test_utils.examples.run_command import ( extend_cmd_parts, run_example_command, @@ -40,6 +41,7 @@ def test_lm_eval_hf(tmp_path): @minimum_sm(89) +@pytest.mark.timeout(480) def test_qwen3_eval_fp8(tmp_path): # Bump max_position_embeddings: TRT-LLM serve rejects prompts longer than # max_seq_len, and the default (32) is shorter than even simple MMLU prompts. @@ -52,7 +54,9 @@ def test_qwen3_eval_fp8(tmp_path): calib=64, lm_eval_tasks="hellaswag,gsm8k", simple_eval_tasks="humaneval", - lm_eval_limit=0.1, + lm_eval_limit=16, + simple_eval_limit=16, + output=128, # Cap generation length: gsm8k/humaneval otherwise generate up to 1024 tokens/sample batch=8, ) finally: diff --git a/tests/examples/llm_ptq/_extensions/test_torch_extensions.py b/tests/examples/llm_ptq/_extensions/test_torch_extensions.py new file mode 120000 index 00000000000..9267cc06186 --- /dev/null +++ b/tests/examples/llm_ptq/_extensions/test_torch_extensions.py @@ -0,0 +1 @@ +../../../gpu/_extensions/test_torch_extensions.py \ No newline at end of file diff --git a/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py b/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py index 7094b29894d..6277b0fa5b4 100644 --- a/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py +++ b/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py @@ -58,7 +58,8 @@ def run_llm_sparsity_ft_command( def data_path(tmp_path_factory): data_path = tmp_path_factory.mktemp("data") run_example_command( - ["python", "data_prep.py", "--save_path", data_path], "llm_sparsity/weight_sparsity" + ["python", "data_prep.py", "--save_path", data_path, "--max_samples", "64"], + "llm_sparsity/weight_sparsity", ) # Copy eval data to train path for faster test @@ -82,7 +83,7 @@ def test_llama_sparsity(tiny_llama_path, tmp_path, sparsity_fmt, dtype): sparsity_fmt=sparsity_fmt, dtype=dtype, calib_size=8, - model_max_length=128, + model_max_length=64, ) @@ -97,7 +98,7 @@ def _test_llama_sparsity_finetune(tiny_llama_path, tmp_path, data_path, sparsity sparsity_fmt=sparsity_fmt, dtype=dtype, calib_size=8, - model_max_length=128, + model_max_length=64, ) # Then do finetuning using the sparsified model @@ -108,7 +109,7 @@ def _test_llama_sparsity_finetune(tiny_llama_path, tmp_path, data_path, sparsity output_dir=finetune_output, data_path=data_path, num_epochs=0.001, - max_length=128, + max_length=64, ) diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 9f84f50c287..fa8233de7a0 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -14,7 +14,6 @@ # limitations under the License. """Tests for prune_minitron.py and distill.py scripts.""" -import subprocess from pathlib import Path from _test_utils.examples.run_command import extend_cmd_parts, run_example_command @@ -26,8 +25,9 @@ def test_distill_and_convert(tmp_path: Path, num_gpus): teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) - train_iters = 5 + train_iters = 2 distill_output_dir = tmp_path / "distill_output" + distilled_hf_path = tmp_path / "distilled_hf" distill_cmd_parts = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"], student_hf_path=teacher_hf_path, @@ -35,7 +35,7 @@ def test_distill_and_convert(tmp_path: Path, num_gpus): output_dir=distill_output_dir, tp_size=num_gpus, pp_size=1, - seq_length=32, + seq_length=16, mbs=1, gbs=4, train_iters=train_iters, @@ -43,28 +43,12 @@ def test_distill_and_convert(tmp_path: Path, num_gpus): eval_interval=5, eval_iters=1, log_interval=1, + hf_export_path=distilled_hf_path, + student_hf_model=teacher_hf_path, ) run_example_command(distill_cmd_parts, example_path="megatron_bridge") - megatron_ckpt_path = distill_output_dir / f"checkpoints/iter_{train_iters:07d}" - assert megatron_ckpt_path.exists() - - # Convert distilled Megatron checkpoint back to HF format - distilled_hf_path = tmp_path / "distilled_hf" - subprocess.run( - [ - "python", - "/opt/Megatron-Bridge/examples/conversion/convert_checkpoints.py", - "export", - "--hf-model", - str(teacher_hf_path), - "--megatron-path", - str(megatron_ckpt_path), - "--hf-path", - str(distilled_hf_path), - ], - check=True, - ) + assert (distill_output_dir / f"checkpoints/iter_{train_iters:07d}").exists() assert (distilled_hf_path / "config.json").exists() @@ -89,7 +73,7 @@ def test_distill_puzzletron_anymodel(tmp_path: Path, num_gpus): output_dir=output_dir, tp_size=num_gpus, pp_size=1, - seq_length=32, + seq_length=16, mbs=1, gbs=4, train_iters=train_iters, diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 0dbdef1c401..33588234dec 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -44,8 +44,8 @@ def test_prune_minitron(tmp_path: Path, num_gpus, create_tiny_model_dir): output_hf_path=pruned_model_path, pp_size=num_gpus, calib_dataset_name="cnn_dailymail", - calib_num_samples=16, - seq_length=32, + calib_num_samples=8, + seq_length=16, prune_target_params=prune_target_params, prune_score_func="mmlu_1pct_bs32", ss_channel_divisor=4, diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index f22a5e73b17..9aee3ac77aa 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -44,9 +44,9 @@ def test_quantize_and_export(tmp_path: Path, num_gpus): recipe="general/ptq/fp8_default-kv_fp8", tp_size=num_gpus, calib_dataset_name="cnn_dailymail", - calib_num_samples=16, + calib_num_samples=4, calib_batch_size=1, - seq_length=32, + seq_length=16, export_megatron_path=megatron_path, ) run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True) diff --git a/tests/examples/speculative_decoding/test_eagle.py b/tests/examples/speculative_decoding/test_eagle.py index ebbbe237d17..00f0adc33c3 100644 --- a/tests/examples/speculative_decoding/test_eagle.py +++ b/tests/examples/speculative_decoding/test_eagle.py @@ -135,6 +135,7 @@ def test_llama_eagle3(tiny_llama_path, f"model.model_name_or_path={tiny_llama_path}", f"data.data_path={tiny_daring_anteater_path}", f"training.output_dir={output_dir}", + "data.sample_size=64", "training.num_train_epochs=0.25", "training.learning_rate=1e-5", "training.training_seq_len=128", @@ -152,14 +153,18 @@ def test_llama_eagle3(tiny_llama_path, def test_resume_training(tiny_daring_anteater_path, eagle_output_dir): """Test resume training of Eagle3.""" - checkpoint_dir = str(eagle_output_dir / "eagle-tinyllama-cp1-mixFalse") + # Resume the compile-free checkpoint (mixTrue) so this test doesn't pay torch.compile; + # the compiled path is smoke-tested once by test_llama_eagle3[1-False]. + checkpoint_dir = str(eagle_output_dir / "eagle-tinyllama-cp1-mixTrue") overrides = [ f"model.model_name_or_path={checkpoint_dir}", f"data.data_path={tiny_daring_anteater_path}", f"training.output_dir={checkpoint_dir}", + "data.sample_size=64", "training.num_train_epochs=0.5", "training.learning_rate=1e-5", "training.training_seq_len=128", + "eagle.eagle_use_torch_compile=false", ] run_example_command( ["./launch_train.sh", "--config", EAGLE3_YAML, *overrides], @@ -172,7 +177,8 @@ def test_ar_validate(eagle_output_dir): run_example_command( [ "python", "./scripts/ar_validate.py", - "--model_path", eagle_output_dir / "eagle-tinyllama-cp1-mixFalse", + # Validate the compile-free checkpoint so AR generation doesn't trigger torch.compile. + "--model_path", eagle_output_dir / "eagle-tinyllama-cp1-mixTrue", "--osl", "10", "--num_samples", "5", "--steps", "3" @@ -186,7 +192,7 @@ def test_export_hf_checkpoint(eagle_output_dir): run_example_command( [ "python", "./scripts/export_hf_checkpoint.py", - "--model_path", eagle_output_dir / "eagle-tinyllama-cp1-mixFalse", + "--model_path", eagle_output_dir / "eagle-tinyllama-cp1-mixTrue", "--export_path", eagle_output_dir / "eagle-tinyllama-export", ], "speculative_decoding", @@ -208,27 +214,34 @@ def test_convert_to_vllm_ckpt(tiny_llama_path, eagle_output_dir): ], "speculative_decoding", ) +# fmt: on @pytest.mark.parametrize( ("model_source", "use_fake_base"), [ - (None, False), # tiny_llama (from fixture), no FakeBase - ("moonshotai/Kimi-K2.5", True), # remote HF repo, FakeBaseModel - pytest.param( - "moonshotai/Kimi-K2-Thinking", True, # remote HF repo, no FakeBaseModel + (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", + True, marks=pytest.mark.manual(reason="skip redundand test, too slow"), ), pytest.param( - "MiniMaxAI/MiniMax-M2.5", True, + "MiniMaxAI/MiniMax-M2.5", + True, marks=pytest.mark.manual(reason="skip redundand test, too slow"), ), ], - ids=["tinyllama", "kimi-k2.5","kimi-k2-thinking","minimax-m2.5"], + ids=["tinyllama", "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, + 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 @@ -256,6 +269,9 @@ def test_offline_eagle3_training( "training.num_train_epochs=0.1", "training.learning_rate=1e-5", "training.training_seq_len=64", + # torch.compile is smoke-tested once by test_llama_eagle3[1-False]; skip its ~2min + # warmup here (the recipe default is true). + "eagle.eagle_use_torch_compile=false", *_TINY_EAGLE_ARCH, ] run_example_command( @@ -294,9 +310,12 @@ def test_eagle3_dry_run(tiny_llama_path, tmp_path, eagle_output_dir): # The dry-run checkpoint must be exportable to the deployment format. run_example_command( [ - "python", "./scripts/export_hf_checkpoint.py", - "--model_path", output_subdir, - "--export_path", export_subdir, + "python", + "./scripts/export_hf_checkpoint.py", + "--model_path", + output_subdir, + "--export_path", + export_subdir, ], "speculative_decoding", ) @@ -333,6 +352,7 @@ def test_offline_resume_training_kimi(tiny_daring_anteater_path, tmp_path, eagle "training.num_train_epochs=0.2", "training.learning_rate=1e-5", "training.training_seq_len=64", + "eagle.eagle_use_torch_compile=false", ] run_example_command( ["./launch_train.sh", "--config", EAGLE3_YAML, *overrides], diff --git a/tests/examples/speculative_decoding/test_eagle_offline_ptq.py b/tests/examples/speculative_decoding/test_eagle_offline_ptq.py index 034a48189c2..33a1ad67fc0 100644 --- a/tests/examples/speculative_decoding/test_eagle_offline_ptq.py +++ b/tests/examples/speculative_decoding/test_eagle_offline_ptq.py @@ -30,6 +30,8 @@ import torch from _test_utils.examples.run_command import MODELOPT_ROOT, run_example_command +from modelopt.torch.export.plugins.hf_spec_export import LLAMA_EAGLE_SINGLE_LAYER + EAGLE3_YAML = str( MODELOPT_ROOT / "modelopt_recipes" / "general" / "speculative_decoding" / "eagle3.yaml" ) @@ -97,6 +99,8 @@ def test_offline_eagle_training(tiny_llama_path, tiny_daring_anteater_path, offl "training.learning_rate=1e-5", "training.training_seq_len=64", "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", *_TINY_EAGLE_ARCH, ] @@ -136,8 +140,6 @@ def test_offline_ptq(offline_ptq_dirs): assert (export_dir / "model.safetensors").exists(), "PTQ export missing model.safetensors" assert (export_dir / "config.json").exists(), "PTQ export missing config.json" - from modelopt.torch.export.plugins.hf_spec_export import LLAMA_EAGLE_SINGLE_LAYER - state_dict = safetensors.torch.load_file(export_dir / "model.safetensors") for key in LLAMA_EAGLE_SINGLE_LAYER["required"] - {"fc", "layers.0.hidden_norm"}: assert f"{key}.weight" in state_dict, f"Missing key '{key}.weight' in exported state dict" diff --git a/tests/examples/speculative_decoding/test_eagle_streaming.py b/tests/examples/speculative_decoding/test_eagle_streaming.py index 291aa0f7929..10836a961ee 100644 --- a/tests/examples/speculative_decoding/test_eagle_streaming.py +++ b/tests/examples/speculative_decoding/test_eagle_streaming.py @@ -125,6 +125,8 @@ def test_streaming_eagle_training( "training.training_seq_len=32", "training.save_steps=1", "training.dataloader_num_workers=0", # enforced by StreamingDataset + # torch.compile is smoke-tested once by test_llama_eagle3[1-False]; skip its warmup here. + "eagle.eagle_use_torch_compile=false", *_TINY_EAGLE_ARCH, ] diff --git a/tests/examples/vlm_ptq/_extensions/test_torch_extensions.py b/tests/examples/vlm_ptq/_extensions/test_torch_extensions.py new file mode 120000 index 00000000000..9267cc06186 --- /dev/null +++ b/tests/examples/vlm_ptq/_extensions/test_torch_extensions.py @@ -0,0 +1 @@ +../../../gpu/_extensions/test_torch_extensions.py \ No newline at end of file diff --git a/tests/gpu/_extensions/test_torch_extensions.py b/tests/gpu/_extensions/test_torch_extensions.py index b44ec369f72..10a26a2280a 100644 --- a/tests/gpu/_extensions/test_torch_extensions.py +++ b/tests/gpu/_extensions/test_torch_extensions.py @@ -14,8 +14,13 @@ # limitations under the License. +import pytest + import modelopt.torch.quantization.extensions as ext +# Override default timeout as these tests JIT-compile the CUDA extensions, which is slow +pytestmark = pytest.mark.timeout(180) + # Compile extensions first so it does not count towards time used to run a test that needs it def test_cuda_ext(): diff --git a/tests/gpu/_extensions/test_torch_kernels.py b/tests/gpu/_extensions/test_torch_kernels.py new file mode 100644 index 00000000000..22c6c584a19 --- /dev/null +++ b/tests/gpu/_extensions/test_torch_kernels.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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. + +"""Pre-compile ModelOpt torch CUDA kernels so the one-time JIT build cost is paid here +rather than landing on the first functional test that uses them (e.g. the conv3d +implicit-GEMM tests). ``tests/gpu/_extensions`` is collected before ``tests/gpu/torch``, so +the module-level kernel cache is warm by the time those tests run in the same process. +""" + +import pytest + +# Override default timeout as these tests JIT-compile the CUDA extensions, which is slow +pytestmark = pytest.mark.timeout(180) + + +def test_conv3d_implicit_gemm(): + """Compile the conv3d implicit-GEMM CUDA extension.""" + from modelopt.torch.kernels.quantization.conv.implicit_gemm_cuda import _get_cuda_module + + assert _get_cuda_module() is not None diff --git a/tests/gpu/onnx/quantization/autotune/test_workflow.py b/tests/gpu/onnx/quantization/autotune/test_workflow.py index 8066766a9cf..eebc42f6c33 100644 --- a/tests/gpu/onnx/quantization/autotune/test_workflow.py +++ b/tests/gpu/onnx/quantization/autotune/test_workflow.py @@ -35,7 +35,13 @@ def simple_conv_model(): return _test_models._create_simple_conv_onnx_model() -@pytest.mark.parametrize("use_trtexec", [True, False]) +@pytest.mark.parametrize( + "use_trtexec", + [ + pytest.param(True, marks=pytest.mark.timeout(240)), # trtexec build needs longer + False, + ], +) def test_export_quantized_model(use_trtexec, simple_conv_model): """Test exporting quantized model with Q/DQ.""" if use_trtexec: diff --git a/tests/gpu/onnx/quantization/test_quantize_onnx_torch_int4_awq.py b/tests/gpu/onnx/quantization/test_quantize_onnx_torch_int4_awq.py index 73302df22ac..57102348d81 100644 --- a/tests/gpu/onnx/quantization/test_quantize_onnx_torch_int4_awq.py +++ b/tests/gpu/onnx/quantization/test_quantize_onnx_torch_int4_awq.py @@ -15,10 +15,12 @@ # NOTE: This test requires modelopt.torch.quantization to be installed as well. +import builtins import copy import os from functools import partial +import numpy import torch from _test_utils.import_helper import skip_if_no_libcudnn from _test_utils.onnx.lib_test_models import SimpleMLP, export_as_onnx, find_init @@ -41,10 +43,6 @@ def test_safe_cupy_array(monkeypatch): """Comprehensive test for safe_cupy_array covering all code paths.""" - import builtins - - import numpy # Import actual numpy for creating int4 tensors - # Test 1: Regular numpy array (should hit line 122) result = int4.safe_cupy_array(numpy.array([1, 2, 3, 4], dtype=numpy.float32)) assert isinstance(result, np.ndarray) diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 92264b894cc..1897f3d6db3 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -32,8 +32,6 @@ def _update_weight_test(rank, size): """Test fsdp2 weight update context for weight update -> only value changed""" - from torch.distributed._composable.fsdp import fully_shard - with patch_fsdp_mp_dtypes(): # Define and shard model model = ToyModel(dims=[4, 4], bias=False).to("cuda") 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 6e0c56bfd1d..f3bfa212f27 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 @@ -97,6 +97,9 @@ def test_unified_hf_export_and_check_safetensors( qformat=qformat, export_path=output_dir, dataset="cnn_dailymail", + # 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, ) # Run the command diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa.py index 7fc3a554c7a..53d98ae0752 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa.py @@ -19,12 +19,7 @@ import torch import torch.nn.functional as F from conftest import make_qkv, make_varlen_meta, sdpa_reference - -pytestmark = [ - pytest.mark.filterwarnings("ignore::UserWarning"), - pytest.mark.filterwarnings("ignore::RuntimeWarning"), - pytest.mark.filterwarnings("ignore::DeprecationWarning"), -] +from transformers import AutoModelForCausalLM, AutoTokenizer from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE @@ -247,9 +242,6 @@ class TestHFIntegration: def test_triton_matches_eager(self, tiny_llama_dir): """Triton attention produces same logits and generated tokens as eager.""" - pytest.importorskip("transformers") - from transformers import AutoModelForCausalLM, AutoTokenizer - tok = AutoTokenizer.from_pretrained(tiny_llama_dir) if tok.pad_token_id is None: tok.pad_token_id = tok.eos_token_id @@ -295,9 +287,6 @@ def test_triton_matches_eager(self, tiny_llama_dir): def test_triton_padded_batch(self, tiny_llama_dir): """Padded batch produces valid logits.""" - pytest.importorskip("transformers") - from transformers import AutoModelForCausalLM, AutoTokenizer - model = AutoModelForCausalLM.from_pretrained( tiny_llama_dir, attn_implementation="modelopt_triton", diff --git a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py index c06d6981959..0a51e48a1c7 100644 --- a/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py +++ b/tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.py @@ -19,12 +19,6 @@ import torch from conftest import make_qkv, make_varlen_meta -pytestmark = [ - pytest.mark.filterwarnings("ignore::UserWarning"), - pytest.mark.filterwarnings("ignore::RuntimeWarning"), - pytest.mark.filterwarnings("ignore::DeprecationWarning"), -] - from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: diff --git a/tests/gpu/torch/kernels/conftest.py b/tests/gpu/torch/kernels/conftest.py index fa4f6177143..77923c93a13 100644 --- a/tests/gpu/torch/kernels/conftest.py +++ b/tests/gpu/torch/kernels/conftest.py @@ -15,10 +15,32 @@ """Shared fixtures and helpers for Triton flash attention tests.""" +from pathlib import Path + import pytest import torch import torch.nn.functional as F +_KERNELS_DIR = Path(__file__).parent + + +def pytest_collection_modifyitems(items): + """Silence noisy third-party warnings (triton/torch) for all kernel tests. + + Consolidated here so individual kernel test modules don't each repeat the same + ``pytest.mark.filterwarnings`` block. Scoped to this directory only — the rest of + the suite keeps surfacing warnings. + """ + ignore_marks = [ + pytest.mark.filterwarnings("ignore::UserWarning"), + pytest.mark.filterwarnings("ignore::RuntimeWarning"), + pytest.mark.filterwarnings("ignore::DeprecationWarning"), + ] + for item in items: + if item.path.is_relative_to(_KERNELS_DIR): + for mark in ignore_marks: + item.add_marker(mark) + def make_qkv(total, num_heads, num_kv_heads, head_dim, device="cuda", dtype=torch.float16): """Create packed Q, K, V tensors.""" diff --git a/tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py b/tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py index 96cc24c2b98..13f3335e1c6 100644 --- a/tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py +++ b/tests/gpu/torch/kernels/quantization/conv/test_implicit_gemm.py @@ -18,6 +18,8 @@ Tests both non-quantized path (vs cuDNN) and FP4-quantized path (vs Triton reference). """ +import math + import pytest import torch import torch.nn.functional as F @@ -552,7 +554,6 @@ def _py_fp4_fake_quant_ref(x_flat, global_amax, block_size): 2. Per block: block_max = max(|x|), scale = fp8_e4m3_roundtrip(block_max / (6 * global_scale)) * global_scale 3. Quantize each element to nearest E2M1 level, then dequantize. """ - import math # E2M1 quantization levels: {0, 0.5, 1, 1.5, 2, 3, 4, 6} # Boundaries (midpoints): <=0.25->0, <0.75->0.5, <=1.25->1, <1.75->1.5, <=2.5->2, <3.5->3, <=5->4, >5->6 diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py b/tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py index 8971d243fdb..35fcb39e42b 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_diffusers_triton_attention.py @@ -18,12 +18,6 @@ import pytest import torch -pytestmark = [ - pytest.mark.filterwarnings("ignore::UserWarning"), - pytest.mark.filterwarnings("ignore::RuntimeWarning"), - pytest.mark.filterwarnings("ignore::DeprecationWarning"), -] - diffusers = pytest.importorskip("diffusers") from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py index fe16559a187..abdb37afe2d 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_calibrate.py @@ -27,12 +27,6 @@ import torch from conftest import make_qkv, make_varlen_meta -pytestmark = [ - pytest.mark.filterwarnings("ignore::UserWarning"), - pytest.mark.filterwarnings("ignore::RuntimeWarning"), - pytest.mark.filterwarnings("ignore::DeprecationWarning"), -] - from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py index d694d130cd9..fc26c5db17c 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_skip_softmax.py @@ -18,14 +18,11 @@ import pytest import torch from conftest import make_varlen_meta +from transformers import AutoModelForCausalLM, AutoTokenizer -pytestmark = [ - pytest.mark.filterwarnings("ignore::UserWarning"), - pytest.mark.filterwarnings("ignore::RuntimeWarning"), - pytest.mark.filterwarnings("ignore::DeprecationWarning"), -] - +import modelopt.torch.sparsity.attention_sparsity as mtsa from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.sparsity.attention_sparsity.methods.flash_skip_softmax import FlashSkipSoftmax if TRITON_KERNEL_AVAILABLE: from modelopt.torch.kernels.common.attention import attention, register_triton_attention @@ -176,10 +173,6 @@ def test_triton_matches_pytorch_reference(self): and applies them to standard softmax attention. The Triton kernel fuses the same skip logic into the online softmax inner loop. """ - from modelopt.torch.sparsity.attention_sparsity.methods.flash_skip_softmax import ( - FlashSkipSoftmax, - ) - batch, seq_len = 1, 256 num_heads, num_kv_heads, head_dim = 4, 4, 64 # MHA for simplicity scale = 1.0 / (head_dim**0.5) @@ -248,11 +241,6 @@ class TestSkipSoftmaxHFIntegration: def test_skip_softmax_via_sparsify(self, tiny_llama_dir): """mtsa.sparsify() with triton_skip_softmax produces finite logits.""" - pytest.importorskip("transformers") - from transformers import AutoModelForCausalLM, AutoTokenizer - - import modelopt.torch.sparsity.attention_sparsity as mtsa - tok = AutoTokenizer.from_pretrained(tiny_llama_dir) if tok.pad_token_id is None: tok.pad_token_id = tok.eos_token_id diff --git a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py index 054428a3a93..651aeab4116 100644 --- a/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py +++ b/tests/gpu/torch/kernels/sparsity/attention/test_triton_fa_sparse_nm.py @@ -20,13 +20,9 @@ import pytest import torch from conftest import make_qkv, make_varlen_meta +from transformers import AutoModelForCausalLM, AutoTokenizer -pytestmark = [ - pytest.mark.filterwarnings("ignore::UserWarning"), - pytest.mark.filterwarnings("ignore::RuntimeWarning"), - pytest.mark.filterwarnings("ignore::DeprecationWarning"), -] - +import modelopt.torch.sparsity.attention_sparsity as mtsa from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: @@ -422,11 +418,6 @@ def test_sparse_disabled_matches_dense(self): def test_sparse_nm_via_sparsify(self, tiny_llama_dir): """mtsa.sparsify() with N:M sparse softmax produces finite logits that differ from dense.""" - pytest.importorskip("transformers") - from transformers import AutoModelForCausalLM, AutoTokenizer - - import modelopt.torch.sparsity.attention_sparsity as mtsa - tok = AutoTokenizer.from_pretrained(tiny_llama_dir) if tok.pad_token_id is None: tok.pad_token_id = tok.eos_token_id diff --git a/tests/gpu/torch/nas/test_distributed_model.py b/tests/gpu/torch/nas/test_distributed_model.py index 1bf31826151..3dc18d27f9a 100644 --- a/tests/gpu/torch/nas/test_distributed_model.py +++ b/tests/gpu/torch/nas/test_distributed_model.py @@ -15,16 +15,19 @@ from functools import partial -import pytest import torch import torch.distributed -from _test_utils.torch.distributed.fsdp_test import run_fsdp2_test, run_fsdp_test +from _test_utils.torch.distributed.fsdp_test import run_fsdp2_test from _test_utils.torch.distributed.utils import synchronize_state_dict from torch import nn from torchvision.models.resnet import Bottleneck from modelopt.torch.nas.search_space import SearchSpace, generate_search_space +# NAS distributed-model tests are for rarely-used fastnas/autonas paths. Redundant +# parametrize cases are marked ``@pytest.mark.manual`` (run with ``--run-manual``); one +# representative case per test stays enabled for sanity. + def _get_test_case(): model = Bottleneck(32, 8) @@ -63,19 +66,6 @@ def _sample_subnet(model): SearchSpace(model).sample(sample_func=min) -@pytest.mark.parametrize("use_orig_params", [False, True]) -def test_fsdp(need_2_gpus, dist_workers, use_orig_params): - dist_workers.run( - partial( - run_fsdp_test, - _get_test_case, - "conv1", - _sample_subnet, - fsdp_kwargs={"use_orig_params": use_orig_params}, - ), - ) - - def test_fsdp2(need_2_gpus, dist_workers): dist_workers.run( partial(run_fsdp2_test, _get_test_case, "conv1", _sample_subnet), diff --git a/tests/gpu/torch/nas/test_search_space_with_vision_models.py b/tests/gpu/torch/nas/test_search_space_with_vision_models.py index f1afe495268..ae97f11f781 100644 --- a/tests/gpu/torch/nas/test_search_space_with_vision_models.py +++ b/tests/gpu/torch/nas/test_search_space_with_vision_models.py @@ -22,11 +22,22 @@ from modelopt.torch.utils import flatten_tree, zero_grad from modelopt.torch.utils.random import _set_deterministic_seed +# NAS search-space tests are for rarely-used fastnas/autonas paths. All but one vision +# model are marked ``@pytest.mark.manual`` (run with ``--run-manual``); one representative +# model per test stays enabled for sanity. models = get_vision_models() +def _model_params(): + """One representative model enabled; the rest marked manual.""" + return [ + pytest.param(fn, id=name, marks=([] if i == 0 else [pytest.mark.manual])) + for i, (name, fn) in enumerate(models.items()) + ] + + @pytest.mark.parametrize("on_gpu", [True]) # just run on GPU, but leave it here for easy debugging -@pytest.mark.parametrize("get_model_and_input", models.values(), ids=models.keys()) +@pytest.mark.parametrize("get_model_and_input", _model_params()) def test_models(get_model_and_input, on_gpu): _set_deterministic_seed() @@ -43,12 +54,10 @@ def test_models(get_model_and_input, on_gpu): # Test subnet forwardpytest.warns out1 = search_space.model(*args, **kwargs) - # Test subnet backward - with torch.autograd.set_detect_anomaly(True): - for t in flatten_tree(out1)[0]: - if isinstance(t, torch.Tensor) and t.requires_grad: - torch.sum(t).backward() - zero_grad(search_space.model) + for t in flatten_tree(out1)[0]: + if isinstance(t, torch.Tensor) and t.requires_grad: + torch.sum(t).backward() + zero_grad(search_space.model) # Test model export subnet = search_space.export() @@ -60,7 +69,7 @@ def test_models(get_model_and_input, on_gpu): # NOTE: we run this test on CPU because of better floating point precision! @pytest.mark.parametrize("on_gpu", [False]) # don't run on GPU but leave it here for easy debugging -@pytest.mark.parametrize("get_model_and_input", models.values(), ids=models.keys()) +@pytest.mark.parametrize("get_model_and_input", _model_params()) def test_dynamic_sorting(get_model_and_input, on_gpu): set_seed() # initialize model diff --git a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py index 49fdf24d145..03d20d26746 100644 --- a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py +++ b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py @@ -20,8 +20,10 @@ import pytest import torch +import torch.nn as nn from _test_utils.torch.transformers_models import create_tiny_llama_dir from accelerate import init_empty_weights, load_checkpoint_and_dispatch +from accelerate.hooks import AlignDevicesHook, add_hook_to_module from transformers import AutoConfig, AutoModelForCausalLM import modelopt.torch.quantization as mtq @@ -30,8 +32,13 @@ from modelopt.torch.quantization.utils import ( enable_weight_access_and_writeback, is_quantized_linear, + persistent_materialization, +) +from modelopt.torch.quantization.utils.layerwise_calib import ( + LayerActivationCollector, + _layer_dir, + _SkipLayer, ) -from modelopt.torch.quantization.utils.layerwise_calib import _layer_dir NVFP4_WEIGHT_MSE_FP8_SWEEP_CFG = { "quant_cfg": [ @@ -467,13 +474,6 @@ def forward(self, x): def test_skip_dummy_has_no_hf_hook(monkeypatch): """Dummies must not carry _hf_hook from the original layer.""" - from accelerate.hooks import AlignDevicesHook, add_hook_to_module - - from modelopt.torch.quantization.utils.layerwise_calib import ( - LayerActivationCollector, - _SkipLayer, - ) - monkeypatch.setattr( LayerActivationCollector, "_decoder_layer_support", @@ -507,11 +507,6 @@ def forward_loop(m): def test_persistent_materialization_cpu_offloaded(tmp_path): """persistent_materialization keeps CPU-offloaded weights on GPU and writes back modifications.""" - import torch.nn as nn - from accelerate.hooks import AlignDevicesHook - - from modelopt.torch.quantization.utils import persistent_materialization - model, config, _, inputs = _make_cpu_offloaded_model(tmp_path) offloaded_layer = model.model.layers[0] @@ -625,11 +620,6 @@ def test_disk_offloaded_tinyllama(tmp_path): def test_persistent_materialization_disk_offloaded(tmp_path): """persistent_materialization keeps disk-offloaded weights on GPU and writes back modifications.""" - import torch.nn as nn - from accelerate.hooks import AlignDevicesHook - - from modelopt.torch.quantization.utils import persistent_materialization - model, config, _, inputs = _make_disk_offloaded_model(tmp_path) offloaded_layer = model.model.layers[0] diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index c5584ece5cf..55648ac26e2 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -23,9 +23,15 @@ import torch.nn as nn from _test_utils.torch.distributed.utils import synchronize_state_dict from torch.distributed._composable.fsdp.fully_shard import fully_shard +from torch.distributed.tensor import DTensor import modelopt.torch.quantization as mtq from modelopt.torch.opt.dynamic import _pytorch_managed +from modelopt.torch.quantization.utils import ( + enable_weight_access_and_writeback, + persistent_materialization, +) +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector def _test_fsdp2_simple_linear(rank, size): @@ -162,8 +168,6 @@ def forward(self, x): def _test_layerwise_calibrate_fsdp2(rank, size): """Layerwise calibration on FSDP2-wrapped model matches non-FSDP reference.""" - from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector - dim = 32 torch.manual_seed(1) model = _SimpleTransformerModel(n_layers=3, dim=dim).cuda() @@ -206,13 +210,6 @@ def test_layerwise_calibrate_fsdp2(dist_workers): def _test_persistent_materialization(rank, size): """persistent_materialization keeps weights accessible and writes back modifications.""" - from torch.distributed.tensor import DTensor - - from modelopt.torch.quantization.utils import ( - enable_weight_access_and_writeback, - persistent_materialization, - ) - dim = 32 torch.manual_seed(1) model = nn.Sequential( diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index c29401a8460..d1d8c0c23d4 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -18,7 +18,7 @@ import pytest import torch -from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_tokenizer +from _test_utils.torch.transformers_models import get_tiny_llama from conftest import requires_triton import modelopt.torch.quantization as mtq @@ -213,23 +213,15 @@ 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): - tokenizer = get_tiny_tokenizer() - model = get_tiny_llama(vocab_size=tokenizer.vocab_size).to("cuda") - - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - tokenizer.padding_side = "left" - - assert tokenizer.pad_token is not None, "Pad token cannot be set!" +def test_gptq_e2e_flow(quant_cfg, tiny_tokenizer): + 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": True} calib_dataloader = get_dataset_dataloader( dataset_name="cnn_dailymail", - tokenizer=tokenizer, + tokenizer=tiny_tokenizer, batch_size=2, num_samples=8, device="cuda", diff --git a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py index 5e64c4a312a..d1eba4987d3 100644 --- a/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py +++ b/tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py @@ -33,6 +33,7 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep from modelopt.torch.quantization.calib import NVFP4MSECalibrator from modelopt.torch.quantization.extensions import get_cuda_ext_mx from modelopt.torch.quantization.nn import TensorQuantizer @@ -205,8 +206,6 @@ def test_reset_allows_recollect(): @requires_triton def test_input_validation(): """``nvfp4_fp8_scale_sweep`` should reject malformed inputs cleanly.""" - from modelopt.torch.kernels.quantization.gemm import nvfp4_fp8_scale_sweep - device = "cuda" x = torch.randn(64, BLOCK_SIZE, device=device) g = x.abs().amax() diff --git a/tests/gpu/torch/quantization/test_real_quantize_cuda.py b/tests/gpu/torch/quantization/test_real_quantize_cuda.py index 72af60ec37c..74ff7e93048 100644 --- a/tests/gpu/torch/quantization/test_real_quantize_cuda.py +++ b/tests/gpu/torch/quantization/test_real_quantize_cuda.py @@ -26,6 +26,7 @@ from _test_utils.torch.transformers_models import create_tiny_llama_dir from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.tensor import DTensor +from transformers import AutoModelForCausalLM import modelopt.torch.quantization as mtq from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights @@ -197,11 +198,6 @@ def forward_loop(model, run_backward=False): @pytest.mark.parametrize("quant_config", [mtq.FP8_DEFAULT_CFG]) def test_real_quantize_linear(quant_config, tmp_path): - try: - from transformers import AutoModelForCausalLM - except ImportError: - pytest.skip("transformers is not installed") - tiny_llama_dir = create_tiny_llama_dir(tmp_path) with init_quantized_weights(quant_config): model = AutoModelForCausalLM.from_pretrained(tiny_llama_dir) diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py b/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py index e97438a4e5a..93e060f0003 100644 --- a/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py +++ b/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py @@ -33,6 +33,10 @@ diffusers = pytest.importorskip("diffusers") +import numpy as np +from diffusers import WanPipeline + +import modelopt.torch.opt as mto from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE if TRITON_KERNEL_AVAILABLE: @@ -56,8 +60,6 @@ def tiny_wan22_path(tmp_path_factory): @pytest.fixture def tiny_wan22_pipe(tiny_wan22_path): """Load a fresh copy of the tiny Wan 2.2 pipeline on CUDA (per test).""" - from diffusers import WanPipeline - pipe = WanPipeline.from_pretrained(tiny_wan22_path, torch_dtype=torch.bfloat16) pipe.to("cuda") return pipe @@ -145,8 +147,6 @@ def test_skip_softmax_pipeline_runs_e2e(self, tiny_wan22_pipe): def test_tight_threshold_matches_dense_within_tolerance(self, tiny_wan22_pipe, tiny_wan22_path): """A near-zero threshold is effectively dense and close to unsparsified.""" - from diffusers import WanPipeline - # Dense run: fresh pipe, no sparsification dense_pipe = WanPipeline.from_pretrained(tiny_wan22_path, torch_dtype=torch.bfloat16) dense_pipe.to("cuda") @@ -159,8 +159,6 @@ def test_tight_threshold_matches_dense_within_tolerance(self, tiny_wan22_pipe, t sparse_frame0 = _run_pipe(tiny_wan22_pipe).frames[0][0] # Both are PIL images — convert to tensor and compare - import numpy as np - d = np.asarray(dense_frame0, dtype=np.float32) s = np.asarray(sparse_frame0, dtype=np.float32) # Pixel-wise MAE should be small for tight threshold (but not bit-exact due to @@ -208,8 +206,6 @@ def test_save_restore_roundtrip(self, tiny_wan22_pipe): """ from _test_utils.torch.diffusers_models import get_tiny_wan22_transformer - import modelopt.torch.opt as mto - _sparsify_both_transformers(tiny_wan22_pipe, _skip_softmax_cfg()) state = mto.modelopt_state(tiny_wan22_pipe.transformer) diff --git a/tests/gpu/torch/utils/test_dataset_utils.py b/tests/gpu/torch/utils/test_dataset_utils.py new file mode 100644 index 00000000000..72cd80b0558 --- /dev/null +++ b/tests/gpu/torch/utils/test_dataset_utils.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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. + +"""Dataset tests that reach the HuggingFace Hub. + +These live under ``tests/gpu`` (networked infra) rather than ``tests/unit``, +which is kept hermetic with toy local fixtures. +""" + +import json + +import pytest +from datasets import load_dataset +from huggingface_hub import get_token + +from modelopt.torch.utils.dataset_utils import ( + SUPPORTED_DATASET_CONFIG, + get_dataset_dataloader, + get_dataset_samples, +) + + +def _write_jsonl(path, rows): + """Write a list of dicts to *path* as JSONL. Returns the path as ``str``.""" + with open(path, "w", encoding="utf-8") as f: + f.writelines(json.dumps(row) + "\n" for row in rows) + return str(path) + + +_NEW_NEMOTRON_KEYS = [ + "nemotron-sft-instruction-following-chat-v2", + "nemotron-science-v1", + "nemotron-competitive-programming-v1", + "nemotron-sft-agentic-v2", + "nemotron-math-v2", + "nemotron-sft-swe-v2", + "nemotron-sft-multilingual-v1", +] + + +@pytest.mark.parametrize("dataset_key", _NEW_NEMOTRON_KEYS) +def test_new_nemotron_registry_shape(dataset_key): + """Shape check on the 7 newly registered nvidia/Nemotron-* entries. + + Complements the gated smoke test below — catches typos in dataset paths or + split names even when the runner has no HF credentials. + """ + assert dataset_key in SUPPORTED_DATASET_CONFIG + entry = SUPPORTED_DATASET_CONFIG[dataset_key] + config = entry["config"] + assert config["path"].startswith("nvidia/Nemotron-") + splits = config["split"] + assert isinstance(splits, list) and splits + assert all(isinstance(s, str) and s for s in splits) + assert len(set(splits)) == len(splits) + assert callable(entry["preprocess"]) + assert entry["chat_key"] == "messages" + + +@pytest.mark.parametrize( + "dataset_key", + [ + # nvidia/Nemotron-SFT-Agentic-v2's ``search`` split currently fails the HF datasets + # JSON schema cast (upstream column drift), so skip its live download smoke; the + # registry shape check above still validates its configuration. + pytest.param(k, marks=pytest.mark.skip(reason="upstream search-split schema cast failure")) + if k == "nemotron-sft-agentic-v2" + else k + for k in _NEW_NEMOTRON_KEYS + ], +) +def test_get_dataset_samples_new_nemotron(dataset_key): + """Smoke-test the 7 newly registered nvidia/Nemotron-* calibration datasets. + + Skipped when no HF token is available because these datasets live behind the HF Hub. + ``huggingface_hub.get_token()`` covers both the ``HF_TOKEN`` env var and tokens + cached by ``hf auth login``. + """ + if not get_token(): + pytest.skip( + "No HF token (env HF_TOKEN or `hf auth login`); skipping gated Nemotron smoke test" + ) + + samples = get_dataset_samples(dataset_key, num_samples=2) + assert isinstance(samples, list) + assert len(samples) == 2 + assert all(isinstance(s, str) and len(s) > 0 for s in samples) + + +# Live HF dataset round-trips. ``hf-internal-testing/dataset_with_data_files`` is a +# 10-row x {train,test} fixture maintained by HF for their own CI — tiny enough to +# download in a test and stable across releases, and ungated (no HF token needed). +_HF_TINY = "hf-internal-testing/dataset_with_data_files" # train, test splits, ``text`` col + + +def _hf_dump_to_jsonl(name: str, split: str, path) -> str: + ds = load_dataset(name, split=split) + ds.to_json(str(path), lines=True) + return str(path) + + +class TestHfTinyDataset: + """End-to-end coverage of the Hub-download branch with a real (tiny) HF dataset.""" + + def test_load_single_split_directly(self): + samples = get_dataset_samples(_HF_TINY, num_samples=4, split="train") + assert len(samples) == 4 + assert all(isinstance(s, str) and s for s in samples) + + def test_load_multiple_splits_directly(self): + """``split=["train", "test"]`` divides ``num_samples`` across both splits.""" + samples = get_dataset_samples(_HF_TINY, num_samples=6, split=["train", "test"]) + assert len(samples) == 6 + # Both splits should contribute; confirm by comparing against direct loads. + train_only = set(get_dataset_samples(_HF_TINY, num_samples=10, split="train")) + test_only = set(get_dataset_samples(_HF_TINY, num_samples=10, split="test")) + assert any(s in train_only for s in samples) + assert any(s in test_only for s in samples) + + def test_default_split_is_train(self): + default_samples = get_dataset_samples(_HF_TINY, num_samples=4) + train_samples = get_dataset_samples(_HF_TINY, num_samples=4, split="train") + assert default_samples == train_samples + + def test_download_to_jsonl_then_load(self, tmp_path): + """Dump the HF dataset to JSONL, then reload it via the local-jsonl path.""" + jsonl_path = _hf_dump_to_jsonl(_HF_TINY, "train", tmp_path / "train.jsonl") + from_jsonl = get_dataset_samples(jsonl_path, num_samples=10) + from_hf = get_dataset_samples(_HF_TINY, num_samples=10, split="train") + assert from_jsonl == from_hf + + def test_dataloader_blending_two_hf_datasets(self, tiny_tokenizer): + """Two HF datasets concatenated via ``get_dataset_dataloader``.""" + loader = get_dataset_dataloader( + dataset_name=[_HF_TINY, "hf-internal-testing/multi_dir_dataset"], + tokenizer=tiny_tokenizer, + batch_size=4, + num_samples=[3, 1], + max_sample_length=16, + ) + batches = list(loader) + assert sum(b["input_ids"].shape[0] for b in batches) == 4 + + def test_dataloader_mixing_hf_and_local_jsonl(self, tmp_path, tiny_tokenizer): + """Live HF dataset blended with a local synthetic JSONL file.""" + local = _write_jsonl(tmp_path / "local.jsonl", [{"text": f"local {i}"} for i in range(2)]) + loader = get_dataset_dataloader( + dataset_name=[_HF_TINY, local], + tokenizer=tiny_tokenizer, + batch_size=5, + num_samples=[3, 2], + max_sample_length=16, + ) + batches = list(loader) + assert sum(b["input_ids"].shape[0] for b in batches) == 5 diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 76c1fbbb0da..405a83db06d 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -51,34 +51,49 @@ def _make_pool(world_size): @pytest.fixture(scope="module") -def dist_workers(): +def _pool_cache(): + """Module-scoped cache of worker pools keyed by world_size. + + Spinning up a pool cold-imports the full torch/megatron/modelopt stack per worker + (~tens of seconds), so fixtures that request the same world_size share one pool + instead of spawning duplicates — e.g. on a 2-GPU runner ``dist_workers`` and + ``dist_workers_size_2`` are both size 2. The cache is module-scoped and torn down at + module end, so workers are never reused across modules (avoids cross-test + state contamination). + """ + pools: dict[int, DistributedWorkerPool] = {} + yield pools + for pool in pools.values(): + pool.shutdown() + + +def _get_pool(cache, world_size): + if world_size not in cache: + cache[world_size] = _make_pool(world_size) + return cache[world_size] + + +@pytest.fixture(scope="module") +def dist_workers(_pool_cache): """Module-scoped pool with world_size=torch.cuda.device_count().""" - pool = _make_pool(torch.cuda.device_count()) - yield pool - pool.shutdown() + return _get_pool(_pool_cache, torch.cuda.device_count()) @pytest.fixture(scope="module") -def dist_workers_size_1(): +def dist_workers_size_1(_pool_cache): """Module-scoped pool with world_size=1 for tests that require a single process.""" - pool = _make_pool(1) - yield pool - pool.shutdown() + return _get_pool(_pool_cache, 1) @pytest.fixture(scope="module") -def dist_workers_size_2(): +def dist_workers_size_2(_pool_cache): if torch.cuda.device_count() < 2: pytest.skip("Need at least 2 GPUs") - pool = _make_pool(2) - yield pool - pool.shutdown() + return _get_pool(_pool_cache, 2) @pytest.fixture(scope="module") -def dist_workers_size_4(): +def dist_workers_size_4(_pool_cache): if torch.cuda.device_count() < 4: pytest.skip("Need at least 4 GPUs") - pool = _make_pool(4) - yield pool - pool.shutdown() + return _get_pool(_pool_cache, 4) diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 4ba5b9e2863..f818cb3594c 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -30,6 +30,7 @@ ) from safetensors import safe_open from safetensors.torch import save_file +from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration import modelopt.torch.quantization as mtq import modelopt.torch.speculative as mtsp @@ -179,8 +180,6 @@ def _test_unified_export_megatron( "vision encoder keys missing from export" ) # try to load the model and run a forward pass - from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration - vl_model = Qwen3VLForConditionalGeneration.from_pretrained( tmp_export_dir, torch_dtype=torch.bfloat16 ).cuda() diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index c34fb2df376..9cd7fb228f0 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -14,6 +14,7 @@ # limitations under the License. import copy +from contextlib import nullcontext from functools import partial import pytest @@ -251,8 +252,6 @@ def _gpt_model_provider( hybrid_override_pattern=None, mamba_head_dim=16, ): - from contextlib import nullcontext - device_ctx = torch.device("meta") if meta_device else nullcontext() with device_ctx: diff --git a/tests/gpu_megatron/torch/quantization/test_mse_calibrator_mixed_precision.py b/tests/gpu_megatron/torch/quantization/test_mse_calibrator_mixed_precision.py index 1412367494f..6e7b988017c 100644 --- a/tests/gpu_megatron/torch/quantization/test_mse_calibrator_mixed_precision.py +++ b/tests/gpu_megatron/torch/quantization/test_mse_calibrator_mixed_precision.py @@ -19,7 +19,10 @@ import pytest import torch +import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_utils import get_weight_scaling_factor from modelopt.torch.quantization.calib import MseCalibrator +from modelopt.torch.quantization.model_calib import mse_calibrate from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer @@ -59,9 +62,6 @@ def test_mixed_nvfp4_fp8_sweep_true_skips_fp8(): """fp8_scale_sweep=True: NVFP4 layer is promoted to NVFP4StaticQuantizer; the FP8 layer is left as a plain TensorQuantizer (no backend factory registered → no MseCalibrator replacement, max-calibrated amax preserved).""" - import modelopt.torch.quantization as mtq - from modelopt.torch.quantization.model_calib import mse_calibrate - device = torch.device("cuda") model = _TwoLayer().to(device) inputs = torch.randn(1, 16, device=device) @@ -79,9 +79,6 @@ def test_mixed_nvfp4_fp8_sweep_false_uses_mse_for_both(): """fp8_scale_sweep=False: both NVFP4 and FP8 layers get an MseCalibrator. NVFP4 layer is still promoted to NVFP4StaticQuantizer (promotion is independent of the sweep flag).""" - import modelopt.torch.quantization as mtq - from modelopt.torch.quantization.model_calib import mse_calibrate - device = torch.device("cuda") model = _TwoLayer().to(device) inputs = torch.randn(1, 16, device=device) @@ -100,8 +97,6 @@ def test_output_layer_nvfp4_promotion_and_forward(): NVFP4StaticQuantizer and its forward dispatches cleanly through the static blockwise FP4 kernel (regression for the lm_head crash that motivated the NVFP4 promotion in mse_calibrate).""" - import modelopt.torch.quantization as mtq - from modelopt.torch.quantization.model_calib import mse_calibrate class _WithOutputLayer(torch.nn.Module): def __init__(self): @@ -146,9 +141,6 @@ def forward(self, x): @pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4 path requires CUDA") def test_output_layer_nvfp4_export_keys(): """A W4A16-quantized output_layer exports with CT-style weight + scale keys.""" - import modelopt.torch.quantization as mtq - from modelopt.torch.export.quant_utils import get_weight_scaling_factor - from modelopt.torch.quantization.model_calib import mse_calibrate class _OutputOnly(torch.nn.Module): def __init__(self): diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py similarity index 99% rename from tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py rename to tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index 2dd2ba587fe..e9f8ee73d95 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -13,16 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for sparse attention vLLM worker compatibility helpers.""" +"""Tests for sparse attention vLLM worker compatibility helpers.""" import math from contextlib import nullcontext import pytest import torch - -pytest.importorskip("vllm") - from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl from modelopt.torch.sparsity.attention_sparsity.plugins import vllm as vllm_plugin diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_vllm_plugin.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py similarity index 99% rename from tests/gpu/torch/sparsity/attention_sparsity/test_vllm_plugin.py rename to tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py index 4c029da751c..fa11b144354 100644 --- a/tests/gpu/torch/sparsity/attention_sparsity/test_vllm_plugin.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_plugin.py @@ -30,9 +30,6 @@ import pytest import torch - -pytest.importorskip("vllm") - from vllm.v1.attention.backends.flash_attn import FlashAttentionImpl from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE diff --git a/tests/regression/torch/speculative/test_dflash.py b/tests/regression/torch/speculative/test_dflash.py index 2e8092d62f5..18e95148fcd 100644 --- a/tests/regression/torch/speculative/test_dflash.py +++ b/tests/regression/torch/speculative/test_dflash.py @@ -73,6 +73,7 @@ def dflash_output_dir(tmp_path_factory): return tmp_path_factory.mktemp("dflash_output") +@pytest.mark.timeout(300) def test_dflash_training(qwen3_model_name, dflash_output_dir): """Train DFlash on Qwen3-0.6B and validate loss convergence.""" output_dir = str(dflash_output_dir / "dflash-qwen3-0.6b") diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 00000000000..f397205f022 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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. + +import contextlib +import os + +# Enforce no HuggingFace Hub network access for unit tests +os.environ["HF_HUB_OFFLINE"] = "1" +os.environ["HF_DATASETS_OFFLINE"] = "1" +os.environ["TRANSFORMERS_OFFLINE"] = "1" + +with contextlib.suppress(ImportError): + import huggingface_hub.constants as _hf_constants + + _hf_constants.HF_HUB_OFFLINE = True diff --git a/tests/unit/onnx/autocast/test_nodeclassifier.py b/tests/unit/onnx/autocast/test_nodeclassifier.py index e604fb7c1cc..002189b6bd3 100644 --- a/tests/unit/onnx/autocast/test_nodeclassifier.py +++ b/tests/unit/onnx/autocast/test_nodeclassifier.py @@ -14,6 +14,7 @@ # limitations under the License. import os +import tempfile from collections import OrderedDict import numpy as np @@ -29,6 +30,7 @@ InitializerRangeRule, IORangeRule, NodeClassifier, + NodeRuleBase, ) from modelopt.onnx.autocast.referencerunner import ReferenceRunner, TensorStats @@ -280,8 +282,6 @@ def test_node_classifier_custom_rule(test_model): "mul_node": [numpy_helper.from_array(np.array([[0.5, 0.5], [0.5, 0.5]], dtype=np.float32))], } - from modelopt.onnx.autocast.nodeclassifier import NodeRuleBase - class CustomRule(NodeRuleBase): def _check_inner(self, node): # Return True if any initializer contains zeros @@ -520,8 +520,6 @@ def test_depth_of_reduction_rule_with_tensor_stats(): ) def test_node_classifier_with_multi_batch_calibration(test_model): """Test NodeClassifier with multi-batch calibration data.""" - import tempfile - node_to_init_map = {key: [] for key in ["add_node", "mul_node"]} # Create multiple batches of calibration data diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 82155e42471..00075714c69 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/tests/unit/onnx/autocast/test_referencerunner.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os import tempfile from collections import OrderedDict @@ -170,8 +171,6 @@ def test_invalid_json(reference_runner): inputs = {"X1": [[1.0, 2.0, 3.0]], "X2": [[4.0, 5.0, 6.0]]} with tempfile.NamedTemporaryFile(suffix=".json", mode="w", delete=False) as f: - import json - json.dump(inputs, f) input_path = f.name try: diff --git a/tests/unit/onnx/quantization/test_convtranspose_qdq.py b/tests/unit/onnx/quantization/test_convtranspose_qdq.py index d06e8fc6d7f..1db505fd4b0 100644 --- a/tests/unit/onnx/quantization/test_convtranspose_qdq.py +++ b/tests/unit/onnx/quantization/test_convtranspose_qdq.py @@ -24,9 +24,14 @@ @pytest.fixture def model_and_input(): - """Create model and dummy input.""" + """Create model and dummy input. + + Input kept small (32x32) — this asserts QDQ placement around ConvTranspose + weights, which is independent of spatial size, so a tiny input keeps the ORT + calibration pass fast. UNet downsamples by 8, so 32 is the minimum sensible size. + """ model = UNet() - dummy_input = torch.randn(2, 1, 256, 256) + dummy_input = torch.randn(1, 1, 32, 32) return model, dummy_input diff --git a/tests/unit/onnx/quantization/test_qdq_utils.py b/tests/unit/onnx/quantization/test_qdq_utils.py index 0ff3686a610..8794066554e 100644 --- a/tests/unit/onnx/quantization/test_qdq_utils.py +++ b/tests/unit/onnx/quantization/test_qdq_utils.py @@ -13,12 +13,25 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings + import numpy as np +import onnx_graphsurgeon as gs +import onnxruntime as ort import pytest from onnx import TensorProto, helper, numpy_helper from modelopt.onnx.export import INT4QuantExporter, MXFP8QuantExporter, NVFP4QuantExporter from modelopt.onnx.export.nvfp4_exporter import _cast_fp4, _cast_fp8 +from modelopt.onnx.quantization.qdq_utils import ( + apply_column_major_transformation, + fp4qdq_to_2dq, + insert_transpose_nodes_for_column_major, + quantize_weights_to_int4, + quantize_weights_to_mxfp8, + replace_zero_scale_with_smallest_nonzero, +) +from modelopt.onnx.quantization.quant_utils import pack_float32_to_4bit_cpp_based def create_test_model_with_int4_dq_reshape_transpose_matmul(constant_scale: bool = False): @@ -637,8 +650,6 @@ def create_test_model_with_int4_dq_matmul(): Returns the model and original weight/scale arrays for verification. """ - from modelopt.onnx.quantization.quant_utils import pack_float32_to_4bit_cpp_based - # Create INT4 quantized weight tensor (K=32, N=16) # Using int8 storage for INT4 values in range [-8, 7] weight_data = np.random.randint(-8, 8, size=(32, 16), dtype=np.int8) @@ -699,13 +710,6 @@ def test_column_major_transformation_graph_structure(self): Verifies: DQ(W) -> MatMul becomes DQ(W^T) -> Transpose -> MatMul """ - import onnx_graphsurgeon as gs - - from modelopt.onnx.quantization.qdq_utils import ( - apply_column_major_transformation, - insert_transpose_nodes_for_column_major, - ) - model, original_weight, original_scale = create_test_model_with_int4_dq_matmul() # Get weights and scales as dicts (simulating what int4.py does) @@ -770,10 +774,6 @@ def test_column_major_transformation_output_equivalence(self): Verifies both produce the same output for the same input. """ - import onnxruntime as ort - - from modelopt.onnx.quantization.quant_utils import pack_float32_to_4bit_cpp_based - # Create original model original_model, original_weight, original_scale = create_test_model_with_int4_dq_matmul() @@ -873,14 +873,6 @@ def test_column_major_gemm_trans_b_flip(self): should have transB flipped to 0 instead of inserting a Transpose node. Also verifies output equivalence between original and transformed models. """ - import onnx_graphsurgeon as gs - import onnxruntime as ort - - from modelopt.onnx.quantization.qdq_utils import ( - apply_column_major_transformation, - insert_transpose_nodes_for_column_major, - ) - from modelopt.onnx.quantization.quant_utils import pack_float32_to_4bit_cpp_based # Original model: weight (N=16, K=32) with Gemm transB=1 # Gemm computes: A @ B^T = (4, 32) @ (16, 32)^T = (4, 16) @@ -1057,8 +1049,6 @@ class TestReplaceZeroScaleWithSmallestNonzero: @pytest.mark.parametrize("dq_op_type", ["DequantizeLinear", "TRT_INT4DequantizeLinear"]) def test_zero_scale_initializer_fed_to_dq_is_patched(self, dq_op_type): - from modelopt.onnx.quantization.qdq_utils import replace_zero_scale_with_smallest_nonzero - model = _build_model_with_zero_scale_initializer(dq_op_type) scale_before = numpy_helper.to_array( next(init for init in model.graph.initializer if init.name == "scale") @@ -1075,7 +1065,6 @@ def test_zero_scale_initializer_fed_to_dq_is_patched(self, dq_op_type): def test_constant_node_scale_path_still_patched(self): """Legacy Constant-node QDQ path must continue to be patched.""" - from modelopt.onnx.quantization.qdq_utils import replace_zero_scale_with_smallest_nonzero scale_data = np.array([1e-3, 0.0, 2e-3], dtype=np.float16) scale_const = helper.make_node( @@ -1120,10 +1109,6 @@ class TestLegacyEdgeLLMShims: """ def test_quantize_weights_to_int4_shim(self): - import warnings - - from modelopt.onnx.quantization.qdq_utils import quantize_weights_to_int4 - model = create_test_model_with_int4_dq_reshape_transpose_matmul() with warnings.catch_warnings(record=True) as caught: @@ -1146,10 +1131,6 @@ def test_quantize_weights_to_int4_shim(self): assert "Transpose" not in node_types def test_quantize_weights_to_mxfp8_shim(self): - import warnings - - from modelopt.onnx.quantization.qdq_utils import quantize_weights_to_mxfp8 - model = create_test_model_with_mxfp8_dq() with warnings.catch_warnings(record=True) as caught: @@ -1173,10 +1154,6 @@ def test_quantize_weights_to_mxfp8_shim(self): @pytest.mark.parametrize("with_transpose", [False, True]) def test_fp4qdq_to_2dq_shim(self, with_transpose): - import warnings - - from modelopt.onnx.quantization.qdq_utils import fp4qdq_to_2dq - model = create_test_model_with_nvfp4_qdq(with_transpose=with_transpose) with warnings.catch_warnings(record=True) as caught: diff --git a/tests/unit/onnx/test_gqa_graph_surgery.py b/tests/unit/onnx/test_gqa_graph_surgery.py index d6d92b90322..e31daa15afe 100644 --- a/tests/unit/onnx/test_gqa_graph_surgery.py +++ b/tests/unit/onnx/test_gqa_graph_surgery.py @@ -29,7 +29,13 @@ from modelopt.onnx.graph_surgery.gqa_replacement import replace_attention_with_gqa -MODEL_ID = "Qwen/Qwen2.5-0.5B" +# Attention shape mirroring ``Qwen/Qwen2.5-0.5B`` — built locally +_QWEN_CONFIG_KWARGS = { + "hidden_size": 896, + "num_attention_heads": 14, + "num_key_value_heads": 2, + "rope_theta": 1000000.0, +} VOCAB_SIZE = 64 SEQ_LEN = 4 MAX_SEQ_LEN = 128 @@ -583,10 +589,34 @@ def _rope(tensor, prefix, cos=cos_out, sin=sin_out): return model -def _get_config(): +def _run_session(model_proto, feeds): + """Run inference directly from an in-memory ModelProto.""" + model_bytes = model_proto.SerializeToString() + sess = ort.InferenceSession(model_bytes, providers=["CPUExecutionProvider"]) + return sess.run(None, feeds) + + +@pytest.fixture(scope="module") +def qwen_config_dir(tmp_path_factory): + """Local Qwen2 config dir so GQA surgery reads its config offline (no Hub download). + + ``replace_attention_with_gqa`` loads the config from ``hf_model_id`` itself + (via ``get_rope_caches``), so a real on-disk config is required — a bare + config object would not suffice. + """ + from transformers import Qwen2Config + + d = tmp_path_factory.mktemp("qwen_config") + Qwen2Config(**_QWEN_CONFIG_KWARGS).save_pretrained(str(d)) + return str(d) + + +@pytest.fixture(scope="module") +def qwen_attention_shapes(qwen_config_dir): + """Attention shape (hidden, heads, kv, head_dim, inv_freq) derived from the config.""" from transformers import AutoConfig - cfg = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=False) + cfg = AutoConfig.from_pretrained(qwen_config_dir, trust_remote_code=False) hidden = cfg.hidden_size heads = cfg.num_attention_heads kv = getattr(cfg, "num_key_value_heads", heads) @@ -596,17 +626,10 @@ def _get_config(): return hidden, heads, kv, hd, inv_freq.numpy() -def _run_session(model_proto, feeds): - """Run inference directly from an in-memory ModelProto.""" - model_bytes = model_proto.SerializeToString() - sess = ort.InferenceSession(model_bytes, providers=["CPUExecutionProvider"]) - return sess.run(None, feeds) - - @pytest.fixture(scope="module") -def models_and_config(): +def models_and_config(qwen_config_dir, qwen_attention_shapes): """Build original model, run GQA surgery, return both protos + config.""" - hidden, heads, kv, hd, inv_freq_np = _get_config() + hidden, heads, kv, hd, inv_freq_np = qwen_attention_shapes orig = _build_toy_model(hidden, heads, kv, hd, inv_freq_np) onnx.checker.check_model(orig) @@ -618,7 +641,7 @@ def models_and_config(): replace_attention_with_gqa( model_path=orig_path, output_path=gqa_path, - hf_model_id=MODEL_ID, + hf_model_id=qwen_config_dir, max_seq_len=MAX_SEQ_LEN, io_dtype="float16", use_external_data=False, diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 4c4e2d07ded..b5c433888d7 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -15,11 +15,16 @@ """Unit tests for modelopt.recipe.loader and modelopt.recipe.loader.load_config.""" +import json +import os import re +import sys +import types from importlib.resources import files import pytest +import modelopt.torch.quantization.config as qcfg from modelopt.recipe.config import ( ModelOptDFlashRecipe, ModelOptEagleRecipe, @@ -27,6 +32,8 @@ RecipeType, ) from modelopt.recipe.loader import _apply_dotlist, load_config, load_recipe +from modelopt.torch.opt.config_loader import _load_raw_config, _schema_type +from modelopt.torch.quantization.config import QuantizerAttributeConfig, normalize_quant_cfg_list # --------------------------------------------------------------------------- # Static YAML fixtures @@ -476,11 +483,6 @@ def test_load_recipe_dflash_field_validation_raises(tmp_path): ) def test_general_ptq_yaml_matches_config_dicts(yaml_path, model_cfg_name, kv_cfg_name): """Each general/ptq YAML's quant_cfg list matches the merged Python config dicts.""" - import json - - import modelopt.torch.quantization.config as qcfg - from modelopt.torch.quantization.config import normalize_quant_cfg_list - model_cfg = getattr(qcfg, model_cfg_name) kv_cfg = getattr(qcfg, kv_cfg_name) recipe = load_recipe(yaml_path) @@ -1277,8 +1279,6 @@ def test_import_circular_via_path_aliases_raises(tmp_path): f" algorithm: max\n" f" quant_cfg: []\n" ) - import os - cwd = os.getcwd() os.chdir(tmp_path) try: @@ -1369,8 +1369,6 @@ def test_builtin_config_snippets_with_modelopt_schema(config_path): def test_modelopt_schema_comment_returns_instance(tmp_path): """A ``modelopt-schema`` comment makes load_config return an instance of that schema.""" - from modelopt.torch.quantization.config import QuantizerAttributeConfig - config_file = tmp_path / "fp8.yaml" config_file.write_text( "# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig\n" @@ -1396,11 +1394,6 @@ def test_modelopt_schema_comment_validation_error(tmp_path): def test_modelopt_schema_reports_circular_resolution(monkeypatch): """A schema missing from an initializing module reports the likely circular import.""" - import sys - import types - - from modelopt.torch.opt.config_loader import _schema_type - module_name = "modelopt._schema_cycle_test" module = types.ModuleType(module_name) module.__spec__ = types.SimpleNamespace(_initializing=True) @@ -1537,8 +1530,6 @@ def test_load_config_multi_doc_dict_dict(tmp_path): """Multi-document YAML with two dicts merges them.""" cfg_file = tmp_path / "multi.yaml" cfg_file.write_text("imports:\n fp8: some/path\n---\nalgorithm: max\n") - from modelopt.torch.opt.config_loader import _load_raw_config - data = _load_raw_config(cfg_file) assert data["imports"] == {"fp8": "some/path"} assert data["algorithm"] == "max" @@ -1548,8 +1539,6 @@ def test_load_config_multi_doc_null_content(tmp_path): """Multi-document YAML where second doc is null treats content as empty dict.""" cfg_file = tmp_path / "multi_null.yaml" cfg_file.write_text("key: value\n---\n") - from modelopt.torch.opt.config_loader import _load_raw_config - data = _load_raw_config(cfg_file) assert data == {"key": "value"} diff --git a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py index 1f2165ed7af..399085ad699 100644 --- a/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py +++ b/tests/unit/torch/deploy/utils/test_torch_onnx_utils.py @@ -37,13 +37,32 @@ deploy_benchmark_all = get_deploy_models() deploy_benchmark_dynamo = get_deploy_models(dynamic_control_flow=False) +# `torch.onnx.export(dynamo=True)` is expensive (~1.5s/export), so the dynamo matrix is +# trimmed to representatives that cover each distinct input/output container shape plus +# the compile-failure path. Full structural and numeric-type coverage still runs in the +# (cheap) non-dynamo ``test_onnx_export_and_inputs`` below. +_DYNAMO_REPRESENTATIVE_MODELS = { + "TensorModel", # plain single tensor + "ListMultiModel", # list of tensors (arg flattening) + "ListDictModel", # mixed list + dict nesting + "NestedModel", # deeply nested inputs + "DictMultiModel", # dict inputs + "ArgsKwargsModel1", # args + kwargs (success) + "ArgsKwargsModel2", # args + kwargs (compile_fail path) + "TwoOutModel", # multiple outputs + "NestedOutModel", # nested outputs +} +deploy_benchmark_dynamo = { + k: v for k, v in deploy_benchmark_dynamo.items() if k in _DYNAMO_REPRESENTATIVE_MODELS +} + @pytest.mark.parametrize( "model", deploy_benchmark_dynamo.values(), ids=deploy_benchmark_dynamo.keys() ) def test_onnx_dynamo_export(skip_on_windows, model: BaseDeployModel): - # try it for all potential numeric types - for active in range(model.get.num_choices): + # One numeric type is enough here — numeric-type coverage is in test_onnx_export_and_inputs. + for active in range(1): # retrieve args model.get.active = active model.get.set_default_counter() diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 450509d8527..1a7a3495158 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -16,6 +16,7 @@ import json import pytest +import torch from _test_utils.torch.diffusers_models import ( get_tiny_dit, get_tiny_flux, @@ -25,7 +26,9 @@ pytest.importorskip("diffusers") +import modelopt.torch.export.unified_export_hf as unified_export_hf from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.diffusers_utils import generate_diffusion_dummy_inputs from modelopt.torch.export.unified_export_hf import export_hf_checkpoint @@ -54,8 +57,6 @@ def test_export_diffusers_unet_quantized_matches_llm_config(tmp_path, monkeypatc model = get_tiny_unet() export_dir = tmp_path / "export_unet_quant" - import modelopt.torch.export.unified_export_hf as unified_export_hf - monkeypatch.setattr(unified_export_hf, "has_quantized_modules", lambda *_: True) fuse_calls = {"count": 0} @@ -93,10 +94,6 @@ def _process_stub(*_args, **_kwargs): def test_flux2_dummy_inputs_shape(): """Verify Flux2-specific dummy input shapes: 4-col RoPE ids, no pooled_projections, guidance.""" - import torch - - from modelopt.torch.export.diffusers_utils import generate_diffusion_dummy_inputs - model = get_tiny_flux2() cfg = model.config inputs = generate_diffusion_dummy_inputs(model, torch.device("cpu"), torch.float32) diff --git a/tests/unit/torch/nas/test_full_algorithms_vision.py b/tests/unit/torch/nas/test_full_algorithms_vision.py index ea8b7228001..d217b79909f 100644 --- a/tests/unit/torch/nas/test_full_algorithms_vision.py +++ b/tests/unit/torch/nas/test_full_algorithms_vision.py @@ -28,13 +28,20 @@ from modelopt.torch.opt.utils import named_hparams from modelopt.torch.utils import random +# These full end-to-end fastnas/autonas vision-algorithm tests are slow (~5-21s each, +# dominated by search/profile/tracing) and exercise rarely-used paths. Heavier cases are +# marked ``@pytest.mark.manual`` (run with ``--run-manual``); one representative case per +# test stays enabled for sanity, preferring the fastnas variant over autonas where a test +# covers both. Core NAS behavior (convert / search-space / save-restore) is also covered by +# the lighter sibling test_nas.py and test_search_space*.py. + @pytest.mark.parametrize( ("get_model_and_input", "variant"), [ - (get_tiny_mobilenet_and_input, 0), - (get_tiny_resnet_and_input, 1), - (get_tiny_resnet_and_input, 2), + pytest.param(get_tiny_mobilenet_and_input, 0, marks=pytest.mark.manual), # autonas + (get_tiny_resnet_and_input, 1), # fastnas — kept enabled for sanity + pytest.param(get_tiny_resnet_and_input, 2, marks=pytest.mark.manual), # autonas ], ) def test_searched_model_constraints(get_model_and_input, variant): @@ -101,8 +108,15 @@ def run_search(): assert check_constraint(subnet_stat, limits) -@pytest.mark.parametrize("flops_only", [True, False]) -@pytest.mark.parametrize("bounded_latency", [True, False]) +@pytest.mark.parametrize( + ("flops_only", "bounded_latency"), + [ + (False, True), # fast representative kept enabled for sanity + pytest.param(True, True, marks=pytest.mark.manual), + pytest.param(True, False, marks=pytest.mark.manual), + pytest.param(False, False, marks=pytest.mark.manual), + ], +) def test_search_constraints(flops_only: bool, bounded_latency: bool): def _fake_latency(self, model, precomputed=None): coeff_params = 0.00 if flops_only else 0.75 @@ -176,7 +190,11 @@ def _fake_latency(self, model, precomputed=None): @pytest.mark.parametrize( - "get_model_and_input", [get_tiny_resnet_and_input, get_tiny_mobilenet_and_input] + "get_model_and_input", + [ + get_tiny_resnet_and_input, # fast representative kept enabled for sanity + pytest.param(get_tiny_mobilenet_and_input, marks=pytest.mark.manual), + ], ) def test_profile_same_max(get_model_and_input): """checks if the max/original subnet of the profiled model is the same as the original model.""" @@ -226,8 +244,8 @@ def _initialize_test_case_profile(get_model_and_input): @pytest.mark.parametrize( ("get_model_and_input", "mode"), [ - (get_tiny_resnet_and_input, "autonas"), - (get_tiny_mobilenet_and_input, "fastnas"), + pytest.param(get_tiny_resnet_and_input, "autonas", marks=pytest.mark.manual), + (get_tiny_mobilenet_and_input, "fastnas"), # fastnas — kept enabled for sanity ], ) def test_profile_search_space(get_model_and_input, mode): diff --git a/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py b/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py index dcd8f24ab13..0d62d3d6520 100644 --- a/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py +++ b/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py @@ -26,6 +26,7 @@ import torch pytest.importorskip("onnx") +pytest.importorskip("diffusers") from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 5a7fe8cd9a6..cb4a4f2cd2e 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -22,6 +22,9 @@ pytest.importorskip("transformers") +import modelopt.torch.quantization as mtq +from modelopt.torch.export.moe_utils import _export_fused_experts +from modelopt.torch.export.quant_utils import get_quant_config from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins.huggingface import ( @@ -32,6 +35,7 @@ register_fused_experts_on_the_fly, register_sparse_moe_on_the_fly, ) +from modelopt.torch.quantization.utils.core_utils import weight_attr_names # --------------------------------------------------------------------------- # Synthetic fused expert module matching the HF transformers 5.0+ pattern @@ -263,9 +267,6 @@ def _cleanup_registry(mod_type): def test_export_creates_per_expert_submodules(self): """_export_fused_experts should create per-expert submodules with standard naming.""" - import modelopt.torch.quantization as mtq - from modelopt.torch.export.moe_utils import _export_fused_experts - model = _TinyMoEModel() expert_type = type(model.moe.experts) self._cleanup_registry(expert_type) @@ -336,8 +337,6 @@ def test_uncalibrated_expert_gate_up_share_amax(self, monkeypatch): tensor before the deepcopies, so gate's clone and up's clone start with the same amax. """ - from modelopt.torch.export.moe_utils import _export_fused_experts - # Build experts where gate and up have very different magnitudes — # any per-half fallback would clearly produce different amaxes. experts = _SyntheticFusedExperts() @@ -420,8 +419,6 @@ def test_per_block_amax_reshape_for_fused_export(self, monkeypatch): per-projection scales. The fix reshapes to ``(fused_total, blocks_per_row)`` before slicing on dim-0 when ``amax.numel() % fused_total == 0``. """ - from modelopt.torch.export.moe_utils import _export_fused_experts - experts = _SyntheticFusedExperts() expert_type = type(experts) if QuantModuleRegistry.get(expert_type) is None: @@ -600,8 +597,6 @@ def _cleanup_registry(mod_type): def test_calibration_populates_all_expert_quantizers(self): """After PTQ, every input/weight quantizer on the fused-experts module has amax set.""" - import modelopt.torch.quantization as mtq - model = _TinyMoEModel() expert_type = type(model.moe.experts) self._cleanup_registry(expert_type) @@ -664,8 +659,6 @@ def test_max_calibrate_populates_dead_static_nvfp4_expert_quantizers(self): ``_amax=None`` unless static NVFP4 finalization bootstraps them from the per-expert weight slice. """ - import modelopt.torch.quantization as mtq - model = _TinyMoEModel() expert_type = type(model.moe.experts) self._cleanup_registry(expert_type) @@ -767,8 +760,6 @@ def _cleanup_registry(mod_type): def test_weight_attr_names_yields_fused_expert_params(self): """weight_attr_names must yield gate_up_proj / down_proj on fused experts even though their quantizers are a plural ModuleList, not singular.""" - from modelopt.torch.quantization.utils.core_utils import weight_attr_names - model = _TinyMoEModel() expert_type = type(model.moe.experts) self._cleanup_registry(expert_type) @@ -788,9 +779,6 @@ def test_mixed_precision_config_export(self): """Mixed-precision recipe (experts FP8 + dense Linear FP8 per-channel) should show both modules in quantized_layers. Using two distinct formats would trigger MIXED_PRECISION; using same-format still exercises enumeration.""" - import modelopt.torch.quantization as mtq - from modelopt.torch.export.quant_utils import get_quant_config - model = _MixedPrecisionModel() expert_type = type(model.moe.experts) self._cleanup_registry(expert_type) diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index ae638c42ee2..34cf6c7f95d 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -12,6 +12,7 @@ # 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. +import copy import os import warnings from contextlib import nullcontext @@ -204,8 +205,6 @@ def test_quantized_transformers_save_restore(tmp_path, model_cls, quant_config): tiny_llama_dir = create_tiny_llama_dir(tmp_path, dtype=torch.float32) # update config to fit test cases if quant_config == mtq.INT4_AWQ_CFG: - import copy - quant_config = copy.deepcopy(quant_config) for entry in quant_config["quant_cfg"]: if entry["quantizer_name"] == "*weight_quantizer": diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index 8ae973bb471..4378d7bbc9e 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -413,7 +413,8 @@ def _test_data_parallel_auto_quantize(rank, size): def test_data_parallel_auto_quantize(skip_on_windows): - spawn_multiprocess_job(4, _test_data_parallel_auto_quantize, backend="gloo") + # 2 ranks fully exercise the cross-rank sync the test asserts; more just adds spawn overhead. + spawn_multiprocess_job(2, _test_data_parallel_auto_quantize, backend="gloo") def test_estimate_quant_compression(): diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index a39ee55d9d2..9e58185e383 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -30,6 +30,7 @@ layerwise_calibrate, ) from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector class _SimpleMLP(nn.Module): @@ -480,8 +481,6 @@ def forward(self, x): def test_layerwise_calibrate_propagates_inputs_without_replaying_full_model(monkeypatch): - from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector - class _ToyLayer(nn.Module): def __init__(self, scale: float, bias: float): super().__init__() @@ -562,7 +561,6 @@ def _pre_hook(_module, args): def test_layerwise_calibrate_handles_inter_layer_logic(monkeypatch): """Verify that parent-level inter-layer logic (e.g. mask selection) works correctly.""" - from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector class _ToyLayer(nn.Module): def __init__(self, scale: float): diff --git a/tests/unit/torch/quantization/test_calibrator.py b/tests/unit/torch/quantization/test_calibrator.py index 19c86b0b9f0..9f7d77ce6f7 100644 --- a/tests/unit/torch/quantization/test_calibrator.py +++ b/tests/unit/torch/quantization/test_calibrator.py @@ -182,9 +182,11 @@ def test_torch_hist(self): class TestEntropyCalibrator: + # ``num_bins=512`` (down from 2048) keeps the entropy KL-search ~8x cheaper while + # still discarding the single outlier; ``start_bin`` scales with num_bins (128->32). def test_one_tensor(self, verbose): hist_calibrator = calib.HistogramCalibrator( - 8, None, False, num_bins=2048, grow_method="stretch" + 8, None, False, num_bins=512, grow_method="stretch" ) x_2 = torch.rand(11, 7, 3, 3) # uniform in (0,1) @@ -192,7 +194,7 @@ def test_one_tensor(self, verbose): hist_calibrator.collect(x_2) # Don't have a better test metric. One outlier 10 should be discarded by KL-divergence - amax = hist_calibrator.compute_amax("entropy", start_bin=128) + amax = hist_calibrator.compute_amax("entropy", start_bin=32) if verbose: print(f"amax={amax.item():.4f}", end=" ") @@ -201,28 +203,24 @@ def test_one_tensor(self, verbose): def test_unsigned(self, verbose): hist_calibrator = calib.HistogramCalibrator( - 8, None, True, num_bins=2048, grow_method="stretch" + 8, None, True, num_bins=512, grow_method="stretch" ) x_2 = torch.rand(11, 7, 3, 3) # uniform in (0,1) x_2[1, 1, 1, 1] = 10.0 # create outlier hist_calibrator.collect(x_2) - amax = hist_calibrator.compute_amax("entropy", start_bin=128) + amax = hist_calibrator.compute_amax("entropy", start_bin=32) if verbose: print(f"amax={amax.item():.4f}", end=" ") assert amax < 1.1 - @pytest.mark.parametrize("torch_hist", [False, True]) - def test_two_tensor(self, torch_hist, verbose): - hist_calibrator = calib.HistogramCalibrator( - 8, None, False, num_bins=2048, torch_hist=torch_hist - ) - - x_2 = torch.rand(11, 7, 3, 3) # uniform in (0,1) - x_2[1, 1, 1, 1] = 10.0 # create outlier + # torch_hist vs numpy histogram equivalence is covered by + # ``TestHistogramCalibrator::test_torch_hist``, so a single case suffices here. + def test_two_tensor(self, verbose): + hist_calibrator = calib.HistogramCalibrator(8, None, False, num_bins=512) x_2 = torch.rand(11, 7, 3, 3) # uniform in (0,1) x_2[1, 1, 1, 1] = 10.0 # create outlier @@ -231,7 +229,7 @@ def test_two_tensor(self, torch_hist, verbose): hist_calibrator.collect(x_3) # Don't have a better test metric. One outlier 10 should be discarded by KL-divergence - amax = hist_calibrator.compute_amax("entropy", start_bin=128) + amax = hist_calibrator.compute_amax("entropy", start_bin=32) if verbose: print(f"amax={amax.item():.4f}", end=" ") @@ -406,13 +404,15 @@ def test_shape_with_axis(self): == test_module.weight_quantizer.amax.shape ) + # ``num_bins=256`` (down from the 2048 default) makes the MSE bin-scan ~8x cheaper; + # the reference calibrator must use the same num_bins for the atol=0 comparison. @pytest.mark.parametrize("method", ["mse", "percentile"]) def test_per_tensor(self, method): test_lenet = QuantConvLinear() - ref_calibrator = calib.HistogramCalibrator(8, None, False) + ref_calibrator = calib.HistogramCalibrator(8, None, False, num_bins=256) - calib.calibrate_weights(test_lenet, method=method, perchannel=False) + calib.calibrate_weights(test_lenet, method=method, perchannel=False, num_bins=256) ref_calibrator.collect(test_lenet.conv1.weight) ref_amax = ref_calibrator.compute_amax(method) assert torch.allclose(ref_amax, test_lenet.conv1.weight_quantizer.amax, rtol=0, atol=0) @@ -421,9 +421,9 @@ def test_per_tensor(self, method): def test_with_axis(self, method): test_lenet = QuantConvLinear() - ref_calibrator = calib.HistogramCalibrator(8, None, False) + ref_calibrator = calib.HistogramCalibrator(8, None, False, num_bins=256) - calib.calibrate_weights(test_lenet, method=method, perchannel=True) + calib.calibrate_weights(test_lenet, method=method, perchannel=True, num_bins=256) ref_calibrator.collect(test_lenet.conv2.weight[1]) ref_amax = ref_calibrator.compute_amax(method) assert torch.allclose(ref_amax, test_lenet.conv2.weight_quantizer.amax[1], rtol=0, atol=0) diff --git a/tests/unit/torch/quantization/test_config_validation.py b/tests/unit/torch/quantization/test_config_validation.py index ce98f989f51..17ee64a6d8a 100644 --- a/tests/unit/torch/quantization/test_config_validation.py +++ b/tests/unit/torch/quantization/test_config_validation.py @@ -18,6 +18,7 @@ import pytest from pydantic import ValidationError +from modelopt.torch.quantization.algorithms import _match_quantizer_cfg from modelopt.torch.quantization.config import ( FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, FP8_DEFAULT_CFG, @@ -438,8 +439,6 @@ class TestMatchQuantizerCfg: def test_wildcard_matches_bare_name(self): """'*weight_quantizer' matches bare 'weight_quantizer'.""" - from modelopt.torch.quantization.algorithms import _match_quantizer_cfg - quant_cfg = normalize_quant_cfg_list( [{"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8}}] ) @@ -449,8 +448,6 @@ def test_wildcard_matches_bare_name(self): def test_star_matches_any_bare_name(self): """'*' matches any bare quantizer name.""" - from modelopt.torch.quantization.algorithms import _match_quantizer_cfg - quant_cfg = normalize_quant_cfg_list([{"quantizer_name": "*", "enable": False}]) matched, enable = _match_quantizer_cfg(quant_cfg, "weight_quantizer") assert matched is None # enable-only entry has cfg=None @@ -458,8 +455,6 @@ def test_star_matches_any_bare_name(self): def test_path_scoped_pattern_matches_matching_suffix(self): """'*mlp*weight_quantizer' matches bare 'weight_quantizer' (suffix match).""" - from modelopt.torch.quantization.algorithms import _match_quantizer_cfg - quant_cfg = normalize_quant_cfg_list( [{"quantizer_name": "*mlp*weight_quantizer", "cfg": {"num_bits": 4}}] ) @@ -468,8 +463,6 @@ def test_path_scoped_pattern_matches_matching_suffix(self): def test_path_scoped_pattern_does_not_match_different_suffix(self): """'*mlp*weight_quantizer' does NOT match bare 'input_quantizer'.""" - from modelopt.torch.quantization.algorithms import _match_quantizer_cfg - quant_cfg = normalize_quant_cfg_list( [{"quantizer_name": "*mlp*weight_quantizer", "cfg": {"num_bits": 4}}] ) @@ -479,8 +472,6 @@ def test_path_scoped_pattern_does_not_match_different_suffix(self): def test_last_match_wins(self): """Later entries override earlier ones.""" - from modelopt.torch.quantization.algorithms import _match_quantizer_cfg - quant_cfg = normalize_quant_cfg_list( [ {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8}}, @@ -492,8 +483,6 @@ def test_last_match_wins(self): def test_no_match_returns_none(self): """No matching entry returns (None, None).""" - from modelopt.torch.quantization.algorithms import _match_quantizer_cfg - quant_cfg = normalize_quant_cfg_list( [{"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8}}] ) @@ -503,8 +492,6 @@ def test_no_match_returns_none(self): def test_bracket_pattern_matches_correctly(self): """'*[kv]_bmm_quantizer' matches 'k_bmm_quantizer' and 'v_bmm_quantizer'.""" - from modelopt.torch.quantization.algorithms import _match_quantizer_cfg - quant_cfg = normalize_quant_cfg_list( [{"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": (4, 3)}}] ) @@ -521,8 +508,6 @@ def test_path_scoped_does_not_overmatch(self): Regression test: the old rsplit('*') logic would strip to 'weight_quantizer' and overmatch any quantizer ending in 'weight_quantizer', but should not match unrelated names. """ - from modelopt.torch.quantization.algorithms import _match_quantizer_cfg - quant_cfg = normalize_quant_cfg_list( [ {"quantizer_name": "*", "enable": False}, diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_flash_skip_softmax.py b/tests/unit/torch/sparsity/attention_sparsity/test_flash_skip_softmax.py index 0724fa9ac03..15ad65f963f 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_flash_skip_softmax.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_flash_skip_softmax.py @@ -17,6 +17,7 @@ import pytest import torch +import torch.nn.functional as F pytest.importorskip("transformers") @@ -348,8 +349,6 @@ def test_get_threshold_info_static(self): def test_get_sparse_context_patches_softmax(self): """get_sparse_context returns an ExitStack that patches F.softmax.""" - import torch.nn.functional as F - method = FlashSkipSoftmax( { "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, @@ -376,8 +375,6 @@ def test_get_sparse_context_patches_softmax(self): def test_calibration_mode_skips_apply(self): """In calibration mode, sparse_softmax wrapper does not apply mask.""" - import torch.nn.functional as F - method = FlashSkipSoftmax( { "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py index 9519856a875..d9565a233d8 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py @@ -62,12 +62,12 @@ def test_generate_target_lengths_stops_at_minimum(self): class TestRulerDatasetBuilder: """Test RULER dataset generation without requiring real tokenizers.""" - def test_builder_initialization(self): + def test_builder_initialization(self, tiny_tokenizer): """Test that builder initializes correctly.""" builder = RulerDatasetBuilder( samples=12, max_seqlen=2048, # Generates: [2048, 1024] - tokenizer_name_or_path="gpt2", + tokenizer_name_or_path=tiny_tokenizer, seed=42, ) @@ -96,12 +96,12 @@ def test_builder_initialization_invalid_config(self): tokenizer_name_or_path="gpt2", ) - def test_dataset_generation_minimal(self): + def test_dataset_generation_minimal(self, tiny_tokenizer): """Test generating small dataset.""" builder = RulerDatasetBuilder( samples=12, # 6 tasks x 2 lengths = need 12 for 1 per task per length max_seqlen=2048, # Generates: [2048, 1024] - tokenizer_name_or_path="gpt2", + tokenizer_name_or_path=tiny_tokenizer, seed=42, ) @@ -111,12 +111,12 @@ def test_dataset_generation_minimal(self): assert len(dataset) == 12 assert all(isinstance(sample, dict) for sample in dataset) - def test_dataset_structure(self): + def test_dataset_structure(self, tiny_tokenizer): """Test that dataset has correct structure.""" builder = RulerDatasetBuilder( samples=6, # Need at least 6 (1 per task) max_seqlen=1024, # Generates: [1024] - tokenizer_name_or_path="gpt2", + tokenizer_name_or_path=tiny_tokenizer, seed=42, ) @@ -135,12 +135,12 @@ def test_dataset_structure(self): assert isinstance(sample["task"], str) assert sample["length"] > 0 - def test_uneven_sample_distribution(self): + def test_uneven_sample_distribution(self, tiny_tokenizer): """Test that samples are distributed evenly (remainder dropped).""" builder = RulerDatasetBuilder( samples=50, # 50 samples across 4 lengths max_seqlen=8192, # Generates: [8192, 4096, 2048, 1024] - tokenizer_name_or_path="gpt2", + tokenizer_name_or_path=tiny_tokenizer, seed=42, ) @@ -472,8 +472,6 @@ def test_calibrate_both_phases_zero(self): def test_calibrate_with_user_forward_loop(self): """User-provided forward_loop skips RULER dataset building.""" - import numpy as np - model = SimpleAttentionModel(hidden_size=64, num_heads=4) # Sparsify first WITHOUT calibration, so we can call calibrate_sparse_attention # ourselves with a user-supplied forward_loop. diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py index 36975e8c08f..e35ac698e76 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py @@ -18,13 +18,12 @@ GPU-dependent tests (training forward, module forward) are in tests/gpu/. """ +import json import os from copy import deepcopy -from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock -import pytest import torch from _test_utils.torch.transformers_models import ( get_tiny_llama, @@ -36,10 +35,13 @@ import modelopt.torch.speculative as mtsp from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG from modelopt.torch.speculative.plugins.hf_dflash import ( + DFlashAttention, DFlashModule, HFDFlashModel, build_target_layer_ids, ) +from modelopt.torch.speculative.utils import AcceptanceRateValidation +from modelopt.torch.utils.plugins.transformers_dataset import LanguageDataCollator BLOCK_SIZE = 4 NUM_DRAFT_LAYERS = 2 @@ -190,8 +192,6 @@ def test_sliding_window_from_config(self): """Test DFlashAttention reads sliding_window from config.layer_types.""" from transformers import PretrainedConfig - from modelopt.torch.speculative.plugins.hf_dflash import DFlashAttention - config = PretrainedConfig( hidden_size=64, num_attention_heads=4, @@ -213,8 +213,6 @@ def test_no_sliding_window_without_config(self): """Test DFlashAttention defaults to no sliding window.""" from transformers import PretrainedConfig - from modelopt.torch.speculative.plugins.hf_dflash import DFlashAttention - config = PretrainedConfig( hidden_size=64, num_attention_heads=4, @@ -234,8 +232,6 @@ class TestValidateOnline: def test_all_accepted(self): """When all draft tokens match posterior, AR = 1 + steps.""" - from modelopt.torch.speculative.utils import AcceptanceRateValidation - validator = AcceptanceRateValidation.__new__(AcceptanceRateValidation) validator.check_data_consistency_across_ranks = lambda x: x @@ -278,8 +274,6 @@ def mock_lm_head(hidden): def test_all_rejected(self): """When no draft tokens match, AR = 1 (base token only + correction).""" - from modelopt.torch.speculative.utils import AcceptanceRateValidation - validator = AcceptanceRateValidation.__new__(AcceptanceRateValidation) validator.check_data_consistency_across_ranks = lambda x: x @@ -353,8 +347,6 @@ def test_export_state_dict_has_no_prefix(self, tmp_path): def test_export_config_fields(self, tmp_path): """Exported config.json should have required DFlash fields.""" - import json - model = get_tiny_llama(num_hidden_layers=4) config = _get_dflash_config() mtsp.convert(model, [("dflash", config)]) @@ -399,37 +391,23 @@ def test_export_tensor_count(self, tmp_path): class TestEnsureGenerationTags: - """Test _ensure_generation_tags with a real tokenizer (Qwen3-0.6B from HF).""" - - @pytest.fixture - def qwen3_tokenizer(self): - from transformers import AutoTokenizer + """Test _ensure_generation_tags masking with the local tiny tokenizer. - return AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") - - @pytest.fixture - def qwen3_chat_template(self): - template_path = ( - Path(__file__).parents[5] - / "tools/launcher/examples/Qwen/Qwen3-8B/chat_template_train.jinja" - ) - return template_path.read_text() + Masking is driven by the injected ``{% generation %}`` blocks and + ``return_assistant_tokens_mask``, so the local tiny tokenizer suffices — no + HF Hub download needed. + """ - def test_chatml_think_template_produces_assistant_mask( - self, qwen3_tokenizer, qwen3_chat_template - ): + def test_chatml_think_template_produces_assistant_mask(self, tiny_tokenizer): """Verify generation-tagged chat template produces correct assistant masks.""" - from modelopt.torch.utils.plugins.transformers_dataset import LanguageDataCollator - collator = LanguageDataCollator( - tokenizer=qwen3_tokenizer, + tokenizer=tiny_tokenizer, train_len=128, return_labels=True, answer_only_loss=True, - chat_template=qwen3_chat_template, ) - # Verify template was replaced with generation-tagged version + # The tiny tokenizer ships a generation-tagged chat template assert "generation" in collator.tokenizer.chat_template # Tokenize a sample conversation @@ -454,19 +432,16 @@ def test_chatml_think_template_produces_assistant_mask( # Decode the non-masked positions to verify they're assistant content non_masked = input_ids[labels != -100] - decoded = qwen3_tokenizer.decode(non_masked) + decoded = tiny_tokenizer.decode(non_masked) assert "The answer is 4" in decoded - def test_multi_turn_masks_only_assistant(self, qwen3_tokenizer, qwen3_chat_template): + def test_multi_turn_masks_only_assistant(self, tiny_tokenizer): """Verify multi-turn: only assistant turns are unmasked.""" - from modelopt.torch.utils.plugins.transformers_dataset import LanguageDataCollator - collator = LanguageDataCollator( - tokenizer=qwen3_tokenizer, + tokenizer=tiny_tokenizer, train_len=256, return_labels=True, answer_only_loss=True, - chat_template=qwen3_chat_template, ) samples = [ @@ -485,7 +460,7 @@ def test_multi_turn_masks_only_assistant(self, qwen3_tokenizer, qwen3_chat_templ input_ids = result["input_ids"] non_masked = input_ids[labels != -100] - decoded = qwen3_tokenizer.decode(non_masked) + decoded = tiny_tokenizer.decode(non_masked) # Both assistant responses should appear in unmasked tokens assert "Hi there" in decoded assert "I am fine" in decoded diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative.py b/tests/unit/torch/speculative/plugins/test_hf_speculative.py index f5a14a08976..4a75ffdd775 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative.py +++ b/tests/unit/torch/speculative/plugins/test_hf_speculative.py @@ -37,8 +37,15 @@ def test_eagle_model_convert_save_and_restore(tmp_path, eagle_config): { "draft_vocab_size": model_ref.config.vocab_size, "hidden_size": model_ref.config.hidden_size, + # Shrink the eagle module so convert + save/restore serialization stay cheap. + "intermediate_size": 32, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "head_dim": 2, } ) + # torch.compile(max-autotune) is a GPU optimization; on CPU it only adds compile time. + config["eagle_use_torch_compile"] = False mtsp.convert(model_ref, mode=[("eagle", config)]) assert isinstance(model_ref, mtsp.plugins.HFEagleModel) diff --git a/tests/unit/torch/speculative/plugins/test_hf_speculative_lora.py b/tests/unit/torch/speculative/plugins/test_hf_speculative_lora.py index 1b03034f525..c4ab2a0266e 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_speculative_lora.py +++ b/tests/unit/torch/speculative/plugins/test_hf_speculative_lora.py @@ -38,6 +38,9 @@ EAGLE_LORA_CONFIG = { "eagle_architecture_config": {**default_eagle_config, **TINY_EAGLE_CFG}, + # torch.compile(max-autotune) is a GPU training-throughput optimization; on CPU it + # adds ~40s of compile time with no benefit, so disable it for these unit tests. + "eagle_use_torch_compile": False, "eagle_base_lora": True, "eagle_base_lora_rank": 4, "eagle_base_lora_alpha": 8.0, @@ -46,7 +49,7 @@ } -@pytest.fixture +@pytest.fixture(scope="module") def lora_eagle_model(): model = get_tiny_llama(num_hidden_layers=4) mtsp.convert(model, mode=[("eagle", deepcopy(EAGLE_LORA_CONFIG))]) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 6508dc67c89..bfcca2975f2 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -18,13 +18,11 @@ import pytest import torch -from huggingface_hub import get_token from torch.utils.data import DataLoader from modelopt.torch.utils import dataset_utils from modelopt.torch.utils.dataset_utils import ( DATASET_COMBOS, - SUPPORTED_DATASET_CONFIG, _disable_use_cache, _forward_loop, _pack_documents_into_rows, @@ -287,22 +285,9 @@ def fake_forward(x): assert result == 4 -@pytest.mark.parametrize("test_local_path", [True, False]) -def test_get_dataset_samples_with_unsupported_minipile_dataset(tmp_path, test_local_path): - pytest.importorskip("datasets") - pytest.importorskip("huggingface_hub") - - from huggingface_hub import snapshot_download - - dataset_name = "nanotron/minipile_100_samples" - if test_local_path: - local_dir = str(tmp_path / dataset_name) - snapshot_download( - repo_id=dataset_name, - repo_type="dataset", - local_dir=local_dir, - ) - dataset_name = local_dir +def test_get_dataset_samples_with_unsupported_dataset(make_toy_hf_dataset): + """A dataset not in ``SUPPORTED_DATASET_CONFIG`` loads via the auto-detect path.""" + dataset_name = make_toy_hf_dataset() # basename never matches a registered key samples = get_dataset_samples(dataset_name, num_samples=5) @@ -576,22 +561,10 @@ def test_pack_documents_into_rows(): # --------------------------------------------------------------------------- -@pytest.fixture -def pad_tokenizer(): - """Real tiny HF tokenizer (vocab=128) shared with other test modules. - - Skips the test if ``transformers`` isn't installed. - """ - pytest.importorskip("transformers") - from _test_utils.torch.transformers_models import get_tiny_tokenizer - - return get_tiny_tokenizer() - - class TestGetDatasetDataloaderBlending: """``get_dataset_dataloader`` accepts a list of sources and concatenates them.""" - def test_single_jsonl(self, tmp_path, pad_tokenizer): + def test_single_jsonl(self, tmp_path, tiny_tokenizer): pytest.importorskip("datasets") path = _write_jsonl( tmp_path / "single.jsonl", @@ -599,7 +572,7 @@ def test_single_jsonl(self, tmp_path, pad_tokenizer): ) loader = get_dataset_dataloader( dataset_name=path, - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, batch_size=2, num_samples=4, max_sample_length=16, @@ -609,7 +582,7 @@ def test_single_jsonl(self, tmp_path, pad_tokenizer): assert batches[0]["input_ids"].shape[0] == 2 assert "attention_mask" in batches[0] - def test_list_of_jsonl_blends(self, tmp_path, pad_tokenizer): + def test_list_of_jsonl_blends(self, tmp_path, tiny_tokenizer): """Two local JSONL files concatenated into a single dataloader.""" pytest.importorskip("datasets") a = _write_jsonl(tmp_path / "a.jsonl", [{"text": f"a{i}"} for i in range(3)]) @@ -617,7 +590,7 @@ def test_list_of_jsonl_blends(self, tmp_path, pad_tokenizer): loader = get_dataset_dataloader( dataset_name=[a, b], - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, batch_size=5, num_samples=[3, 2], max_sample_length=16, @@ -626,7 +599,7 @@ def test_list_of_jsonl_blends(self, tmp_path, pad_tokenizer): assert len(batches) == 1 assert batches[0]["input_ids"].shape[0] == 5 - def test_mixed_formats_blended(self, tmp_path, pad_tokenizer): + def test_mixed_formats_blended(self, tmp_path, tiny_tokenizer): """Mixing a text-column JSONL with a prompt/completion JSONL — both should flow.""" pytest.importorskip("datasets") plain = _write_jsonl(tmp_path / "plain.jsonl", [{"text": "hello"}]) @@ -634,7 +607,7 @@ def test_mixed_formats_blended(self, tmp_path, pad_tokenizer): loader = get_dataset_dataloader( dataset_name=[plain, pc], - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, batch_size=2, num_samples=[1, 1], max_sample_length=16, @@ -643,7 +616,7 @@ def test_mixed_formats_blended(self, tmp_path, pad_tokenizer): assert len(batches) == 1 assert batches[0]["input_ids"].shape[0] == 2 - def test_length_mismatch_raises(self, tmp_path, pad_tokenizer): + def test_length_mismatch_raises(self, tmp_path, tiny_tokenizer): """``dataset_name`` and ``num_samples`` lists must align.""" pytest.importorskip("datasets") a = _write_jsonl(tmp_path / "a.jsonl", [{"text": "x"}]) @@ -651,13 +624,13 @@ def test_length_mismatch_raises(self, tmp_path, pad_tokenizer): with pytest.raises(AssertionError, match="same length"): get_dataset_dataloader( dataset_name=[a, b], - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, num_samples=[1], max_sample_length=16, ) -def test_multi_source_pack_shuffles_to_avoid_dominance(monkeypatch, pad_tokenizer): +def test_multi_source_pack_shuffles_to_avoid_dominance(monkeypatch, tiny_tokenizer): """With ``pack=True`` and 2+ sources, samples are shuffled so a long-doc source can't silently exhaust the row budget and drop the other sources. @@ -676,7 +649,7 @@ def _fake(name, num_sample, **_kwargs): loader = get_dataset_dataloader( dataset_name=["src_a", "src_b"], - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, batch_size=4, num_samples=[4, 4], max_sample_length=64, @@ -687,8 +660,8 @@ def _fake(name, num_sample, **_kwargs): all_ids = torch.cat([b["input_ids"] for b in batches], dim=0) assert all_ids.shape[1] == 64 # Tokenize the source tags so we can check both sources appear in the packed rows - src_a_id = pad_tokenizer("src_a", add_special_tokens=False).input_ids[0] - src_b_id = pad_tokenizer("src_b", add_special_tokens=False).input_ids[0] + src_a_id = tiny_tokenizer("src_a", add_special_tokens=False).input_ids[0] + src_b_id = tiny_tokenizer("src_b", add_special_tokens=False).input_ids[0] flat = all_ids.flatten().tolist() assert src_a_id in flat, "source A tokens missing from packed rows" assert src_b_id in flat, ( @@ -715,11 +688,11 @@ def _fake(name, num_sample, **_kwargs): monkeypatch.setattr(dataset_utils, "get_dataset_samples", _fake) return calls - def test_combo_expands_evenly(self, monkeypatch, pad_tokenizer): + def test_combo_expands_evenly(self, monkeypatch, tiny_tokenizer): calls = self._record_calls(monkeypatch) get_dataset_dataloader( dataset_name="cnn_nemotron_v2_mix", - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, num_samples=8, batch_size=1, max_sample_length=16, @@ -727,11 +700,11 @@ def test_combo_expands_evenly(self, monkeypatch, pad_tokenizer): members = DATASET_COMBOS["cnn_nemotron_v2_mix"] assert calls == [(members[0], 4), (members[1], 4)] - def test_combo_remainder_distributed_to_earlier_members(self, monkeypatch, pad_tokenizer): + def test_combo_remainder_distributed_to_earlier_members(self, monkeypatch, tiny_tokenizer): calls = self._record_calls(monkeypatch) get_dataset_dataloader( dataset_name="nemotron-post-training-v3", - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, num_samples=10, batch_size=1, max_sample_length=16, @@ -741,11 +714,11 @@ def test_combo_remainder_distributed_to_earlier_members(self, monkeypatch, pad_t expected_counts = [2, 2, 2, 1, 1, 1, 1] assert calls == list(zip(members, expected_counts)) - def test_plain_and_combo_compose(self, monkeypatch, pad_tokenizer): + def test_plain_and_combo_compose(self, monkeypatch, tiny_tokenizer): calls = self._record_calls(monkeypatch) get_dataset_dataloader( dataset_name=["cnn_dailymail", "nemotron-post-training-v3"], - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, num_samples=[3, 7], batch_size=1, max_sample_length=16, @@ -753,12 +726,12 @@ def test_plain_and_combo_compose(self, monkeypatch, pad_tokenizer): members = DATASET_COMBOS["nemotron-post-training-v3"] assert calls == [("cnn_dailymail", 3)] + [(m, 1) for m in members] - def test_combo_overlapping_with_member_raises(self, monkeypatch, pad_tokenizer): + def test_combo_overlapping_with_member_raises(self, monkeypatch, tiny_tokenizer): self._record_calls(monkeypatch) with pytest.raises(ValueError, match="combo 'cnn_nemotron_v2_mix'"): get_dataset_dataloader( dataset_name=["cnn_dailymail", "cnn_nemotron_v2_mix"], - tokenizer=pad_tokenizer, + tokenizer=tiny_tokenizer, num_samples=[2, 4], batch_size=1, max_sample_length=16, @@ -770,15 +743,32 @@ def test_get_dataset_samples_rejects_combo_name(self): # --------------------------------------------------------------------------- -# Live HF dataset round-trips. ``hf-internal-testing/dataset_with_data_files`` -# is a 10-row x {train,test} fixture maintained by HF for their own CI — tiny -# enough to download in a unit test and stable across releases. +# Arbitrary-dataset round-trips. We build a tiny on-disk dataset directory +# (``train``/``test`` parquet with a ``text`` column) and load it through the +# same non-registered ``load_dataset(path=...)`` branch a downloaded HF dataset +# would hit — keeping these unit tests fully offline rather than fetching +# ``hf-internal-testing/*`` fixtures from the Hub. # --------------------------------------------------------------------------- -_HF_TINY = "hf-internal-testing/dataset_with_data_files" # train, test splits, ``text`` col +@pytest.fixture(scope="module") +def make_toy_hf_dataset(tmp_path_factory): + """Factory returning a local dataset-directory path with ``train``/``test`` splits.""" + pytest.importorskip("datasets") + from datasets import Dataset + + def _make(prefix: str = "toy", num_rows: int = 10) -> str: + d = tmp_path_factory.mktemp(prefix) + for split in ("train", "test"): + Dataset.from_dict( + {"text": [f"{prefix} {split} sample {i}" for i in range(num_rows)]} + ).to_parquet(str(d / f"{split}.parquet")) + return str(d) + + return _make -def _hf_dump_to_jsonl(name: str, split: str, path) -> str: + +def _dump_dataset_to_jsonl(name: str, split: str, path) -> str: from datasets import load_dataset ds = load_dataset(name, split=split) @@ -786,49 +776,47 @@ def _hf_dump_to_jsonl(name: str, split: str, path) -> str: return str(path) -@pytest.mark.integration -class TestHfTinyDataset: - """End-to-end coverage with a real (tiny) HF dataset.""" +class TestLocalDatasetDirRoundTrips: + """End-to-end coverage of the auto-detected (non-registered) dataset path.""" - def test_load_single_split_directly(self): - pytest.importorskip("datasets") - samples = get_dataset_samples(_HF_TINY, num_samples=4, split="train") + def test_load_single_split_directly(self, make_toy_hf_dataset): + dataset = make_toy_hf_dataset() + samples = get_dataset_samples(dataset, num_samples=4, split="train") assert len(samples) == 4 assert all(isinstance(s, str) and s for s in samples) - def test_load_multiple_splits_directly(self): + def test_load_multiple_splits_directly(self, make_toy_hf_dataset): """``split=["train", "test"]`` divides ``num_samples`` across both splits.""" - pytest.importorskip("datasets") - samples = get_dataset_samples(_HF_TINY, num_samples=6, split=["train", "test"]) + dataset = make_toy_hf_dataset() + samples = get_dataset_samples(dataset, num_samples=6, split=["train", "test"]) assert len(samples) == 6 # Default per-split is num_samples // n + remainder; for 6/2 → 3 from each. # We can't assert exact origin without re-reading, but both splits should # contribute, which we'll confirm by comparing against direct loads below. - train_only = set(get_dataset_samples(_HF_TINY, num_samples=10, split="train")) - test_only = set(get_dataset_samples(_HF_TINY, num_samples=10, split="test")) + train_only = set(get_dataset_samples(dataset, num_samples=10, split="train")) + test_only = set(get_dataset_samples(dataset, num_samples=10, split="test")) assert any(s in train_only for s in samples) assert any(s in test_only for s in samples) - def test_default_split_is_train(self): - pytest.importorskip("datasets") - default_samples = get_dataset_samples(_HF_TINY, num_samples=4) - train_samples = get_dataset_samples(_HF_TINY, num_samples=4, split="train") + def test_default_split_is_train(self, make_toy_hf_dataset): + dataset = make_toy_hf_dataset() + default_samples = get_dataset_samples(dataset, num_samples=4) + train_samples = get_dataset_samples(dataset, num_samples=4, split="train") assert default_samples == train_samples - def test_download_to_jsonl_then_load(self, tmp_path): - """Dump the HF dataset to JSONL, then reload it via the local-jsonl path.""" - pytest.importorskip("datasets") - jsonl_path = _hf_dump_to_jsonl(_HF_TINY, "train", tmp_path / "train.jsonl") + def test_download_to_jsonl_then_load(self, tmp_path, make_toy_hf_dataset): + """Dump the dataset to JSONL, then reload it via the local-jsonl path.""" + dataset = make_toy_hf_dataset() + jsonl_path = _dump_dataset_to_jsonl(dataset, "train", tmp_path / "train.jsonl") from_jsonl = get_dataset_samples(jsonl_path, num_samples=10) - from_hf = get_dataset_samples(_HF_TINY, num_samples=10, split="train") - assert from_jsonl == from_hf + from_dir = get_dataset_samples(dataset, num_samples=10, split="train") + assert from_jsonl == from_dir - def test_dataloader_blending_two_hf_datasets(self, pad_tokenizer): - """Two HF datasets concatenated via ``get_dataset_dataloader``.""" - pytest.importorskip("datasets") + def test_dataloader_blending_two_datasets(self, tiny_tokenizer, make_toy_hf_dataset): + """Two datasets concatenated via ``get_dataset_dataloader``.""" loader = get_dataset_dataloader( - dataset_name=[_HF_TINY, "hf-internal-testing/multi_dir_dataset"], - tokenizer=pad_tokenizer, + dataset_name=[make_toy_hf_dataset("a"), make_toy_hf_dataset("b")], + tokenizer=tiny_tokenizer, batch_size=4, num_samples=[3, 1], max_sample_length=16, @@ -836,68 +824,17 @@ def test_dataloader_blending_two_hf_datasets(self, pad_tokenizer): batches = list(loader) assert sum(b["input_ids"].shape[0] for b in batches) == 4 - def test_dataloader_mixing_hf_and_local_jsonl(self, tmp_path, pad_tokenizer): - """Live HF dataset blended with a local synthetic JSONL file.""" - pytest.importorskip("datasets") + def test_dataloader_mixing_dir_and_local_jsonl( + self, tmp_path, tiny_tokenizer, make_toy_hf_dataset + ): + """Dataset directory blended with a local synthetic JSONL file.""" local = _write_jsonl(tmp_path / "local.jsonl", [{"text": f"local {i}"} for i in range(2)]) loader = get_dataset_dataloader( - dataset_name=[_HF_TINY, local], - tokenizer=pad_tokenizer, + dataset_name=[make_toy_hf_dataset(), local], + tokenizer=tiny_tokenizer, batch_size=5, num_samples=[3, 2], max_sample_length=16, ) batches = list(loader) assert sum(b["input_ids"].shape[0] for b in batches) == 5 - - -_NEW_NEMOTRON_KEYS = [ - "nemotron-sft-instruction-following-chat-v2", - "nemotron-science-v1", - "nemotron-competitive-programming-v1", - "nemotron-sft-agentic-v2", - "nemotron-math-v2", - "nemotron-sft-swe-v2", - "nemotron-sft-multilingual-v1", -] - - -@pytest.mark.parametrize("dataset_key", _NEW_NEMOTRON_KEYS) -def test_new_nemotron_registry_shape(dataset_key): - """Always-on shape check on the 7 newly registered nvidia/Nemotron-* entries. - - Complements the gated smoke test below — catches typos in dataset paths or - split names even when the runner has no HF credentials. - """ - assert dataset_key in SUPPORTED_DATASET_CONFIG - entry = SUPPORTED_DATASET_CONFIG[dataset_key] - config = entry["config"] - assert config["path"].startswith("nvidia/Nemotron-") - splits = config["split"] - assert isinstance(splits, list) and splits - assert all(isinstance(s, str) and s for s in splits) - assert len(set(splits)) == len(splits) - assert callable(entry["preprocess"]) - assert entry["chat_key"] == "messages" - - -@pytest.mark.integration -@pytest.mark.parametrize("dataset_key", _NEW_NEMOTRON_KEYS) -def test_get_dataset_samples_new_nemotron(dataset_key): - """Smoke-test the 7 newly registered nvidia/Nemotron-* calibration datasets. - - Skipped when no HF token is available because these datasets live behind the HF Hub. - ``huggingface_hub.get_token()`` covers both the ``HF_TOKEN`` env var and tokens - cached by ``hf auth login``. - """ - pytest.importorskip("datasets") - if not get_token(): - pytest.skip( - "No HF token (env HF_TOKEN or `hf auth login`); skipping gated Nemotron smoke test" - ) - - samples = get_dataset_samples(dataset_key, num_samples=2) - - assert isinstance(samples, list) - assert len(samples) == 2 - assert all(isinstance(s, str) and len(s) > 0 for s in samples)