mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Chenjie LuoandClaude Opus 5.5
3091b8ff69
### What does this PR do? Type of change: new feature **Second of two PRs adding IQ2_S** (2.5625 bits per weight). #2512 landed the PyTorch codec; this PR adds its **CUDA encoder** and makes the format reachable: - the CUDA encoder, its binding and extension build wiring, plus the CUDA path in `quantize_iq2_s` - an `IQFormat` record and **one `IQ_FORMAT_REGISTRY` entry**, so backend dispatch, both exporters and `convert_hf_config` take it from there - the `ggml` package export - the `general/ptq/iq2_s` recipe, its presets, `ptq.md` and a CHANGELOG entry The kernel lands with the registration so every registered format keeps a CUDA encoder. On the mixed-precision checkpoint #2511 measured (`unsloth/Qwen3.8-27B-GGUF`), IQ2_S covers **9 tensors and 0.6 B parameters**. ### The kernel IQ2_S's **1024-entry codebook is twice IQ2_XS's**, which makes its search the most expensive in the family. The codebook and its norms take 36 KiB of shared memory, the most of any IQ kernel but inside the 48 KiB static limit, so they are declared statically like the IQ2_XS and IQ2_XXS kernels. That cost is why the kernel matters more here than anywhere else: | | torch | CUDA | | |---|---|---|---| | IQ2_S, 5632×2048 weight | 0.8 M elem/s | **725.7 M elem/s** | **907×** | | extrapolated to a 27B model | ~9.8 hours | **~37 s** | | ### Usage ```bash python examples/hf_ptq/hf_ptq.py --pyt_ckpt_path <model> --recipe general/ptq/iq2_s ``` ### Testing Registering the format brings it under every registry-driven test with no IQ2_S-specific test code: backend dispatch and weight caching, the `num_bits` guard, `convert_hf_config` metadata (uniform and mixed precision), all 9 Megatron export tests, and the two `TensorQuantizer` tests in the shared battery. The shared CUDA battery gains one row. - `tests/unit/torch/quantization/test_ggml_backend.py`, `test_iq_formats.py`, `tests/unit/torch/export/test_convert_hf_config.py`, `tests/unit/recipe/test_presets.py`: **134 passed** - broader unit sweep (`-k 'ggml or iq or gguf or registry'` over quantization, export and recipe tests): **192 passed**. The one failure, `test_export_registry.py::test_builtin_dispatch_covers_all_handler_shapes`, is a `torchvision` import error in my environment, unrelated to IQ. - `tests/gpu/torch/quantization/test_iq_formats_cuda.py`, `test_iq1_s_cuda.py`, `test_iq2_xs_cuda.py`: **42 passed** on RTX PRO 6000 Blackwell (sm_120). 7 of them are IQ2_S: CUDA-vs-PyTorch encoder parity, determinism, reconstruction at scale, zero and non-finite policy, float64 input and the fallback path. - `tests/gpu_megatron/torch/export/test_unified_export_megatron.py -k 'iq or ggml'`: **36 passed** (9 tests × 4 formats) in `nvcr.io/nvidia/nemo:26.08` - `tests/examples/hf_ptq/test_llm_ptq.py -k iq2_s`: **passed**. TinyLlama PTQ through unified HF export writes `quant_algo: IQ2_S`, `block_payload_bytes: 82`, and `down_proj` packed as `(2048, 22, 82)` uint8. - `general/ptq` now holds 30 recipes. - The shared-memory change in `b7739d5d0` leaves the packed bytes identical (same hash on a 5632×2048 weight), and packing runs at 849.1 M elem/s against 825.7 before on RTX PRO 6000. The GPU battery was rerun: 42 passed. All of the above was rerun after rebasing onto `main` at `c2aaa44f6`. That base adds a Q8_0 packer to the same GGML extension (#2515), and changes the hf_ptq example and the export code this format goes through. The packed IQ2_S bytes still hash the same. On this RTX PRO 6000 (sm_120), two of #2515's own Q8_0 tests in `tests/gpu/_extensions/test_torch_extensions.py` fail: `test_cuda_ext_q8_0_zero_and_roundf_layout` and `test_cuda_ext_q8_0_dequantizes_with_small_error`. They fail identically on a clean `main` checkout, so they are not from this PR. ### 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`: ✅ No new code sources or dependencies. - Did you write any new necessary tests?: ✅ - Did you update Changelog?: ✅ - Did you get Claude approval on this PR?: ❌ Not yet run. ### Additional Information Merge order: #2511 (IQ2_XXS, merged) → #2525 (format registry, merged) → #2512 (IQ2_S codec, merged) → **this** → #2513 (IQ1_M). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added IQ2_S weight-only quantization for eligible linear layers, at 2.5625 bits per weight. * Added a PTQ recipe that requires no calibration data. Weights must meet the existing 256-value block-size constraint. * Added CUDA-accelerated packing for CUDA weights, with a Python fallback when the CUDA extension is unavailable. * **Documentation** * Updated the PTQ recipe catalog and IQ-format size tradeoffs. <!-- 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>