[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). Slices 1, 2, 5, 8, 9 and 10 (#2845, #2846, #2849, #2852, #2853, #2854) have merged. The remaining slices (#2847, #2848, #2850, #2851, #2855) are independent PRs on main: each applies to main on its own, and merging all of them onto main in different orders gives the same tree. The times above are from the series' CI runs before 2026-10-21, when all slices formed one linear stack; a PR number in a method note names that PR's run from then.

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 (merged)
2. [OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI (merged)
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 (merged)
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 (merged)
9. [OMNIML-6180] [9/11] Trim the torch_onnx example lane #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane (merged)
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 (merged)
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.
  • Standardized preprocessing tests on a small local tokenizer.
  • Restricted model analysis configurations to a batch size of 96.
  • Reduced sample dataset sizes for 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]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 503eff3
• Latest commit, not reviewed: 55e1c27

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

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: f28ce766-7767-4452-94ce-973bc6d0b3d0

📥 Commits

Reviewing files that changed from the base of the PR and between 126279b and 503eff3.

📒 Files selected for processing: 2

• tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py
• tests/gpu/torch/puzzletron/test_nas_convert.py

💤 Files with no reviewable changes: 1

• tests/gpu/torch/puzzletron/test_nas_convert.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 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, danielkorzekwa

Merge Risk: ⚪ Minimal · up to 503ef

The test-only changes are mergeable; normal GPU CI remains the appropriate validation.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

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

Passed checks: 5 passed

Check name	Status	Explanation
Title check		The title clearly and concisely summarizes the main changes: using local tokenizers and calibration data in GPU tests, and reducing Puzzletron test data and workload.
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		No security anti-patterns were introduced. The reviewed diff changes only test utilities, GPU test files, and YAML configurations; it contains no Python changes under modelopt/ or examples/. The d…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

No security anti-patterns were introduced. The reviewed diff changes only test utilities, GPU test files, and YAML configurations; it contains no Python changes under modelopt/ or examples/. The diff adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec comments. It also changes no pyproject.toml or requirements.txt files.

✨ 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 78.91%. Comparing base (6d76b59) to head (55e1c27).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2850      +/-   ##
==========================================
+ Coverage   70.16%   78.91%   +8.75%     
==========================================
  Files         659      659              
  Lines       74081    74081              
==========================================
+ Hits        51978    58462    +6484     
+ Misses      22103    15619    -6484     

Flag	Coverage Δ
examples-diffusers	20.04% <ø> (-0.01%)	⬇️
examples-gpt-oss	13.51% <ø> (-0.01%)	⬇️
examples-hf_ptq	23.19% <ø> (-0.04%)	⬇️
examples-llm_distill	13.59% <ø> (-0.01%)	⬇️
examples-llm_eval	17.24% <ø> (+<0.01%)	⬆️
examples-llm_qat	17.48% <ø> (-0.01%)	⬇️
examples-llm_sparsity	15.76% <ø> (-0.01%)	⬇️
examples-megatron_bridge	26.43% <ø> (-0.13%)	⬇️
examples-specdec_bench	13.29% <ø> (-0.01%)	⬇️
examples-speculative_decoding	17.47% <ø> (-0.01%)	⬇️
examples-torch_onnx	21.39% <ø> (-0.01%)	⬇️
examples-torch_trt	15.24% <ø> (-0.01%)	⬇️
examples-vllm_serve	15.08% <ø> (-0.01%)	⬇️
gpu	59.13% <ø> (+37.00%)	⬆️
regression	15.10% <ø> (-0.01%)	⬇️
unit	60.60% <ø> (+<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

[cjluo-nv]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

LGTM: the test reductions preserve the relevant assertions; human sign-off is required for the justified test consolidation and input changes.

Needs action:

• Confirm the justified test edits: test_nas_search.py retains both deleted FFN-conversion assertions, and local calibration/tokenizer inputs preserve export, GPTQ, and preprocessing coverage.

No action needed:

• The design reuses existing tokenizer fixtures, dataset loaders, and filesystem guards; the PR explains why Hub data and larger inputs are unnecessary.
• The new helper’s header matches LICENSE_HEADER.
• Tests were inspected, not executed during this review.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-21 19:59 UTC