Avoid CUDA indexing failures in iMatrix grid search updates - #2997
Avoid CUDA indexing failures in iMatrix grid search updates#2997yasu-oh wants to merge 2 commits into
Conversation
Replace boolean-indexed in-place updates in the iMatrix grid search with shape-validated torch.where selections. This avoids the CUDA advanced-indexing path that produced inconsistent same-mask selection sizes or IndexKernel out-of-bounds assertions in GPU integration runs, while preserving the intended element-wise update semantics. Signed-off-by: yasu-oh <84763339+yasu-oh@users.noreply.github.com>
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merge Protections🔴 2 of 2 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
🔴 Require two reviewsWaiting for
This rule is failing.PRs labelled "two-reviews" must have at least two approving reviews before merging.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces shape validation checks and updates the tensor assignment logic using torch.where within the _grid_search function of the iMatrix observer, accompanied by a new unit test to verify that shape mismatches raise a RuntimeError. The feedback suggests optimizing this validation by checking only the shape of err on the first iteration, thereby avoiding unnecessary Python overhead in a hot loop.
Check the error shape only on the first grid-search iteration before computing the improvement mask. The remaining update tensor shapes are fixed by construction, so this avoids repeated dictionary allocation and redundant checks in the hot loop. Signed-off-by: yasu-oh <84763339+yasu-oh@users.noreply.github.com>
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
brian-dellabetta
left a comment
brian-dellabetta
left a comment
There was a problem hiding this comment.
Hi @yasu-oh , the changes look fine to me, but could you provide a minimal reproducible example script that triggers this "CUDA IndexKernel.cu out-of-bounds device assertion" error you are reporting? I've not seen it before
|
Thanks. I was able to reproduce the E4B failure again, although it appears to be intermittent rather than deterministic. On the same system and configuration, it reproduced once in six fresh-process runs. In the failed run, the CUDA IndexKernel.cu out-of-bounds assertion was followed by a traceback that surfaced at: in llmcompressor/observers/imatrix.py, ending with a CUDA device-side assert. Because CUDA operations may be asynchronous, I am not yet claiming that this exact assignment is definitively the original source of the failure. I am currently preparing a standalone reproduction script and collecting the relevant logs, including runs with CUDA_LAUNCH_BLOCKING=1, and will share them once they are ready. |
|
A quick update on the E4B case: although I observed the IndexKernel assertion once, I have not been able to reproduce it again after more than 20 additional fresh-process runs under the same conditions. Given the lack of reliable reproducibility, please disregard the E4B-specific report for now. I do not want it to distract from the consistently reproducible E2B failure. I will revisit the E4B case separately if I can reproduce it reliably and collect stronger evidence. |
|
Hi @yasu-oh , so you're seeing it pretty frequently with E2B? do you have a minimal reproduction i can try on my side? |
|
A further update: I can no longer reproduce the E2B failure either with my current environment:
I also noticed that compressed-tensors recently disabled the Triton Given that I can no longer reliably reproduce either the E2B or E4B issue, I don't want to overstate the evidence that the masked assignments themselves were the root cause. That said, I still think the I appreciate you taking the time to review and investigate this. If you think this defensive change is still worthwhile, I would be happy to keep the PR open; if you would prefer not to carry it without a currently reproducible failure, I completely understand and am also happy to close it. |
|
Hi @yasu-oh , thanks for the report. We disabled the triton kernels and changed some of the logic away from torch.where, to fix a regression we were hitting in nvfp4 regression tests. We will re-enable after the releases, can we revisit at that time? It'd be good to see an MRE so we can understand root cause before merging |
2 participants
Summary
Replace boolean-indexed in-place updates in the iMatrix grid search with shape-validated torch.where selections.
The change:
Problem
The current update path can fail during CUDA quantization runs:
python
best_error[improved] = err[improved]
best_min[improved] = shrink_min[improved]
best_max[improved] = shrink_max[improved]
The following failures were observed with GPTQ, NVFP4, and the imatrix_mse observer:
The exact lower-level cause has not been confirmed, so this change does not assume a specific PyTorch defect. The added shape validation did not trigger during successful patched GPU runs.
Validation
The development environment used to prepare this change is CPU-only.
CPU validation:
GPU validation was performed separately using GPTQModifier with NVFP4 imatrix_mse, 1,024 calibration samples, max_seq_length=8192, FP8 static KV-cache calibration, and the basic calibration pipeline for Gemma 4 E models:
The same patched environment was used to produce the Gemma 4 QAT NVFP4 collection covering 31B, 26B-A4B, 12B, E4B, and E2B:
https://huggingface.co/collections/yasu-oh/gemma-4-qat-nvfp4
Trade-off
torch.where materializes full output tensors, so it may use slightly more temporary memory and memory bandwidth than masked in-place assignment.