[Fix]: Fallback to KleidiAI channelwise kernel groupsize isnt suitable #1647
Conversation
Description: 1. Some models can have certain odd shapes which can not be used with blocked quantization. Fallback to channelwise quantization for those shapes. 2. Fix Formatting issue with experimental tests Signed-off-by: Nikhil Gupta <nikhil.gupta2@arm.com>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1647
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
@digantdesai @metascroy can you please help with reviewing this |
| if torch.backends.kleidiai.is_available(): | ||
| if isinstance(granularity, PerGroup): | ||
| if weight.shape[-1] != group_size and group_size % 32 == 0: |
There was a problem hiding this comment.
weight.shape[-1] != group_size I don't understand this constraint.
IIRC the constraint is k % group_size == 0 and group_size % 32 == 0
| @@ -628,15 +621,28 @@ def apply(weight, bias: Optional[torch.Tensor] = None): | |||
| assert ( | |||
| TORCH_VERSION_AT_LEAST_2_6 | |||
| ), "aten target is requires torch version > 2.6.0" | |||
| # Fallback to Channelwise scheme if group_size is too big | |||
| if weight.shape[-1] < group_size: | |||
There was a problem hiding this comment.
I'd recommend adding a top-level configuration option which clearly tells the user "group size can be changed for certain weight shapes" to handle this case, and throwing an exception if that config setting isn't on.
There was a problem hiding this comment.
I don't understand why we are doing a fallback here. If group_size doesn't divide weight.shape[-1], why not raise exception and let user explicitly move to channelwise?
There was a problem hiding this comment.
Because groupwise is more accurate and we might want to selectively fallback for certain unsupported layers only to maintain accuracy. But I am closing this, user can always filter nn module and use multiple calls to quantizer with correct configuration.
Description: