Support for other ops on MXTensor

[avizon-aws]

Hi,
Currently there is support for some ops when using a MXTensor, however its still in the early phases I think, for e.g. if I try to use the operations like concat, chunk on a MXTensor, then it would fail.

I thought about their implementations and I think that we might need an additional attribute to MXTensor i.e. a dim attribute. The motivation for this attribute is to convey which dimension is the contraction dimension (the dimension to be quantized), currently its assumed that -1 (the last dim) is the quantization dimension. However, in some cases e.g. transpose, that might not hold true and if we want to perform an operation like chunk, then the contraction dimension must be known and there are additional restrictions for some ops, an example is shown below.

temp_tensor =MXTensor(....)
torch.chunk(temp_tensor, num_chunks) #We need to ensure that the size of each chunk on a device is a multiple of 32 in the quantization dimension, this is to ensure that the scales are also chunked correctly, in order to do this, we need to know the quantization dimension.

Want to know your thoughts.

[slabhs-aws]

@danielvegamyhre @vkuzo
Is there a roadmap to support these manipulation ops on MXTensor like chunk, split, cat?

[vkuzo]

hi there, sorry for missing this, we monitor issues but didn't know about the github discussions feature.

MXTensor's scale.shape or block_size attributes should be helpful in implementing proper chunk|cat - we should be able to infer the contraction dim by either looking scale.shape or block_size. Let me know your thoughts?

Is there a roadmap to support these manipulation ops on MXTensor like chunk, split, cat?

I think it's ok to add ops on MXTensor as long as they are useful and have well defined semantics similar to their high precision versions. If this is important for your use case, happy to work to either find someone on PyTorch to help, or help review any contributions from your side.

[avizon-aws]

@vkuzo , you are right that we can infer the contraction dim using scale.shape and block size, but I think that would lead to unnecessary if-else, I think it might be better to add an attribute for contraction dim. I think for the quantization function, the contraction dim in the API makes sense. Want to know your thoughts.

[vkuzo]

One thing to consider is square blocked formats such as 32x32 blocks, which we plan to support in the near future. There isn't a single contraction dim in that case.

[avizon-aws]

I think it would also be in line with what we are planning for the quantize aten op in another discussion.

[vkuzo]

right. But we can already represent this with scale.shape and block_size as well. Not sure I buy the case of adding a third variable to the public API for this use case. Seems like the tradeoff is between minimizing BC surface (use existing fields) versus make implementation code somewhat cleaner (introduce new var), minimizing the BC surface seems more attractive to me in this case. Curious if you see it differently?

[vkuzo]

I guess I would add that if there is a consistent design of "contraction_dim" or a similar alternative that everyone agrees upon in PyTorch Core as the best way to represent this, torchao should adopt it to stay consistent with PT Core.

[avizon-aws]

Hmm, I think once the design of the quantize aten op is finalized, we can get back to this discussion. (That aten op has the dim field).

[vkuzo]

sounds good!