Repository navigation
Launch the fused NVFP4 GPTQ kernel on its tensors' GPU - #2877
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. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2877 +/- ##
==========================================
- Coverage 79.49% 79.08% -0.41%
==========================================
Files 660 660
Lines 74742 74743 +1
==========================================
- Hits 59415 59112 -303
- Misses 15327 15631 +304
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:
|
|
End-to-end check on oci-jhb (4 x GB300,
|
7fe0584 to
e77ff57
Compare
b0e76a7 to
79eddeb
Compare
|
e77ff57 to
8c958af
Compare
79eddeb to
38a8fc9
Compare
**Stack (merge in order):** [1/3] #2840 GPTQ on HF fused MoE experts ← [2/3] #2864 speed-up ← [3/3] #2877 fused NVFP4 GPTQ kernel device fix. ### What does this PR do? Type of change: new feature GPTQ only collected Hessians for modules with a 2-D `weight` and one `weight_quantizer` (`is_quantized_linear`), so HF fused MoE experts (one 3-D `gate_up_proj`/`down_proj` parameter with a `ModuleList` of per-expert weight quantizers) were skipped. They fell back to round-to-nearest, or to imatrix in the IQ recipes, which gets only part of GPTQ's gain: on a dense model, imatrix alone closes 53–63% of the gap from the plain encode to GPTQ. In an MoE those experts are most of the weights. - `gptq()` now builds one `GPTQHelper` per expert and projection. Its Hessian comes from a forward hook on the projection's shared input quantizer, keyed by `_current_expert_idx`, so it holds exactly the tokens routed to that expert, after input quantization, as the dense path's Hessian does. - `GPTQHelper` takes an optional `weight_quantizer` and `weight`, gains `accumulate()`, and `update_weights()` writes the result back in place, so an expert's slice is updated inside the 3-D parameter. `GGMLGPTQHelper` pins each expert's codes to its own quantizer. - Native Transformers FP8 experts store weights in 8 bits and hand calibration dequantized copies, so an in-place update would be lost. GPTQ and imatrix now refuse them before any calibration state changes and ask for `unpack_weight()`, as folding already does. - An expert no calibration token reaches has a zero Hessian and gets the plain search, the same fallback a dense layer calibration never reaches already gets. - Each expert Hessian is `in_features²` float32: 22.5 GiB per MoE layer for GLM-5.3-Flash's experts and 13.3 GiB for Qwen3.8-Flash-Next's, 945 and 638 GiB for the whole model. Large MoEs therefore need layerwise calibration, which the IQ recipes already enable; `gptq()` reports the size and warns before allocating when it exceeds the free memory of the GPUs holding the experts. - Behavior change: the built-in IQ1/IQ2 recipes (`[gptq (layerwise), imatrix]`) now run GPTQ on the fused experts too, and imatrix then only encodes weights GPTQ leaves. The `<format>_experts_only` recipes stay imatrix-only. - Recipe comments, `ptq.md`, config docstrings and the changelog no longer say GPTQ skips the experts. The toy fused-MoE test model moved to `_test_utils` so the GPTQ and imatrix tests share it. ### Usage ```python import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe # [gptq (layerwise), imatrix]: fused experts now get per-expert GPTQ. config = load_recipe("general/ptq/iq2_xxs").quantize.model_dump() mtq.quantize(moe_model, config, forward_loop) ``` ### Testing - `test_gptq_quantizes_each_fused_expert_as_a_linear_fed_only_its_tokens`, run with IQ2_XXS (`GGMLGPTQHelper`) and INT4 (the base `GPTQHelper`): every expert's result, both projections, is bit-identical to dense GPTQ on an `nn.Linear` holding that expert's weight and fed only that expert's tokens; for IQ2_XXS the pinned GGML payload matches too. At the toy width a GGML weight is a single 256-column group, whose GPTQ encode does not use the Hessian, so the INT4 case is the one that checks routing: it fails if every token goes to expert 0's Hessian. Both fail on `main`. - `test_layerwise_gptq_calibrates_fused_experts_on_the_earlier_layers_output`: on a two-layer toy MoE, layerwise GPTQ writes every layer's experts, layer 0 matches a single pass, and layer 1 is calibrated on layer 0's quantized output. 3.8 s on CPU. - `test_native_fp8_experts_reject_packed_weight_updates` (GPTQ and imatrix): packed FP8 experts are refused with the state untouched, and the updates persist after `unpack_weight()`. Fails without the guard. - `tests/unit/torch/quantization/{test_ggml_gptq,test_ggml_imatrix,test_gptq,test_layerwise_calibrate,test_calib}.py`, `plugins/test_fused_experts.py` and `tests/unit/recipe` on CPU with transformers 5.18: 761 passed at the top of the stack. - Real MoEs on 4x GB300 with #2864 and #2877 on top: layerwise GPTQ on the routed experts of Qwen3.8-Flash-Next and GLM-5.3-Flash at IQ2_XXS, IQ2_XS and IQ2_S, including Qwen's NVFP4 `down_proj` experts through the fused Triton kernel. KL to BF16 and downstream accuracy against imatrix are in the comments. - Not run: the GPU unit tests (no local GPU). ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (new optional arguments; dense GPTQ writes its result back with `copy_` instead of replacing `.data`) - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information Follows #2735 (imatrix calibration for the GGML formats). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * GPTQ calibration now supports fused Mixture-of-Experts layers using routed tokens and per-expert calibration. * Experts receiving no routed tokens are handled with a fallback calibration path. * Imatrix calibration applies to weights GPTQ leaves unquantized while preserving already encoded weights. * Layerwise calibration helps limit Hessian memory use for large MoE models. * **Bug Fixes** * Calibration now safely rejects unsupported packed native FP8 expert weights. * **Documentation** * Updated IQ quantization recipe guidance for expert-level GPTQ and imatrix coverage. <!-- 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>
9b8f2be to
fb5079d
Compare
fb5079d to
c921d52
Compare
Triton launches on the current CUDA device. A model that device_map="auto" spreads over several GPUs keeps later decoder layers off the current one, so GPTQ with fused: true launched the kernel on the wrong GPU with pointers into another GPU's memory: a CUDA illegal memory access on NVLink systems (GB300), a Triton pointer error on PCIe ones. Launch it under torch.cuda.device of the weight block, as the other quantization kernels do. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
c921d52 to
1244f17
Compare
meenchen
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
Approving: this is a small, correct bug fix. It launches the GPTQ kernel under torch.cuda.device(w_block.device), the same way the repo's other Triton quantization wrappers do, and it comes with a 2-GPU regression test.
No action needed:
global_scaleis passed to the kernel as a Python float, so it isn't a pointer that could point at the wrong GPU. The kernel's tensor inputs and outputs all sit onw_block's device.- The new test
test_fused_nvfp4_gptq_runs_on_the_weights_gpumoves the tensors tocuda:1whilecuda:0stays current, and checks the result matches thecuda:0run exactly. It uses the sharedneed_2_gpusfixture fromtests/conftest.py. No existing tests were changed. I read the tests but didn't run them. - No licensing changes. The CHANGELOG entry is added.
|
/ok to test 1244f17 |
### What does this PR do? Type of change: new feature (performance) **Stack (merge in order):** [1/3] #2840 GPTQ on HF fused MoE experts (merged) ← [2/3] #2864 speed-up ← [3/3] #2945 Hessian placement by measured headroom. The fused NVFP4 GPTQ kernel device fix, #2877, is standalone on `main` and can merge in any order. Layerwise GPTQ on a large HF fused-MoE model took about 25 minutes per decoder layer: 20 hours for Qwen3.8-Flash-Next's gate_up experts (512 experts, 48 layers) and over a day for GLM-5.3-Flash. Profiling one Qwen-shaped layer showed the GPTQ weight update was only 9 s of that. The time went to forward passes, for three reasons this PR removes: 1. **A forward nobody used.** `gptq()` ran `max_calibrate` with the calibration forward before its own Hessian forward. Weight amax comes from the weights (`weight_only_quantize`), so the forward only feeds activation quantizers. GPTQ now passes `skip_forward_without_activation_calib=True`, as `max` and `imatrix` already do, and weight-only recipes run one forward instead of two. 2. **Hessians on the CPU while GPUs sat idle.** `GPTQHelper` moves a Hessian to the CPU when its weight's GPU is over 65% used. With the model sharded across 4 GB300s, the GPU holding the current layer was 79% used and two others were empty, so every expert Hessian went to the CPU. Each hook call then copied the whole matrix (26 MB for Qwen) to the host and rescaled it there. A single process now puts such a Hessian on the emptiest GPU it already holds memory on that is under the limit and has the Hessian's bytes free, and uses the CPU only otherwise. Only GPUs the process already uses are candidates, so no other job's GPU is touched; distributed runs keep the old rule, since another GPU may belong to another rank. `update_hessian` also folds its incremental average into one `addmm_` instead of a separate scale pass and add pass over the matrix. 3. **Batch 1.** `hf_ptq.py` picked `--batch_size 1` for layerwise runs because its auto batch-size probe runs a full-model forward. A layerwise batch only holds one layer's activations, so the default is now a token budget: 8,192 tokens, 16 samples at `--calib_seq 512`. HF fused experts run a Python loop over experts on every forward, so per-sample cost is almost flat in batch size: 250 ms per sample at batch 1. 4. **A logged error per expert.** Each expert's GPTQ helper logged its Hessian-weighted error, which costs two weight-by-Hessian matmuls and a host sync: 576 times per GLM-5.3-Flash layer. Expert helpers now skip it, and `gptq()` logs one line per experts module with how many experts it updated and how many calibration never reached. Dense layers keep their line. ### Usage No API change. `hf_ptq.py` users who want the old batching pass `--batch_size 1`: ```bash python hf_ptq.py --pyt_ckpt_path <model> --recipe <layerwise GPTQ recipe> --batch_size 1 ``` ### Testing **Phase timing on one Qwen3.8-Flash-Next-shaped layer**, measured locally on an RTX PRO 6000. The layer has 512 experts, hidden 2560, intermediate 640 and top-10 routing; gate_up is IQ2_XXS GPTQ with 32 samples of 512 tokens: | | GPTQ total | forwards | expert updates | |---|---|---|---| | before, batch 1 | 29.6 s | 2 x 9.8 s | 9.2 s | | after, batch 1 | 19.0 s | 1 x 8.9 s | 9.3 s | | after, batch 16 | 11.2 s | 1 x 1.0 s | 9.3 s | A GLM-5.3-Flash-shaped layer at batch 16 (288 experts, hidden 4096, intermediate 2048, gate_up and down quantized): 21.4 s, of which 19.8 s is the expert updates. **End to end on oci-jhb** (GB300 x4, Qwen3.8-Flash-Next, gate_up IQ2_XXS, layerwise GPTQ, 1,024 x 512 calibration tokens): | Qwen3.8-Flash-Next unless noted | GPTQ time per layer | layerwise calibration (48 layers; GLM 45) | whole job (load, calibrate, KL scoring) | |---|---|---|---| | before (batch 1) | 1,206-1,663 s on linear-attention layers, up to 1,862 s on full-attention ones | ~20 h | ~21 h | | this PR (batch 16) | 29-35 s (first layer 65 s, incl. warm-up) | **44 min** | 70 min | | GLM-5.3-Flash before (batch 1, gate_up + down, 288 experts) | 739-2,131 s | ~25 h (projected) | | | GLM-5.3-Flash, this PR (batch 16) | 37-44 s; 126-136 s on the last 8 MoE layers | **55 min** | 79 min | With only the first and third changes, a layer whose Hessians went to the CPU took 474 s against about 80 s for one whose Hessians stayed on a GPU, which is what the second change fixes. Qwen3.8-Flash-Next's 12 full-attention layers also need a fix outside ModelOpt. transformers' reference sparse-attention indexer (`Qwen4ExpTextQSAIndexer`) loops in Python over every (sample, query) pair, which adds about 340 s per pass over 1,024 x 512 tokens at any batch size. At 512-token calibration, every visible token fits its 2,048-token budget, so its output equals the attention mask it receives. The run above patched the indexer to return that mask directly. The patch is bit-identical to the loop, checked on padded and float masks, and is not part of this PR. Without it, those layers took 374 s each instead of 33 s. KL divergence against BF16 for the same experiment (wikitext-2, 100 x 512 tokens, second half scored, 25,500 tokens; differences paired per chunk): | gate_up IQ2_XXS | mean KL | vs max | |---|---|---| | max | 0.2971 | | | imatrix | 0.2863 | -0.0108 +- 0.0023 | | GPTQ, this PR (batch 16) | **0.2492** | -0.0479 +- 0.0033 (GPTQ - imatrix: -0.0371 +- 0.0034) | On GLM-5.3-Flash (all experts), GPTQ in this PR reaches KL 0.6298 against 0.7211 for imatrix and 0.7402 for max; both models' results are on #2840. On GLM, the last 8 MoE layers took about 3x longer. By then every GPU this process uses was over the 65% limit, so their Hessians went to the CPU. A limit based on absolute headroom would keep them on a GPU: 65% of a 284 GB GB300 leaves about 100 GB unused. That is a possible follow-up. There is no batch-1 KL to compare against: the batch-1 GPTQ runs kept being preempted and were cancelled. Batch size changes results only through floating-point rounding: `hf_ptq.py` pads every calibration row to one length before batching, so batch 1 and batch 16 see the same pad tokens. On phi-2 (NVFP4 W4A4, layerwise, 128 x 512 cnn_dailymail calibration; KL to BF16 on wikitext-2, 40 x 512 tokens, paired per chunk), batch 16 minus batch 1 is -0.0058 +- 0.0032 for GPTQ and +0.0058 +- 0.0052 for local-Hessian, neither significant. **Unit tests:** - The related unit tests (`tests/unit/torch/quantization/{test_gptq,test_ggml_gptq,test_ggml_imatrix,test_layerwise_calibrate,test_calib}.py`, `plugins/test_fused_experts.py`, `tests/unit/recipe`) and `tests/examples/hf_ptq/test_hf_ptq_args.py`: all pass on CPU with transformers 5.18. - New `test_gptq_runs_a_max_calibration_forward_only_for_activation_amax` counts the calibration forwards: 1 for weight-only, 2 with an input quantizer. It fails without the change. - New `test_layerwise_calibration_batch_reaches_the_dataloader` drives `quantize_main` with the layerwise `general/ptq/iq2_xxs` recipe up to the calibration dataloader, without a checkpoint, and checks the batch it receives: 16 at 8,192 tokens, capped at the sample count, 1 for sequences over 8,192 tokens, and an explicit `--batch_size 1` kept as is. It fails with the old batch-1 default. `--calib_seq` must now be positive (`test_calib_seq_must_be_positive`). - New GPU tests, which need 2 GPUs. CI's GPU runner has one and skips them, so I ran them on 2x RTX PRO 6000: - `test_hessian_of_a_full_gpu_goes_to_an_idle_gpu`: a Hessian whose GPU reports over 65% used is allocated on another GPU and accumulates the same values. It fails without the change. - `test_hessian_of_a_full_gpu_falls_back_to_the_cpu`: the Hessian goes to the CPU when the idle GPU is full, when it lacks the Hessian's bytes, and on any rank of a multi-rank run. The too-small case fails without the capacity check. - All 14 tests in `tests/gpu/torch/quantization/test_gptq.py` pass, and all 15 with #2877 on top. - #2840's bit-exact fused-vs-dense GPTQ test still passes. The fused `addmm_` changes the Hessian's float rounding for every GPTQ user, but both paths share it. - Skipping the per-expert error log, on GLM-5.3-Flash (gate_up and down IQ2_XXS, layerwise, 1,024 x 512 tokens, 4x GB300): GPTQ time per MoE layer 41.4 s against 43.5 s before. KL to BF16 does not change: a Qwen3.8-Flash-Next run with this commit gives 0.249200, as before. **Not run:** multi-rank GPTQ end to end (it keeps the old CPU rule), and `hf_ptq.py` layerwise on a dense model at the new default batch. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ for the API. `hf_ptq.py`'s layerwise default batch changes from 1 to 16 at `--calib_seq 512`. Batches pad each sequence to the longest in the batch, and pad tokens feed Hessians and activation amax, so results move slightly, as they already do in non-layerwise runs. `--batch_size 1` restores the old behavior. - 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information The expert updates are now the largest remaining phase on GLM, about 20 s per layer. Encoding one 256-column group for all experts in a single call is bit-exact with per-expert calls and 18x faster per group. Batched Cholesky and batched GEMMs are not bit-exact, so a follow-up would batch only the encoder. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Layerwise calibration now defaults to batches totaling up to 8,192 tokens, limited by available samples. Use `--batch_size 1` for the previous batch size. * GPTQ skips an unnecessary calibration pass when activation quantization does not require one. * With CPU offloading enabled, GPTQ can use a less-utilized GPU for Hessian calculations when sufficient memory is available; otherwise, it uses the CPU. * GPTQ now reports how many expert weights were updated and how many received no calibration tokens. * **Bug Fixes** * Calibration sequence lengths of zero or less are now rejected. * **Documentation** * Updated calibration guidance to explain the default batch size, performance tradeoffs, and batch size used in the documented run. <!-- 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>
What does this PR do?
Type of change: Bug fix
Standalone on
main, split out of the GPTQ-for-MoE stack (#2840 merged ← #2864 ← #2945). It touches only the fused kernel, so it can merge in any order with that stack.GPTQ with
fused: trueruns the NVFP4 column update in a Triton kernel (gptq_fused_block_scalar, in modelopt since 0.45). Triton launches on the current CUDA device, and the wrapper launched without selecting the device of the tensors it was given. A model thatdevice_map="auto"spreads over several GPUs keeps later decoder layers off the current GPU. Once layerwise GPTQ reached such a layer, the kernel ran on GPU 0 with pointers into GPU 1:CUDA error: an illegal memory access was encountered, reported later at an unrelated op.ValueError: Pointer argument cannot be accessed from Triton.The fix launches the kernel under
torch.cuda.device(w_block.device), as the other quantization Triton wrappers (fp4_kernel.py,nvfp4_fp8_sweep.py,bmm2_qdq.py) already do.Found while running layerwise GPTQ on Qwen3.8-Flash-Next with the experts'
down_projin static NVFP4 (fused: true). The model sits on 4 x GB300 with layers 0-22 on GPU 0. Three runs on three different nodes all died right after layer 22.Usage
No API change.
Testing
test_fused_nvfp4_gptq_runs_on_the_weights_gpu(2 GPUs): the fused update with every tensor oncuda:1andcuda:0current matches the same update run oncuda:0bit for bit. It fails without the fix with the Triton pointer error on the PCIe GPUs I ran it on.main,tests/gpu/torch/quantization/test_gptq.pypasses on 2x RTX PRO 6000 (11 tests, including the existing fused-vs-unfused NVFP4 ones).Before your PR is "Ready for review"
CONTRIBUTING.md: N/A🤖 Generated with Claude Code
Summary by CodeRabbit