Intx Quantization Tensor Class #468
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/468
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 1505bca with merge base de4a1fb ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
Not an expert on QAT, but does this mean we would have to enable support for autograd |
torchao/prototype/intx/bitpacking.py
Outdated
| 2 bit shard: [0b00100111, 00010001] | ||
| 4 bit shard: [0b00000000, 0b01101001, 0b10010111, 0b00100101] |
There was a problem hiding this comment.
so this seems to be interleaved packing, that is packing two elements far away together, will packing the adjacent elements together be more efficient because of data locality? I guess this might be covered by setting different pack_dim values, but still want to see if we want to explicitly test this
There was a problem hiding this comment.
Yea it is interleaved. I think its something worth exploring when making optimized kernels. The issue is if you pack adjacent elements then you have to perform interleaved shifting and bit-wise or. Not sure which is faster without using memory/compute profiler.
torchao/prototype/intx/Intx.py
Outdated
| shards = [shard.to(torch.uint8) for shard in shards] | ||
| self.shard = shards | ||
| for i, atrib in enumerate(self.bits_to_shard[bit_size]): | ||
| setattr(self, atrib, shards[i]) |
There was a problem hiding this comment.
should we assert len(shards) == len(bits_to_shard[self.bit_size])?
There was a problem hiding this comment.
Users can only pass in the data to be quantized and the element_size. The pack function (called by the constructor) will automatically create the correct amount of shards based on the element size. I don't really think it would ever break
|
|
||
| aten = torch.ops.aten | ||
|
|
||
| class UintxTensor(torch.Tensor): |
There was a problem hiding this comment.
nit: UIntx feels better I think, same for the filename
There was a problem hiding this comment.
I was thinking about that, but I felt it looks too much like user interface UI
There was a problem hiding this comment.
oh OK, I was mainly concerned about naming convention, but Uintx seems fine for naming convention as well (assuming it means uintx datatype and capitalize u) although it looks a bit weird.
| setattr(self, attrib, shards[i]) | ||
|
|
||
| self.packed_shape = packed_shape | ||
| self.bit_size = bit_size |
There was a problem hiding this comment.
also I feel it would be helpful to add an assert for accepted bit_size, also mayby rename this to bit_width? feels more commen
There was a problem hiding this comment.
What would the assert be checking?
There was a problem hiding this comment.
it should one of 1 to 7 right?
|
merging now since the remaining are minor nits |
PR fulfilling #439
benchmark results:
Performance with dtypes that aren't multiples of 2 is significantly worse, but that was to be expected without custom kernels.