Skip to content

Launch the fused NVFP4 GPTQ kernel on its tensors' GPU - #2877

Merged
cjluo-nv merged 2 commits into
mainfrom
chenjiel/fix-gptq-fused-kernel-device
Oct 23, 2026
Merged

cjluo-nv merged 2 commits into
mainfrom
chenjiel/fix-gptq-fused-kernel-device

Conversation

@cjluo-nv

@cjluo-nv cjluo-nv commented Oct 19, 2026 •

Copy link
Copy Markdown
Collaborator

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: true runs 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 that device_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:

  • NVLink systems (GB300): CUDA error: an illegal memory access was encountered, reported later at an unrelated op.
  • PCIe systems: 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_proj in 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

  • New GPU test test_fused_nvfp4_gptq_runs_on_the_weights_gpu (2 GPUs): the fused update with every tensor on cuda:1 and cuda:0 current matches the same update run on cuda:0 bit for bit. It fails without the fix with the Triton pointer error on the PCIe GPUs I ran it on.
  • Existing tests: on main, tests/gpu/torch/quantization/test_gptq.py passes on 2x RTX PRO 6000 (11 tests, including the existing fused-vs-unfused NVFP4 ones).
  • End to end: the Qwen3.8-Flash-Next layerwise GPTQ that died after layer 22 completes with this change on oci-jhb; details in the comment below.

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?: ✅
  • Did you update Changelog?: ✅ (bug present since 0.45)
  • Did you get Claude approval on this PR?: ❌

🤖 Generated with Claude Code

Summary by CodeRabbit

  • Bug Fixes
    • Fixed GPTQ fused quantization failures when calibration processes layers located on a different GPU from the current device. Multi-GPU models split automatically across devices can now complete this quantization workflow without CUDA or Triton errors.
    • Confirmed that quantization results remain consistent when processing data on a GPU other than the current device.

@copy-pr-bot

copy-pr-bot Bot commented Oct 19, 2026

Copy link
Copy Markdown

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.

@coderabbitai

coderabbitai Bot commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration
📥 Commits

Reviewing files that changed from the base of the PR and between 1244f17 and 0797396.

📒 Files selected for processing: 1
  • CHANGELOG.rst
🚧 Files skipped from review as they are similar to previous changes: 1
  • CHANGELOG.rst

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

📝 Walkthrough

The fused GPTQ kernel now launches within the CUDA device context of its weight block. A Triton-gated test checks output equality with inputs on a second GPU. The changelog describes the multi-GPU failure scenario.

Changes: GPTQ device context

Layer / File(s) Summary
Select the tensor device and verify multi-GPU behavior
modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py, tests/gpu/torch/quantization/test_gptq.py, CHANGELOG.rst
The kernel launch uses the device of w_block. The test checks exact output equality when inputs are on cuda:1 while cuda:0 remains the current device. The changelog records the multi-GPU failure scenario.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix · Severity of issue fixed: Medium

Suggested labels: cherry-pick-0.48.0

Suggested reviewers: meenchen

Merge Risk: ⚪ Minimal · up to 07973

No actionable issue is identified that would prevent merging after normal checks.

Pre-merge checks: Passed 6
Check name Status Explanation
Description Check Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check Passed The title clearly and concisely describes the main fix: launching the fused NVFP4 GPTQ kernel on the GPU that owns its tensors.
Docstring Coverage Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 2 files. (1 skipped: 1 …
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns Passed PASS. The pull request changes only the fused GPTQ kernel, its GPU test, and the changelog. The added modelopt code only adds a torch.cuda.device(w_block.device) context around the Triton launch. No…
Full details: Docstring Coverage

Explanation

Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 2 files. (1 skipped: 1 unsupported.)

Full details: Security Anti-Patterns

Explanation

PASS. The pull request changes only the fused GPTQ kernel, its GPU test, and the changelog. The added modelopt code only adds a torch.cuda.device(w_block.device) context around the Triton launch. No added torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), # nosec, or dependency changes are present.

✨ Finishing Touches
🧪 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.

@cjluo-nv
cjluo-nv added this pull request to stack #2875 October 19, 2026 20:50
@codecov

codecov Bot commented Oct 19, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.08%. Comparing base (8f435e8) to head (0797396).

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     
Flag Coverage Δ
examples-diffusers 19.94% <0.00%> (-0.06%) ⬇️
examples-gpt-oss 13.47% <0.00%> (-0.03%) ⬇️
examples-hf_ptq 23.38% <0.00%> (-0.12%) ⬇️
examples-llm_distill 13.55% <0.00%> (-0.04%) ⬇️
examples-llm_eval 17.20% <0.00%> (-0.05%) ⬇️
examples-llm_qat 17.37% <0.00%> (-0.09%) ⬇️
examples-llm_sparsity 16.79% <0.00%> (-0.05%) ⬇️
examples-megatron_bridge 26.45% <0.00%> (-0.11%) ⬇️
examples-specdec_bench 13.25% <0.00%> (-0.03%) ⬇️
examples-speculative_decoding 17.45% <0.00%> (-0.11%) ⬇️
examples-torch_onnx 21.26% <0.00%> (-0.09%) ⬇️
examples-torch_trt 15.18% <0.00%> (-0.04%) ⬇️
examples-vllm_serve 17.42% <0.00%> (-0.06%) ⬇️
gpu 58.84% <100.00%> (-1.06%) ⬇️
regression 15.07% <0.00%> (+0.04%) ⬆️
unit 61.03% <0.00%> (-0.01%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@cjluo-nv

Copy link
Copy Markdown
Collaborator Author

End-to-end check on oci-jhb (4 x GB300, device_map="auto"): layerwise GPTQ of Qwen3.8-Flash-Next, routed experts gate_up IQ2 and down_proj static NVFP4 with fused: true, 1,024 calibration samples.

  • Before this fix: three runs, on three different nodes, died right after layer 22 with CUDA error: an illegal memory access was encountered. Layer 23 is the first layer device_map placed on GPU 1.
  • With this fix: IQ2_XXS and IQ2_XS calibrated all 48 layers and exported, about 1 h 7 min each. The IQ2_S run is past layer 43 with no error.

@cjluo-nv
cjluo-nv force-pushed the chenjiel/fix-gptq-fused-kernel-device branch from 7fe0584 to e77ff57 Compare October 21, 2026 06:28
@cjluo-nv
cjluo-nv force-pushed the chenjiel/gptq-speedup branch from b0e76a7 to 79eddeb Compare October 21, 2026 06:28
@github-actions

github-actions Bot commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-23 16:11 UTC

@cjluo-nv
cjluo-nv force-pushed the chenjiel/fix-gptq-fused-kernel-device branch from e77ff57 to 8c958af Compare October 21, 2026 16:23
@cjluo-nv
cjluo-nv force-pushed the chenjiel/gptq-speedup branch from 79eddeb to 38a8fc9 Compare October 21, 2026 16:23
cjluo-nv added a commit that referenced this pull request Oct 22, 2026
**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>
@cjluo-nv
cjluo-nv force-pushed the chenjiel/fix-gptq-fused-kernel-device branch 3 times, most recently from 9b8f2be to fb5079d Compare October 22, 2026 18:53
@cjluo-nv cjluo-nv changed the title [3/3] Launch the fused NVFP4 GPTQ kernel on its tensors' GPU [3/4] Launch the fused NVFP4 GPTQ kernel on its tensors' GPU Oct 22, 2026
@cjluo-nv
cjluo-nv force-pushed the chenjiel/fix-gptq-fused-kernel-device branch from fb5079d to c921d52 Compare October 22, 2026 20:16
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>
@cjluo-nv
cjluo-nv removed this pull request from stack #2875 October 22, 2026 20:35
@cjluo-nv
cjluo-nv force-pushed the chenjiel/fix-gptq-fused-kernel-device branch from c921d52 to 1244f17 Compare October 22, 2026 20:36
@cjluo-nv
cjluo-nv changed the base branch from chenjiel/gptq-speedup to main October 22, 2026 20:36
@cjluo-nv cjluo-nv changed the title [3/4] Launch the fused NVFP4 GPTQ kernel on its tensors' GPU Launch the fused NVFP4 GPTQ kernel on its tensors' GPU Oct 22, 2026
@cjluo-nv
cjluo-nv marked this pull request as ready for review October 22, 2026 20:37
@cjluo-nv
cjluo-nv requested review from a team as code owners October 22, 2026 20:37

@meenchen meenchen left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_scale is 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 on w_block's device.
  • The new test test_fused_nvfp4_gptq_runs_on_the_weights_gpu moves the tensors to cuda:1 while cuda:0 stays current, and checks the result matches the cuda:0 run exactly. It uses the shared need_2_gpus fixture from tests/conftest.py. No existing tests were changed. I read the tests but didn't run them.
  • No licensing changes. The CHANGELOG entry is added.

@cjluo-nv
cjluo-nv enabled auto-merge (squash) October 22, 2026 20:53
@cjluo-nv

Copy link
Copy Markdown
Collaborator Author

/ok to test 1244f17

@cjluo-nv
cjluo-nv merged commit f17be25 into main Oct 23, 2026
54 checks passed
@cjluo-nv
cjluo-nv deleted the chenjiel/fix-gptq-fused-kernel-device branch October 23, 2026 16:11
cjluo-nv added a commit that referenced this pull request Oct 23, 2026
### 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants