[Reland] Unify get_block_size #3059
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3059
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 23c2bc1 with merge base 5e90c47 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
Xia-Weiwen
added
the
topic: not user facing
torchao/quantization/granularity.py
Outdated
| :func:`~torchao.quantization.quant_primitives.quantize_affine` for docs for | ||
| `block_size` | ||
| Attributes: | ||
| block_size (Tuple[int, ...]): The size of each quantization group |
There was a problem hiding this comment.
Thanks. Updated.
torchao/quantization/utils.py
Outdated
| ) | ||
| for i in range(len(block_size)): | ||
| assert input_shape[i] % block_size[i] == 0, ( | ||
| f"Block size {block_size} does not divide input shape {input_shape}" |
There was a problem hiding this comment.
nit: Not all shapes in input shape {input_shape} are divisible by block size {block_size}
There was a problem hiding this comment.
Thanks. Updated.
torchao/quantization/utils.py
Outdated
| return (1,) * (len(input_shape) - 1) + (input_shape[-1],) | ||
| elif isinstance(granularity, PerGroup): | ||
| assert input_shape[-1] % granularity.group_size == 0, ( | ||
| f"Group size {granularity.group_size} does not divide input shape {input_shape}" |
There was a problem hiding this comment.
same for this one Last dimension of input {input_shape[-1]} is not divisible by group size {granularity.group_size}
There was a problem hiding this comment.
Thanks. Updated.
| from .. import Granularity, PerAxis, PerBlock, PerGroup, PerRow, PerTensor, PerToken | ||
| from ..utils import get_block_size |
There was a problem hiding this comment.
this does feel a bit weird, does importing from full path works here?
There was a problem hiding this comment.
Thanks. Updated
Summary
It's a reopen of #3039 but retains backward compatibility. And UTs are not changed to ensure backward compatibility.
This PR merges the two definitions of
get_block_sizein ao/torchao/quantization/pt2e/observer.py and torchao/quantization/observer.py and put it in torchao/quantization/utils.pyIt also removes duplicate definitions of granularities in the pt2e path.
Test plan
This util function is used in many places so it should have been covered by all kinds of UT, especially