Add GPTQ to prototype #3517
Open
Add GPTQ to prototype #3517
+1,152
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- Add GPTQ quantization algorithm implementation in torchao/prototype/gptq - Add ObserverTensor for activation tracking during calibration - Add unified GPTQConfig that handles both observation and quantization phases - Add gptq_quantize function for weight quantization with Hessian - Add comprehensive test suite for GPTQ and ObserverTensor - Add example script demonstrating GPTQ quantization workflow - All Int4 helper functions are self-contained in gptq module
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3517
Note: Links to docs will display an error until the docs builds have been completed. ❌ 12 New FailuresAs of commit 5b99ce2 with merge base c4273fe ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
jcaip
added
topic: improvement
vkuzo
reviewed
Dec 19, 2025
| @@ -0,0 +1,453 @@ | |||
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |||
Contributor
There was a problem hiding this comment.
nit: put in test/prototype/gptq folder to match codebase convention
vkuzo
reviewed
Dec 19, 2025
| from .observer import ObserverTensor | ||
|
|
||
|
|
||
| @dataclass |
Contributor
There was a problem hiding this comment.
nit: put logic in some_file.py instead of __init__.py
vkuzo
reviewed
Dec 19, 2025
Comment on lines
+160
to
+170
| dead = torch.diag(H) == 0 | ||
| H[dead, dead] = 1 | ||
| W[:, dead] = 0 | ||
|
|
||
| damp = percdamp * torch.mean(torch.diag(H)) | ||
| diag = torch.arange(columns, device=device) | ||
| H[diag, diag] += damp | ||
| H = torch.linalg.cholesky(H) | ||
| H = torch.cholesky_inverse(H) | ||
| H = torch.linalg.cholesky(H, upper=True) | ||
| Hinv = H |
Contributor
There was a problem hiding this comment.
nit: add some comments linking to the right page of the paper
vkuzo
reviewed
Dec 19, 2025
|
|
||
| all_qparams = [] | ||
|
|
||
| for W_quantize_block, block_start in zip( |
Contributor
There was a problem hiding this comment.
nit: add comments as needed to explain how this code is following the algorithm in the paper
vkuzo
reviewed
Dec 19, 2025
| def _calculate_hessian(inputs, device=None): | ||
| """Calculate Hessian matrix from input activations for GPTQ. | ||
|
|
||
| DEPRECATED: This function is kept for backward compatibility in tests only. |
Contributor
There was a problem hiding this comment.
backward compatibility with what?
vkuzo
reviewed
Dec 19, 2025
| from torchao.utils import TorchAOBaseTensor | ||
|
|
||
|
|
||
| class ObserverTensor(TorchAOBaseTensor): |
Contributor
There was a problem hiding this comment.
the logic here is specific to GPTQ, can we ensure the name specifies that
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds GPTQ support to
torchao.prototype.gptq.It is exposed via a new config, GPTQConfig, which can have two steps, "observe" and "convert".
When
quantize_(model, GPTQConfig(step="observe"))is run, observer tensors are attached to the weight tensors, which keep track of linear / torch.bmm ops and updates the hessian matrix based on the observed inputs.When
quantize_(model, GPTQConfig(step="convert"))is run, we will find any observer tensors, take the Hessian and do int4 GPTQ quantization to find the weights. The core of this function is ingptq_quantize.Currently only int4 is hardcoded, but if we enable dequantization and the ability to create a tensor with existing qparams, we should be able to do this for any config.
Also included is an example script,
gptq_example.pythat does sequential / nonsequential quantization on helllaswag for a simple example.