mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-14 05:00:00 +08:00
main
6
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Chenjie LuoandClaude Opus 5.5
|
af07754681 |
Speed up the diffusers example tests (#2724)
### What does this PR do?
Type of change: Tests (CI speed-up)
The diffusers example lane spends ~29 min in pytest, almost all of it in
~45 example-script subprocesses on tiny models. This PR removes
redundant and avoidable work without dropping coverage:
- **Export HF checkpoints from runs that already happen.** The SDXL
int8/fp8 and Wan 2.2 14B transformer int8/fp8 configs in
`test_export_diffusers_hf_ckpt.py` used exactly the same arguments as
runs in `test_diffusers.py`. Those quantize runs now also pass
`--hf-ckpt-dir` (`hf_export=True`) and check the export, and the four
duplicate export tests are removed (30 → 26 tests in these files).
- **Calibrate from a local prompts file.** `quantize.py` loaded the Hub
prompt datasets on every calibration run. For Wan 2.2 that is all of
`nkp37/OpenVid-1M`, which cost ~30s per run. Tests now pass
`--prompts-file tests/examples/diffusers/calib_prompts.txt`, which also
removes a network dependency.
- **Tiny PixArt for `test_pixart_cachify`.** It downloaded the full-size
`PixArt-XL-2-1024-MS`. It now uses a tiny local PixArt-Alpha pipeline,
with the same tiny config as diffusers' own PixArt tests and built like
the existing tiny Wan 2.2 fixture. `cachify` still wraps both
transformer blocks. `PIXART_PATH` is removed because nothing else uses
it.
- **Cheaper `torch.compile` tests.** `diffusion_trt.py` compiles with
`mode="max-autotune"`. The compile tests now set
`TORCHINDUCTOR_MAX_AUTOTUNE_{GEMM,CONV}_BACKENDS=ATEN`, so Inductor
skips benchmarking Triton GEMM/conv templates. The `--torch-compile`
path still runs.
### Usage
N/A (tests only).
### Testing
I ran these on an RTX PRO 6000 (same GPU as the CI runner) in
`nvcr.io/nvidia/pytorch:26.07-py3`, set up like the CI lane: editable
`.[onnx,hf,dev-test]`, `diffusers==0.39.0`, apex uninstalled. Coverage
was on, the same way CI collects it (`COVERAGE_PROCESS_START`, `--cov`).
- Branch: `pytest tests/examples/diffusers
--ignore=test_diffusers_onnx_trt.py --ignore=fastgen --cov` → **40
passed in 22m37s**. fastgen needs `nemo_automodel`, which I didn't
install; this PR doesn't touch it.
- `main`, same machine: `test_diffusers.py`,
`test_export_diffusers_hf_ckpt.py`, `test_cache_diffusion.py -k "not
pixart"` → **29 passed in 28m21s**. This set is smaller than the branch
run: no sparsity, PixArt or deprecation tests (~1m50s on the branch).
Per-test call times (s), with CI from run 37945140337:
| Test | CI (main) | local main | local branch |
|---|---:|---:|---:|
| `test_wan22_quantization[*]` (7 tests, each) | 66–86 | 79–98 | 49–52 |
| `test_wan22_hf_ckpt_export` (2 tests) | 44 + 46 | 54 + 54 | removed |
| `test_diffusers_hf_ckpt_export[sdxl int8 / fp8]` | 43 + 37 | 32 + 31 |
removed |
| `test_pixart_cachify` | 86 | not run | 1.8 |
| `test_diffusion_trt_torch[sdxl_1.0_torch_compile]` | 132 | 156 | 119 |
| `test_diffusion_trt_torch[flux_schnell_torch_compile]` | 60 | 71 | 57
|
| `test_diffusion_trt_torch[sd3_medium_torch_compile]` | 51 | 58 | 49 |
On the same test set this saves about 7.5 min locally, plus about 84s
for PixArt. Scaled to CI timings, I expect the lane's pytest step to
drop from ~29 min to ~21 min.
The ATen restriction cuts the SDXL compile run from 108s to 63s without
coverage (run on its own). Coverage tracing of Inductor's Python
accounts for most of what remains.
`pre-commit` passes on the changed files.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: N/A (tests only)
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A (the tiny
PixArt config values mirror diffusers' own tests, as the existing tiny
Wan 2.2 helper does)
- Did you write any new necessary tests?: N/A
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A (tests only)
- Did you get Claude approval on this PR?: ❌
### Additional Information
Possible follow-ups not in this PR:
- Run fewer restore subprocesses: one per format instead of all 13,
about 3 min, but it gives up some coverage.
- Run the example scripts in-process like `megatron_example_runner.py`:
about 20s of startup saved per subprocess, but a larger change.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Tests**
* Expanded diffusion-model test coverage for PixArt-Alpha, including
cache behavior and checkpoint exports.
* Added validation that exported checkpoints include configuration and
model-weight files.
* Updated quantization tests to use local calibration prompts and cover
checkpoint export for selected SDXL and Wan 2.2 runs.
* Adjusted compiled test runs to use compatible autotuning settings.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
|
||
|
Keval MorabiaandClaude Opus 4.8
|
0081473861 |
Speed up slow unit/gpu/example tests (#1616)
### What does this PR do?
Type of change: test infrastructure / test speedups + CI stabilization
Make the test suite faster, `tests/unit` hermetic, and the CI lanes
stable, without losing coverage. Most changes are mechanical test/infra
edits; the buckets below cover the diff broadly.
**Unit tests — hermetic (no HF Hub):** toy local datasets/configs + the
local tiny tokenizer (with a checked-in chat template) replace Hub
assets; `tests/unit/conftest.py` enforces offline mode. Genuinely-HF
tests moved to `tests/gpu*` (e.g. the new
`tests/gpu/torch/utils/test_dataset_utils.py`). `CONTRIBUTING.md`
documents the hermetic-unit-test expectation.
**Unit-test speedups (no coverage loss):** speculative (disable CPU
torch.compile), calibrator (fewer histogram bins), ONNX conv/dynamo
(smaller shapes + representative subset), Ruler/sparse-attention (local
tokenizer), data-parallel autoquant (world size 4→2). Shared
`tiny_tokenizer` fixture. The distributed test helper now uses a private
`spawn` context instead of mutating the global start method (avoids
cross-test contamination).
**Rarely-used autonas/fastnas tests:** heavy parametrize cases marked
`@pytest.mark.manual`, one representative kept per test (fastnas
preferred); lighter sibling tests still cover core behavior. The legacy
FSDP1 NAS distributed test is also dropped: FastNAS/AutoNAS aren't used
with either FSDP1 or FSDP2, and FSDP1 is superseded by the newer FSDP2
API — so we keep a single FSDP2 case as a sanity check and drop FSDP1,
leaving the suite leaner.
**gpu_megatron:** deduplicate distributed worker pools by world_size
within a module (saves a redundant pool spin-up in multi-pool files;
module-scoped, no cross-module reuse).
**Example tests:** reduce per-test work via args that default to current
behavior (tests pass the fast values) — torch_onnx TRT optimization
level, diffusers calibration/inference steps, eagle `sample_size`,
megatron_bridge iters/calib, llm_sparsity data slice, export
safetensors-structure `calib_size`. Also enable the recently added
`gpt-oss` example tests in CI.
**Per-test timeouts:** `pytest-timeout` with a default per-directory
timeout (60s unit / 300s gpu+example) enforced in `tests/conftest.py`
(`timeout_func_only` in `pyproject.toml`), so a new test cannot silently
exceed the budget — an unmapped test dir crashes collection. A few
inherently slow tests carry explicit higher per-test overrides
(CUDA-compile, autotune, dflash).
**CUDA kernel pre-compilation:** a dedicated `tests/gpu/_extensions`
test JIT-builds the conv3d implicit-GEMM kernel up front (collected
before the functional tests in the same process) so the one-time build
cost no longer lands on — and time out — the first functional test that
uses it. Mirrored into the `llm_ptq`/`vlm_ptq` example lanes.
**Test relocation & optional-dependency guards:** vLLM sparsity plugin
test moved to `tests/gpu_vllm` (drops the in-test `importorskip`);
diffusers-dependent unit test guarded with `importorskip("diffusers")`
for partial-install lanes; `gpt_oss` example test dir renamed to
`gpt-oss` to match the CI matrix.
**Diffusers test models:** shared model-path constants in
`tests/_test_utils/examples/models.py` consolidated/renamed and point at
tiny `hf-internal-testing` test pipes (SDXL/SD3/FLUX) so
cachify/quantize/export tests run on toy weights; `local_id`s
normalized.
**Shared dataset utils:** `examples/llm_sparsity/.../hf_pts.py` now uses
`get_dataset_dataloader` (drops the bespoke cnn_dailymail-only
`get_calib_dataloader`; supports any registered/HF/JSONL dataset,
includes attention_mask); `data_prep.py` gains `--max_samples`.
**CI workflows:** container image bumps (pytorch 26.04→26.05, TRT-LLM
rc16→rc17) and tightened lane timeouts (unit 30→15 min, gpu lanes
trimmed, onnx example lane 45 min).
**Imports at top of file:** in-function imports across the test suite
are moved to module top per the coding guideline, conservatively —
optional deps stay guarded (in-function or behind a module-level
`importorskip`) in `tests/unit` since the partial-install lane runs
without them, and build/hardware-availability imports (apex, triton,
megatron/transformer_engine, tensorrt_llm) plus `_test_utils` lazy
guards are left in place.
**Kernel warning filters:** the repeated `filterwarnings` blanket-ignore
in six `tests/gpu/torch/kernels/**` modules is consolidated into a
scoped hook in `tests/gpu/torch/kernels/conftest.py` (kernel tests only
— the rest of the suite keeps surfacing warnings).
**Eagle example speedups:** `torch.compile` (eagle recipe default) added
~2 min to every eagle training test; it's now disabled in the eagle
example tests except one smoke (`test_llama_eagle3[1-False]`), and the
downstream resume / AR-validate / export tests point at the compile-free
checkpoint. Measured: `test_ar_validate` 139s→17s, offline training
142s→22s, streaming 140s→23s — the compile path is still smoke-tested
once.
**Example lanes install editable (`-e`):** so example scripts launched
as subprocesses resolve `modelopt` to the same source path as the test
process and reuse the pre-compiled CUDA-extension cache instead of
recompiling (~2 min/test); verified in the TRT-LLM container.
**Tiny test tokenizer:** `get_tiny_tokenizer` defaults to left padding
(what decoder-LM calibration expects) and ships a terse
generation-tagged chat template — replacing a verbose ChatML one that
inflated tokenized length on the 128-vocab tokenizer and broke the
offline-PTQ example tests' `max-seq-len` filter.
**Restored Hub-download coverage:** the live (ungated) HF dataset
round-trips exercising `get_dataset_samples`' download branch now live
in `tests/gpu/torch/utils/test_dataset_utils.py` (they had been dropped
from the hermetic unit file without a counterpart).
Individual file changes not explicitly called out above fall under this
general test/CI cleanup.
### Testing
Unit + the touched gpu_megatron files validated locally; example/GPU
lanes validated in CI.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅ (tests + example CLI args
default to prior behavior)
- If you copied code from any other sources or added a new PIP
dependency: N/A
- Did you write any new necessary tests?: N/A (optimizes/relocates
existing tests)
- Did you update Changelog?: N/A
- Did you get Claude approval on this PR?: ❌ (pending)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
## Release Notes
* **Chores**
* Updated container image versions for PyTorch (26.04→26.05),
TensorRT-LLM (1.3.0rc16→1.3.0rc17), and ONNX/TensorRT (26.04→26.05).
* **Tests**
* Enhanced test isolation: unit tests now run hermetically without
HuggingFace Hub access.
* Optimized test runtime via smaller model/dataset parameters and
parallel test caching.
* Added CUDA extension availability tests and extended dataset utility
coverage.
* **Documentation**
* Updated testing guidelines in `CONTRIBUTING.md` to emphasize offline
test design.
* **Chores**
* Added pytest timeout configuration and improved CI/CD workflow
efficiency with editable installs.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
||
|
Keval Morabia
|
1ccd945a51 |
Remove unused diffusers/cache_diffusion/pipeline and cuda-python dependency (#996)
`cuda-python` has mixed license and needs EStaff approval for usage. And till 0.42, it was only used in `examples/diffusers/cache_diffusion/pipeline` which has not been updated in 9 months and not used anymore hence removing. Also cherry-picked to `release/0.42.0` branch: https://github.com/NVIDIA/Model-Optimizer/pull/984 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Removed TensorRT/ONNX deployment and inference tooling, related model export/configuration, and runtime helpers from the cache-optimized diffusion examples; removed the cuda-python example dependency. * **Tests** * Removed the example benchmarking script and its associated benchmark test. * **Documentation** * Strengthened dependency-review, security, and PR guidance; updated PR template and contributing documentation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Keval Morabia
|
cb8f9efa50 |
Refactor test utils (#476)
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Keval Morabia
|
cafa7f60ed | Update for 0.33.0 release | ||
|
Keval Morabia
|
7a047435ae | Update for 0.29.0 release |