Commit Graph
6 Commits
Author SHA1 Message Date
Keval Morabia 19e0121d0a Fix 2-GPU test_model_load_utils hang; test import and fixture cleanup (#2079)
### What does this PR do?

Type of change: Bug fix + new tests (test-suite only; changes are
confined to `tests/`)

- **Fix `test_parallel_load_and_export` hang on 2 GPUs.** The temp paths
were built from `os.getpid()` *inside* the workers, so each rank got a
different `ckpt_dir`; rank 1 then took `_resolve_checkpoint_dir`'s Hub
branch and blocked on an extra barrier while rank 0 ran the loader's
broadcasts. The checkpoint is now built once in the parent under
`tmp_path` and passed in, so all ranks agree on the path.
- **Imports moved to module top** across `tests/gpu*` and
`tests/examples`; function-local imports kept only where guarded
(`importorskip`/`try`), where the import *is* the test (JIT compile), or
where it must follow `sys.path` setup.
- **Reuse `_test_utils` instead of local copies:** added
`get_tiny_mixtral`; deduped `assert_nodes_are_quantized` (5 copies), the
accelerate-offload/layerwise config helpers, `make_quant_attention`,
`get_dflash_config`, the NVFP4 amax assertions, and 3 copies of the
`tiny_wan22_path` fixture.
- **Shared model-dir fixtures assert they were not modified**
(`assert_unmodified_tree`): a file manifest is compared on teardown, so
a test that writes into a session-scoped fixture directory fails instead
of silently changing what later tests see.
- **Dropped dependency guards the CI env already guarantees**
(diffusers, tensorrt_llm in `gpu_trtllm`, transformer_engine in
`gpu_megatron`, transformers in examples) so a missing dep fails loudly
instead of skipping.
- **`test_heterogenous_sharded_state_dict` is skipped on Blackwell**
(sm_120), matching the existing marker and its TE/CUDA-13 rationale —
same tracking issue as #1901.

### Testing

Local, 2x RTX 6000 Ada: `tests/gpu/torch/utils/test_model_load_utils.py`
passes on 2 GPUs and on 1 GPU (previously hung on 2). Also ran the
touched files in `tests/unit` (608 passed) and `tests/gpu` (~340
passed).

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: N/A — test-only PR
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A
- Did you get Claude approval on this PR?: ❌ — not yet run


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Tests**
* Added reusable validation utilities for generated files, quantization
behavior, attention modules, model fixtures, offloading, and speculative
decoding.
* Expanded coverage for tiny Wan, Mixtral, Llama, and related model
scenarios.
* Consolidated duplicated setup and assertions across ONNX, GPU,
quantization, export, and sparsity tests.
* Improved fixture integrity checks, NVFP4 validation, and handling of
identity inputs.
* Reduced unnecessary dependency-based skips and isolated known
platform-specific flakiness.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-08-05 21:00:42 +05:30
jingyu-mlandClaude Opus 4.8 6b4ad85849 Qwen-Image diffusers PTQ: FP8 / NVFP4 / NVFP4-SVDQuant HF checkpoints (#1706)
### What does this PR do?

Type of change: New feature

Adds **Qwen-Image** (`Qwen/Qwen-Image`, `QwenImageTransformer2DModel`)
to the diffusers quantization example and exports HuggingFace
checkpoints in three precisions — **FP8**, **NVFP4**, and **NVFP4 +
SVDQuant** — through the unified HF export.

- Registers `--model qwen-image` (lazy diffusers import; no
`trust_remote_code`).
- Transformer-block-range recipe: quantizes only the linears under
`transformer_blocks`, keeping the **first 2 / last 2** blocks (and
everything outside `transformer_blocks`) in original precision. Applied
**before** calibration so SVDQuant never mutates the excluded blocks.
Expressed with the top-level `enable` `QuantizerCfgEntry` field
(disable-all → re-enable `transformer_blocks` → disable first/last-N).
- SVDQuant export (AWQ-style): promotes quantizer-owned tensors to clean
module-level safetensors keys at export time —
`weight_quantizer.svdquant_lora_a/b → <module>.svdquant_lora_a/b` and
`input_quantizer._pre_quant_scale → <module>.pre_quant_scale` — with a
documented `NVFP4_SVD` `quantization_config` (`group_size`,
`has_zero_point: false`, `pre_quant_scale: true`, `lora_rank`). **Core
SVDQuant quantization code (`modelopt/torch/quantization`) is
unchanged.**
- Shared export-path change — **intentionally global** (applies to all
diffusers exports — SDXL / Flux / Wan, not just Qwen; the full export
suite was verified green on GB200): `hide_quantizers_from_state_dict`
now strips quantizer state from *all* modules (not just quant-linears)
so calibrated norm-layer input quantizers no longer leak
`input_quantizer._amax`. (An earlier `max_shard_size` workaround was
dropped after merging `main`: #1794 makes the ComfyUI layerwise-metadata
post-processing a no-op unless explicitly opted in, so a default sharded
export no longer hits the unsupported-sharded path.)

### Usage

```bash
python examples/diffusers/quantization/quantize.py \
    --model qwen-image --override-model-path <Qwen-Image> --model-dtype BFloat16 \
    --format fp4 --quant-algo svdquant --lowrank 32 \
    --calib-size 64 --n-steps 20 \
    --hf-ckpt-dir <out>
# FP8:   --format fp8 --quant-algo max
# NVFP4: --format fp4 --quant-algo max
```

### Testing

- Focused unit + example tests pass on GB200 (sm_100): block-range
recipe, `NVFP4_SVD` config schema, SVDQuant forward/fold (LoRA stays on
`weight_quantizer`), Qwen dummy-input / strict-QKV-fusion / promotion,
pipeline loading, and the diffusers HF-export test for Qwen FP8 / NVFP4
/ SVDQuant.
- Full `tests/examples/diffusers/test_export_diffusers_hf_ckpt.py` is
green (SDXL, Flux, Qwen, Wan2.2) — confirms the shared export changes do
not regress other models.
- End-to-end on the real `Qwen/Qwen-Image` (~20B): all three formats
export valid HF checkpoints — only `transformer_blocks` 2..57 quantized,
nothing outside, no quantizer/`_amax` leak, correct
`weight_scale`(`_2`)/`input_scale`, promoted SVDQuant keys
(rank-consistent shapes), and the expected `quantization_config`.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ <!-- live-model LoRA storage
unchanged; existing exports unaffected -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update Changelog?: ❌ <!-- new-example feature; add a
CHANGELOG.rst entry if required -->
- Did you get Claude approval on this PR?: ❌ <!-- run /claude review -->

### Additional Information

All changes are confined to the diffusers example
(`examples/diffusers/quantization`) plus the shared export path
(`modelopt/torch/export`); the core quantization library is untouched.

### Follow-up (next step): fused-QKV SVDQuant for sglang / Nunchaku

This export keeps attention `q/k/v` (and `add_q/k/v_proj`) as
**separate** projections — the diffusers-native layout. That matches
sglang's bf16 / FP8 / plain-NVFP4 paths (which also keep QKV separate)
and ModelOpt/TRT-LLM consumers, so those load 1:1.

sglang's **NVFP4-SVDQuant (Nunchaku)** path, however, builds a **fused**
`to_qkv` with a *single* fused rank-r LoRA in Nunchaku-native format
(`proj_down`/`proj_up`, `smooth_factor`, `wscales`/`wtscale`). Our
per-projection tensors (`svdquant_lora_a/b` + `pre_quant_scale`; three
independent rank-r decompositions) are not directly loadable there — and
cannot be fused at load time, because the fp16 weight residual needed to
derive a single fused rank-r is not preserved after export.

**Planned next step:** an opt-in fused-QKV SVDQuant export mode that
fuses q/k/v **before** SVDQuant calibration (yielding one rank-r over
the fused weight) and emits a Nunchaku-compatible layout, enabling
lower-latency fused-QKV inference in sglang. Tracked as a separate
follow-up.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added Qwen-Image (`QWEN_IMAGE`) model quantization and Diffusers
export support
  * Added NVFP4_SVD (SVDQuant) export configuration support
* Added transformer block-range quantization recipes (exclude first/last
blocks)
* **Bug Fixes**
  * Improved missing-pipeline error messaging for Qwen-Image
* Prevented quantizer-related tensor/buffer leakage by promoting and
cleaning quantizer outputs during export
* **Tests**
  * Added Qwen-Image HF checkpoint export tests and offline fixtures
  * Added unit coverage for SVDQuant promotion/clean state-dict keys
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-07 17:21:31 -07:00
jingyu-mlandClaude Opus 4.8 aec72ffa68 Add DMD2 distillation for Qwen-Image (fastgen) (#1326)
### What does this PR do?

**Type of change:** New example + new `modelopt.torch.fastgen` library
module.

Adds **DMD2 (Distribution Matching Distillation) for Qwen-Image** —
distilling the base model into a few-step (1–4) generator. Includes the
framework-agnostic `modelopt.torch.fastgen` loss library (DMD pipeline,
EMA, optional GAN discriminator) and a NeMo AutoModel–based training
example with a mock-data smoke config, a real-data config, and inference
/ export scripts.

**Noted**: the example script will be migrated to AutoModel repo

### Usage

```bash
# Mock-data wiring smoke — runs end-to-end with no dataset to prepare
torchrun --nproc-per-node=8 \
    examples/diffusers/fastgen/dmd2_finetune.py \
    --config examples/diffusers/fastgen/configs/dmd2_qwen_image_smoke.yaml
```

See `examples/diffusers/fastgen/README.md` for real-data training and
inference.

### Testing

Unit tests under `tests/unit/torch/fastgen/`; `pre-commit` /
code-quality clean.

### Before your PR is "*Ready for review*"

- Backward compatible?: ✅ (new, additive module)
- Followed `CONTRIBUTING.md` for any copied code / new deps: ✅
- New tests added?: ✅
- Updated Changelog?: N/A


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Adds a FastGen-based distillation framework (DMD2) with
student/fake-score training, EMA support, GAN discriminator branch,
inference pipeline, and export utilities.
* Qwen-Image integration with latent packing and feature-capture for
plugin-enabled pipelines.

* **Documentation**
* New README, example configs, and runnable example scripts for
Qwen-Image distillation and inference.

* **Tests**
* Comprehensive unit tests covering math parity, gradient routing,
plugins, hooks, EMA, and recipe setup.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-08 16:03:35 -07:00
jingyu-ml 26ae8da517 [2/3] Implicit Gemm NVFP4 (#1227)
### What does this PR do?

Type of change: new feature <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

- Add Conv3D implicit GEMM kernel with BF16 WMMA tensor cores and fused
NVFP4 activation quantization for video diffusion VAE layers
- Integrate into _QuantConv3d via QuantModuleRegistry — automatically
dispatched when NVFP4 quantization is applied to nn.Conv3d
- Move kernel from `experimental/conv/ to modelopt/torch/kernels/conv/`;
move tests to `tests/gpu/torch/quantization/kernels/`

### Testing
<!-- Mention how have you tested your change if applicable. -->

- Added test cases to measure the difference between cuDNN and our CUDA
implicit GEMM kernel
- Added an NVFP4 fake quantization test using CUDA code

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ <!--- If ❌, explain why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ <!---
Mandatory -->
- Did you write any new necessary tests?: ✅ <!--- Mandatory for new
features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ <!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Per-backbone quantization/export in a single run with per-backbone
checkpoints and backbone-aware quant filters
* Configurable NVFP4 block-size via CLI/config; improved NVFP4 Conv3D
inference path and Wan 2.2 quantization support
* **Bug Fixes**
* Video-model calibration now respects extra params and forces video
decoding during calibration
* **Documentation**
* Added comprehensive Conv3D implicit‑GEMM kernel documentation; removed
experimental Conv3D prototype docs/benchmark
* **Tests**
* New Wan 2.2 quantization/export tests and expanded Conv3D/FP4 kernel
test coverage
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
2026-04-19 12:20:14 +05:30
Keval Morabia 4c611e47a6 Update files on GitHub
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2025-07-31 22:48:09 +05:30
Keval Morabia 7a047435ae Update for 0.29.0 release 2025-05-08 23:43:51 +05:30