Skip to content

[PT2E] Fix per-tensor observer issue with varying shape & rank #2177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
+15 −4
Open

[PT2E] Fix per-tensor observer issue with varying shape & rank #2177

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2ac41fb
[PT2E] Fix per-tensor observer issue with varying shape & rank
Xia-Weiwen May 6, 2025
330e2c0
block_size = [-1] for per-tensor quantization
Xia-Weiwen May 8, 2025
b03b1e6
Refine code
Xia-Weiwen May 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion torchao/quantization/pt2e/observer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1793,7 +1793,7 @@ def get_block_size(
"Please provide an instance of Granularity, not subclass of it"
)
if isinstance(granularity, PerTensor):
return input_shape
return (-1,) * len(input_shape)
elif isinstance(granularity, PerAxis):
block_size = list(input_shape)
block_size[granularity.axis] = 1
Expand Down Expand Up @@ -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(
Copy link
Contributor

@jerryzh168 jerryzh168 May 6, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

when does this happen? can you give an example? I thought using -1 for dynamic dimensions will be enough?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

are you saying the rank / number of dimension changes for input as well? can we use a single -1 to represent this case?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

observer_node.args[0].meta["tensor_meta"].shape, self.granularity
)
if hasattr(self, "is_dynamic") and self.is_dynamic:
choose_qparams_affine = model.graph.call_function(
torch.ops.torchao.choose_qparams_affine,
Expand Down
Loading