Repository navigation
[OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data - #2850
Open
shengliangxu wants to merge 4 commits into
Open
[OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data#2850shengliangxu wants to merge 4 commits into
shengliangxu wants to merge 4 commits into
Conversation
|
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. |
Contributor
|
Warning Review limit reachedEnable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. Next included review available in 42 minutes. View limit details
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 30.77% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 13 functions across 6 files. (9 skipped: … | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation 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 changed Python files are under modelopt/orexamples/. The PR adds nopyproject.tomlorrequirements.txtchanges. No prohibited security patterns were added in the changed Python files.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.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 30.77% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 13 functions across 6 files. (9 skipped: 9 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
✨ Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- 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
added this pull request to stack #2856
October 17, 2026 01:41
This was referenced Oct 17, 2026
shengliangxu
changed the title
This was referenced Oct 17, 2026
Collaborator
Author
|
/ok to test 2d82a59 |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## shengliangx/ci-5-onnx-tests #2850 +/- ##
============================================================
Coverage 78.94% 78.94%
============================================================
Files 658 658
Lines 74023 74023
============================================================
Hits 58435 58435
Misses 15588 15588
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
shengliangxu
changed the title
shengliangxu
removed this pull request from stack #2856
October 17, 2026 06:58
shengliangxu
force-pushed
the
shengliangx/ci-6-local-data-tests
branch
from
October 17, 2026 07:02
2d82a59 to
7383ca1
Compare
shengliangxu
added this pull request to stack #2859
October 17, 2026 07:02
shengliangxu
changed the title
Collaborator
Author
|
/ok to test 7383ca1 |
shengliangxu
changed the title
shengliangxu
force-pushed
the
shengliangx/ci-6-local-data-tests
branch
from
October 18, 2026 23:15
7383ca1 to
5f3497e
Compare
shengliangxu
marked this pull request as ready for review
October 19, 2026 03:44
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
| from safetensors import safe_open | ||
|
|
||
|
|
||
| @pytest.fixture(scope="module") |
Collaborator
There was a problem hiding this comment.
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
force-pushed
the
shengliangx/ci-6-local-data-tests
branch
from
October 19, 2026 21:34
5f3497e to
1b417b8
Compare
Six tests in test_megatron_preprocess_data.py tokenize with gpt2 from the Hugging Face Hub. megatron_preprocess_data loads the tokenizer once for the vocabulary size and again in every worker, so each test makes several hub round trips, and none of them depends on gpt2's vocabulary: they check output naming, non-empty .bin and .idx files, and that max_tokens shortens the output. Switch them to the tiny_tokenizer_path fixture, the checked-in byte-level tokenizer that the blend test already uses. Three of the six read only local data (the jsonl path, the jsonl max_tokens and the gzip tests), so the tokenizer download was their only network access. The PR history lists the jsonl max_tokens test, with the HF split max_tokens test, among the lane's 120 s timeouts (3 failures each); it no longer has a network dependency to time out on. The other three (HF split max_tokens, HF resume and the HF streaming warning) still read the real HF dataset, which is the code path they exist to cover; only their tokenizer changes. The Qwen3 tests keep their real tokenizer: they need its chat template and EOS handling. The blend test already uses the tiny tokenizer, and its HF split source is the path it covers, so it is unchanged. No test ids, assertions or markers change. Measured on a CPU host with a cold Hugging Face cache, the nine tests in the module take 38.5 s before and 24.4 s after: jsonl path 3.3 s to 0.01 s, jsonl max_tokens 3.6 s to 0.13 s, gzip 1.6 s to 0.07 s, HF split max_tokens 7.0 s to 4.9 s, HF resume 5.0 s to 3.1 s, HF streaming warning 4.4 s to 2.8 s. CI jobs start without a warm hub cache, so a similar saving is expected. Testing: ruff check and ruff format with the pinned ruff 0.15.20, and py_compile. Ran the module on CPU, next to a copy of the tiny_tokenizer_path fixture from tests/gpu_megatron/conftest.py, with the Hub reachable: 4 passed and 5 xpassed before and after, and the collected test ids are identical. Not run: the gpu_megatron lane itself, which needs GPUs and prebuilds CUDA extensions in its conftest. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…other - save_dummy_dataset wrote 2500 copies of one conversation into each of the train and valid splits. The tiny test tokenizer turns a conversation into 1424 tokens, so that is 1738 blocks of 2048 tokens, while the pruning scores read 100 blocks and the scoring and realize steps 2 each. Every dataloader build tokenizes the whole split, and a run builds three of them. The 2500 came from an older 8192-token Llama-3 setup. Use 400 conversations (278 blocks). On a CPU the three builds plus the dataset save took 2.8 s per test with 2500 conversations and 0.36 s with 400. The scoring step checks lm_loss with a tolerance of 0.15; on random-init toy Llamas the validation loss moved by at most 0.02 between the two sizes, because the same text repeats and only the choice of blocks changes. - calc_subblock_stats ran 12 passes (batch sizes 64, 96 and 128 times four dtype sets), but the MIP and the teacher memory and parameter checks only read the entry for the batch_size in mip.subblock_stats_args, which is 96 in all nine configs. batch_sizes: [96] gives 4 passes (0.56 s to 0.18 s for a toy Llama on a CPU; MoE configs have more subblocks per pass). - test_nas_convert_ffn_pruning ran the same mtn.convert call on the same toy Llama and config as test_nas_search (21.6 s in the CI logs, with its own 2-rank spawn). Its two assertions move into test_nas_search, which checks them with the other rank-0 assertions after the search. Left alone: the Nemotron-3 toy keeps its 128 experts, because shrinking it changes EXPECTED_TEACHER_MEMORY_MIB and EXPECTED_TEACHER_NUM_PARAMS, which need a GPU runner with mamba_ssm to regenerate. Testing: timed the dataloader builds, the subblock-stats passes and the loss sensitivity on CPU with the repository's own helpers and dataloader code, collected tests/gpu/torch/puzzletron (16 tests, one fewer than before), and ran ruff. Not run: the puzzletron tests themselves, which need GPUs and mamba_ssm. Expect the gpu job to lose roughly 50 s (the removed test is 22 s of it); the CI run will show the real number. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…or batch probing Each of the 11 test_unified_hf_export_and_check_safetensors params starts a fresh hf_ptq process, 22 s per param in the CI logs. The test only checks the structure of the exported safetensors, so three parts of that run are not needed: - Calibration text came from cnn_dailymail on the Hub, a network round trip per param. Use a local jsonl of 64 random-word texts instead, written by a new helper, write_calib_jsonl, in _test_utils/torch/calib_data.py. Each text fills the 512-token calibration window and differs from the others, so the activation statistics stay varied. - hf_ptq probed for the largest calibration batch. --batch_size 64 puts the whole calibration set (calib_size 64) in one batch. - hf_ptq ran model.generate() before and after quantization. --skip_generate skips those previews; the export is not affected. Testing: no GPU here, so I ran the test on a CPU host through a python shim that lets hf_ptq pass its GPU guard and passes --device cpu. 8 of the 11 params pass with the new command line, scale-equality assertions included: fp8, nvfp4, int4_awq, w4a8_awq_beta, int8_weight_only, w4a16_nvfp4, the Qwen3 MoE nvfp4 and the GPT-OSS fp8. The other three (nvfp4 AWQ-lite, SVDQuant and the MSE sweep) need CUDA or Triton and did not run. The dataloader returns one (64, 512) batch with a full attention mask from the helper's file. Not run: the GPU job itself. From the per-param CI timings this should save about 38 s per `gpu` job. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
test_gptq_e2e_flow streamed cnn_dailymail from the Hub for each of its three params, a network round trip per param (about 6 s per job in the CI logs) and a test that can fail on a Hub hiccup. The GPTQ calibration only needs some text, and the test has no accuracy assertion, so it now reads 8 texts from the local jsonl that write_calib_jsonl (added with the unified-export change) creates. Testing: ruff check and format with the pinned ruff 0.15.20. On a CPU host the GPTQ flow completes on that jsonl for the FP8 and INT4 blockwise configs; the NVFP4 config needs CUDA. Not run: the GPU test itself. Expected saving is a few seconds per gpu job. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
force-pushed
the
shengliangx/ci-6-local-data-tests
branch
from
October 20, 2026 01:16
1b417b8 to
413e380
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
preprocess_datatests use the tiny local tokenizer (tests/gpu_megatron). They tokenized withgpt2from the Hub, andmegatron_preprocess_dataloads 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-leveltiny_tokenizer_pathfixture that the blend test already uses. Three of the six read only local data, so the tokenizer download was their only network access (the jsonlmax_tokenstest is among the lane's 120 s timeouts, 3 failures). Measured on a CPU host with a cold Hub cache, the nine tests of the module take 38.5 s before and 24.4 s after (37 s and 23 s in CI). The Qwen3 tests keep their real tokenizer (chat template and EOS handling).save_dummy_datasetwrote 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_statsran 12 passes where only the entry forbatch_size96 is read:batch_sizes: [96]runs 4.test_nas_convert_ffn_pruningran the samemtn.converton the same toy Llama and config astest_nas_search(21.6 s with its own 2-rank spawn); its two assertions move intotest_nas_search.hf_ptqprocess (22 s per param). It now reads a local jsonl of 64 random-word texts (new helperwrite_calib_jsonl) instead ofcnn_dailymailfrom 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.cnn_dailymailfor 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.
tests/gpu/torch/puzzletron(18 tests, 17 after the folded NAS test)preprocess_datamoduleSo 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 thetiny_tokenizer_pathfixture with the Hub reachable: 4 passed and 5 xpassed before and after, with identical test ids.tests/gpu/torch/puzzletron(16 tests, one fewer than before). The Nemotron-3 toy keeps its 128 experts, because shrinking it changesEXPECTED_TEACHER_MEMORY_MIBandEXPECTED_TEACHER_NUM_PARAMS, which need a GPU runner withmamba_ssmto regenerate.hf_ptqpass 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).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.).CONTRIBUTING.md: no new dependencywrite_calib_jsonl(tests/_test_utils/torch/calib_data.py), used by two tests; no new test functionsAdditional 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.
Summary by CodeRabbit