[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. ❌ 1 New FailureAs of commit 330e2c0 with merge base 07ca637 ( NEW FAILURE - The following job has failed:
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
| @@ -1891,6 +1891,10 @@ def convert(self, model: torch.fx.GraphModule, observer_node: Node): | |||
| assert self.original_dtype is not None, ( | |||
| "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?
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.