Add support for int4 weight-only QAT #383
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/383
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 2f91cd3 with merge base aef7e09 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
Cool! Do you think we could also merge a table showing accuracy and speedup in the QAT folder? I'd like to use that in a new training section this week |
Yeah setting up the experiments now, hoping to get those numbers soon |
andrewor14
force-pushed
the
4w_qat
branch
6 times, most recently
from
6d9a205
to
2ac2250
Compare
|
What happened with this PR? Are we still planning on merging it? |
Yes, Jerry and I discussed offline. I plan to refactor the existing quant primitives to make a common |
andrewor14
force-pushed
the
4w_qat
branch
4 times, most recently
from
d56774b
to
6cf40b7
Compare
| else: | ||
| _convert_qat_linear_4w(child) | ||
|
|
||
| class Int4WeightOnlyQATLinear(torch.nn.Linear): |
There was a problem hiding this comment.
are you going to follow up to refactor this to use quantize_ API?
There was a problem hiding this comment.
Yes, that will come later when I refactor QAT to use tensor subclasses
| ) | ||
|
|
||
| Args: | ||
| input (torch.Tensor): original float32, float16 or bfloat16 Tensor |
There was a problem hiding this comment.
please dedup the args, also Args should come before Returns I think
Summary: In QAT, we often wish to filter out the gradients corresponding to values outside the expected quantization range, for example: ``` q = _quantize_affine_no_dtype_cast(...) dq = _dequantize_affine_no_dtype_check(...) mask = torch.logical_and((q >= quant_min), (q <= quant_max)) grad = grad * mask ``` The existing `fake_quantize_affine` returns the dequantized values only, so callers do not have access to this mask. This commit adds the variant to this op that returns both the dequantized values and the mask, similar to `fake_quantize_per_tensor_affine_cachemask` in core. Test Plan: python test/quantization/test_quant_primitives.py -k test_fake_quantize_affine_cachemask
Summary: This commit adds support for int4 weight-only QAT, which simulates the numerics of the existing Int4WeightOnlyQuantizer. The main motivation for this is to provide an end-to-end path for running QAT and lowering to the efficient int4 tinygemm cuda kernel. To enable this, we have to add new fake quantization primitives to match the numerics of the tinygemm kernel, and this required refactoring existing quant primitives to skip dtype casting. Test Plan: python test/quantization/test_qat.py -k test_qat_4w_linear Reviewers: jerryzh168, msaroufim Subscribers: jerryzh168, msaroufim, HDCharles, supriyar
Summary: This commit adds support for int4 weight-only QAT, which simulates the numerics of the existing
Int4WeightOnlyQuantizer. The main motivation for this is to provide an end-to-end path for running QAT and lowering to the efficient int4 tinygemm cuda kernel. To enable this, we have to add new fake quantization primitives to match the numerics of the tinygemm kernel, and this required refactoring existing quant primitives to skip dtype casting.
Test Plan:
python test/quantization/test_qat.py -k test_qat_4w_linear
Reviewers: jerryzh168, msaroufim
Subscribers: jerryzh168, msaroufim, HDCharles, supriyar