[shengliangxu]

What does this PR do?

Type of change: CI (test runtime and reliability, tests only)

Less data and fewer Hub round trips in four GPU-lane tests:

• Six preprocess_data tests use the tiny local tokenizer (tests/gpu_megatron). They tokenized with gpt2 from the Hub, and megatron_preprocess_data loads the tokenizer once for the vocabulary size and again in every worker, so each test made several round trips; none depends on gpt2's vocabulary. They use the checked-in byte-level tiny_tokenizer_path fixture that the blend test already uses. Three of the six read only local data, so the tokenizer download was their only network access (the jsonl max_tokens test is among the lane's 120 s timeouts, 3 failures). Measured on a CPU host with a cold Hub cache, the 11 tests of the module take 42.8 and 44.0 s before and 35.0 and 36.1 s after (two runs each; the nine tests of the earlier base took 38.5 s and 24.4 s, and 37 s and 23 s in CI). The Qwen3 tests keep their real tokenizer (chat template and EOS handling).
• The puzzletron test data shrinks, and a duplicate NAS test is folded into another. save_dummy_dataset wrote 2500 copies of one conversation into each split (1738 blocks of 2048 tokens), while the pruning scores read 100 blocks and the later steps 2 each; 400 conversations (278 blocks) take the three dataloader builds plus the save from 2.8 s to 0.36 s per test, and the validation loss moved by at most 0.02 on random-init toy Llamas (the scoring tolerance is 0.15). calc_subblock_stats ran 12 passes where only the entry for batch_size 96 is read: batch_sizes: [96] runs 4. test_nas_convert_ffn_pruning ran the same mtn.convert on the same toy Llama and config as test_nas_search (21.6 s with its own 2-rank spawn); its two assertions move into test_nas_search.
• The unified HF export test calibrates on local text, with no previews or batch probing. Each of its 11 params starts a fresh hf_ptq process (22 s per param). It now reads a local jsonl of 64 random-word texts (new helper write_calib_jsonl, shared through the session-wide calib_jsonl fixture in tests/gpu/conftest.py, which fails the session if a test writes into its directory) instead of cnn_dailymail from the Hub, passes --batch_size 64 (the whole calibration set in one batch instead of probing for the largest batch) and --skip_generate. The test only checks the structure of the exported safetensors.
• The GPTQ end-to-end test reads 8 texts from the same local jsonl instead of streaming cnn_dailymail for each of its three params (and a test that could fail on a Hub hiccup).

Time improvement (measured in CI)

Time of the changed tests; before = the runs of #2846-#2849, after = #2850, #2852, #2854 and #2855 (median). #2849 and #2850 ran on runners 1.26x and 1.29x slower for every test, so the change after speed scaling is in brackets.

lane	tests	before	after	change
gpu	tests/gpu/torch/puzzletron (18 tests, 17 after the folded NAS test)	290 s	202 s	-88 s (-70 s, -26%)
gpu	unified HF export, 11 params	228 s	153 s	-75 s (-67 s, -30%)
gpu	GPTQ end-to-end, 3 params	7.2 s	1.1 s	-6 s
gpu_megatron	preprocess_data module	37 s	23 s	-14 s (-10 s, -27%)

So about 170 s (144 s after scaling) off the gpu job and 14 s off the gpu_megatron job. This PR's own run was on a slow runner (puzzletron 310 s); the three later runs agree with each other (199-203 s). My first estimates were 50 s, 38 s and 6 s for the first three rows.

Method: every PR of the series has its own CI run, so "before" is the run of the PR(s) below that lack the change and "after" the runs that carry it. Times come from the job logs (--durations, or the gap between consecutive result lines). Where runner speed matters, the tests that no PR of the series touches scale each run, and the raw and the scaled medians are both shown. One test varies by 15-30% from run to run on its own, so smaller differences are not evidence.

Usage

Nothing user-facing.

Testing

• preprocess_data: ran the module on CPU next to a copy of the tiny_tokenizer_path fixture with the Hub reachable: 11 passed before and after (main has dropped the HF xfail marks), with identical test ids.
• calib_jsonl: with the real fixture and the repo's get_dataset_dataloader on a CPU host, the 8 samples read from the shared 64-text file equal those of the old 8-text file, and two reads leave the directory unmodified (the fixture's teardown check passes). --setup-plan shows the unified export test using the conftest fixture; the GPTQ and export tests themselves need a GPU.
• puzzletron: timed the dataloader builds, the subblock-stats passes and the loss sensitivity on CPU with the repo's own helpers; collected tests/gpu/torch/puzzletron (16 tests, one fewer than before). The Nemotron-3 toy keeps its 128 experts, because shrinking it changes EXPECTED_TEACHER_MEMORY_MIB and EXPECTED_TEACHER_NUM_PARAMS, which need a GPU runner with mamba_ssm to regenerate.
• unified export: no GPU here, so I ran the test on a CPU host through a shim that lets hf_ptq pass its GPU guard and passes --device cpu. 8 of the 11 params pass with the new command line, scale-equality assertions included (fp8, nvfp4, int4_awq, w4a8_awq_beta, int8_weight_only, w4a16_nvfp4, the Qwen3 MoE nvfp4 and the GPT-OSS fp8); the other three need CUDA or Triton. GPTQ: the flow completes on the jsonl for the FP8 and INT4 blockwise configs (NVFP4 needs CUDA).
• Not run: the GPU tests themselves; this PR's gpu and gpu_megatron jobs are the check. Pre-commit passes.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: N/A (CI and tests)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: no new dependency
• Did you write any new necessary tests?: ✅ one small helper, write_calib_jsonl (tests/_test_utils/torch/calib_data.py), and one session fixture, calib_jsonl (tests/gpu/conftest.py), used by two tests; no new test functions
• Did you update Changelog?: N/A (CI and tests)
• Did you get Claude approval on this PR?: N/A (optional, not run)

Additional Information

Slice 6 of 11 of the GPU-CI speed-up work, split out of the reference draft #2730 (not for merge; it holds the whole series with its measurements). Merge in order; each PR stands on its own and passes CI without the later ones.

1. [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist #2845 [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist
2. [OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI
3. [OMNIML-6179] [3/11] Run manual-marked tests on nightly and on-demand GPU runs #2847 [OMNIML-6179] [3/11] Run manual-marked tests on nightly and on-demand GPU runs
4. [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time #2848 [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time
5. [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models #2849 [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models
6. [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data #2850 [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data (this PR)
7. [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests #2851 [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests
8. [OMNIML-6180] [8/11] Trim the DFlash regression test #2852 [OMNIML-6180] [8/11] Trim the DFlash regression test
9. [OMNIML-6180] [9/11] Trim the torch_onnx example lane #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane
10. [OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies #2854 [OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies
11. [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules #2855 [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules

Summary by CodeRabbit

• Tests
  • Updated model export and quantization checks to use locally generated calibration data instead of a public dataset.
  • Standardized preprocessing tests on a small local tokenizer.
  • Restricted model analysis test configurations to a batch size of 96.
  • Reduced the size of sample datasets used by pruning-related checks.
  • Strengthened conversion verification by checking generated pruning outputs.
  • Removed the FFN pruning conversion test while retaining attention-pruning coverage.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 5af47c9f-7c1b-4c6b-99dc-f4e2aa3e6be4

📥 Commits

Reviewing files that changed from the base of the PR and between dba99af and bd0525c.

📒 Files selected for processing: 2

• tests/gpu/conftest.py
• tests/gpu/torch/quantization/test_gptq.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 4 remain after this review.

📝 Walkthrough

The changes add deterministic local calibration data, adjust Puzzletron workloads and NAS assertions, and replace hard-coded tokenizer and dataset identifiers in test fixtures.

Changes: Local calibration data

Layer / File(s)	Summary
Generate and use local calibration data tests/_test_utils/torch/calib_data.py, tests/gpu/conftest.py, tests/gpu/torch/export/..., tests/gpu/torch/quantization/test_gptq.py	A deterministic JSONL calibration-data helper and fixture are added. The export and GPTQ tests use the generated file instead of cnn_dailymail. The export test sets batch size to 64 and adds --skip_generate.

Puzzletron test workloads

Layer / File(s)	Summary
Update Puzzletron workload and NAS tests tests/_test_utils/torch/puzzletron/utils.py, tests/gpu/torch/puzzletron/resources/configs/..., tests/gpu/torch/puzzletron/test_nas_convert.py, tests/gpu/torch/puzzletron/test_nas_search.py	The dummy dataset repetition count changes from 2,500 to 400. Listed calc_subblock_stats.batch_sizes values are restricted to [96]. The FFN-pruning conversion test is removed, and the NAS search test checks conversion outputs.

Megatron preprocessing tests

Layer / File(s)	Summary
Use the tiny tokenizer fixture tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py	JSONL, Hugging Face, cached-resume, gzip, and streaming-warning tests use tiny_tokenizer_path instead of "gpt2".

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~12 minutes

Change: Other

Suggested reviewers: kevalmorabia97

Merge Risk: ⚪ Minimal · up to bd052

No concrete merge-blocking issue is established; the changes are mergeable with normal CI validation.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 35.29% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 17 functions across 7 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely summarizes the main changes: using local tokenizers and calibration text in GPU tests and reducing Puzzletron test data.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The pull request changes only test helpers, GPU test files, and YAML configurations. It adds no Python changes under modelopt/ or examples/, and it changes no pyproject.toml or `requiremen…

Full details: Security Anti-Patterns

Explanation

PASS. The pull request changes only test helpers, GPU test files, and YAML configurations. It adds no Python changes under modelopt/ or examples/, and it changes no pyproject.toml or requirements*.txt files. Therefore, none of the specified security anti-patterns is introduced in the reviewed scope.

✨ Finishing Touches: 💡 2

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[shengliangxu]

/ok to test 2d82a59

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 76.90%. Comparing base (2d43ca2) to head (bd0525c).

Additional details and impacted files

@@                       Coverage Diff                       @@
##           shengliangx/ci-5-onnx-tests    #2850      +/-   ##
===============================================================
- Coverage                        76.90%   76.90%   -0.01%     
===============================================================
  Files                              658      658              
  Lines                            74029    74029              
===============================================================
- Hits                             56933    56932       -1     
- Misses                           17096    17097       +1     

Flag	Coverage Δ
unit	60.54% <ø> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test 7383ca1

[kevalmorabia97]

can you make this session wide fixture in tests/{gpu,gpu_megatron}/conftest.py and reuse in all tests since we dont expect calibration data to be modified by a test (can make this directory read only to restrict any tests modifying it)

[shengliangxu]

Done in 546bff3: calib_jsonl is a session-wide fixture in tests/gpu/conftest.py (the 64 texts are written once). I used assert_unmodified_tree, the check the shared tiny-tokenizer fixture already uses, instead of chmod: the CI containers run as root, so a read-only directory would not stop a write (the helper's docstring has the details); a test that writes into it now fails at teardown. The unified HF export test and the GPTQ end-to-end test both use it. The GPTQ test still asks the dataloader for 8 samples, and the shared file starts with the same 8 texts as before.

I did not add a copy to tests/gpu_megatron/conftest.py: none of its tests calibrates on text (they use model.get_input()), so it would be unused. If one needs it later, the fixture can move up to tests/conftest.py.