[PT2E] Fix per-tensor observer issue with varying shape & rank #2177
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2177
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit b03b1e6 with merge base 07ca637 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
Xia-Weiwen
added
the
topic: not user facing
Xia-Weiwen
changed the title
Xia-Weiwen
force-pushed
the
fix_per_tensor_quant
branch
from
87f1249 to
2ac41fb
Compare
| "Expecting original_dtype to be populated" | ||
| ) | ||
| # Since input shape & rank may change (e.g. Resnet18), here we need to update block_size for each input | ||
| self.block_size = get_block_size( |
There was a problem hiding this comment.
when does this happen? can you give an example? I thought using -1 for dynamic dimensions will be enough?
There was a problem hiding this comment.
To reproduce the issue, you may run the code here: #2094 (comment)
You will have to using -1 for block_size without updating of self.block_size here.
There was a problem hiding this comment.
are you saying the rank / number of dimension changes for input as well? can we use a single -1 to represent this case?
There was a problem hiding this comment.
are you saying the rank / number of dimension changes for input as well?
Yes
can we use a single -1 to represent this case?
I think it's doable. But there are checks to guard len(self.block_size) == len(input.shape). We need to handle the special case for per-tensor quant at these locations. Is it ok?
|
@jerryzh168 Could you please review this PR? Thanks. |
|
Hi @jerryzh168 @drisspg Could you please review this PR? I am not sure if the current implementation is what you expected. Thanks. |
| shape_for_reduction: (3, 3, 5, 2, 10) | ||
| reduction_dim: [0, 1, 3, 4] | ||
| """ | ||
| assert len(block_size) == len(input_size) |
There was a problem hiding this comment.
is this still used? we should be using the code in quant_primitives.py I think
There was a problem hiding this comment.
Yes. It's still used when running the prepared model (model after prepare_pt2e). Is it a bug? Do I need to fix it, too?
There was a problem hiding this comment.
I am using the observers defined here: torchao/quantization/pt2e/_affine_quantization.py
There was a problem hiding this comment.
Hi @jerryzh168 May I know your suggestion on this? Thanks.
There was a problem hiding this comment.
I think we should be using the ones in torchao/quantization/observer.py eventually
only occurrence seems to be
ao/test/quantization/pt2e/test_quantize_pt2e.py
Line 2531 in 554cb60
so if you are adding new things I'd recommend use the one from torchao.quantization
There was a problem hiding this comment.
@Xia-Weiwen sorry for the delay, please feel free to work on this
There was a problem hiding this comment.
I thought we already use the one from torchao:
ao/torchao/quantization/pt2e/prepare.py
Line 35 in 212d912
There was a problem hiding this comment.
@jerryzh168 Thanks for the reply. I did not mean torch.ao. I meant there are two versions of such utilities in torchao, torchao.quantization.pt2e and torchao.quantization. For example,
ao/torchao/quantization/pt2e/observer.py
Line 84 in 5153bd3
and
ao/torchao/quantization/observer.py
Line 32 in 5153bd3
The PT2E flow in torchao uses those in
torchao.quantization.pt2e while you said you wanted to switch to torchao/quantization/observer.py.So, I was asking whether you would switch to
torchao/quantization/observer.py in PT2E flow first. Do you have any suggestions on that? Thanks.
There was a problem hiding this comment.
oh I see, yeah for now use torchao/quantization/observer.py would be better I think, we haven't finalized the folder structure for this one yet
There was a problem hiding this comment.
@jerryzh168 Am I supposed to wait until you finalize the folder structure? Thanks.
|
would be good to add a test for this |
|
This PR is stale and looks like it's not an urgent task to switch to Torchao's observer for PT2E. So, let me close it and open a new one if needed in the future. |
3 participants
Fixes #2094 and #2112
We may find inputs with varying shapes and ranks, e.g. when running Resnet18. The current implementation is based on block_size, which is not enough for such cases. The fix is simple: use block_size = -1 for each dimension for per-tensor quantization and update block_size for each input when inserting q/dq in
convert.