Add cachemask variant for fake_quantize_affine #500
Merged
Conversation
This file contains 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
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/500
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit d70f92c with merge base aef7e09 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
jerryzh168
reviewed
Jul 12, 2024
| ) | ||
|
|
||
| Args: | ||
| input (torch.Tensor): original float32, float16 or bfloat16 Tensor |
There was a problem hiding this comment.
we can dedup the comments by: https://github.com/pytorch/pytorch/blob/06ebf87a1eca6c345f5e3e39b63c2ef487695043/torch/ao/quantization/observer.py#L582
e.g.
:func:`~torchao.quantization.quant_primitives.fake_quantize_affine`
?
andrewor14
force-pushed
the
fake_quantize_affine_cachemask
branch
from
4e4dd12
to
ab7f401
Compare
jerryzh168
reviewed
Jul 16, 2024
| General fake quantize op for quantization-aware training (QAT). | ||
| This is equivalent to calling `quantize_affine` + `dequantize_affine` | ||
| but without the dtype casts. | ||
|
|
There was a problem hiding this comment.
please add a section for Args and link to fake_quantize_affine
jerryzh168
reviewed
Jul 16, 2024
| outlier mask for intermediate quantized values | ||
| ) | ||
|
|
||
| Please refer to :func:`~torchao.quantization.quant_primitives.fake_quantize_affine` |
There was a problem hiding this comment.
oh, I think we can move this to before Returns
jerryzh168
approved these changes
Jul 16, 2024
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
andrewor14
force-pushed
the
fake_quantize_affine_cachemask
branch
from
ab7f401
to
d70f92c
Compare
dbyoung18
pushed a commit
to dbyoung18/ao
that referenced
this pull request
Jul 31, 2024
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
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.
Summary: In QAT, we often wish to filter out the gradients corresponding to values outside the expected quantization range, for example:
The existing
fake_quantize_affinereturns 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 tofake_quantize_per_tensor_affine_cachemaskin core.Test Plan:
python test/quantization/test_quant_primitives.py -k test_fake_quantize_affine_cachemask