Supporting MXFP8 All gather for Activations

[avizon-aws]

This issue is related to pytorch/torchtitan#2160 , I was looking into this and I think we might need to update the MXTensor in order to be able to support the MXFP8 all gather for the Linear layers, it will help to improve the E2E model performance. I think we will need to add qdata_dim0 and scale_dim0 to MXTensor as below (We can update the naming as per feedback):

class MXTensor(TorchAOBaseTensor):
    tensor_data_names = ["qdata", "scale", "qdata_dim0", "scale_dim0"]
    tensor_attribute_names = [
        "_elem_dtype",
        "block_size",
        "_orig_dtype",
        "kernel_preference",
        "act_quant_kwargs",
        "_is_swizzled_scales",
    ]


self.qdata_dim0 = None or qdata_dim0  #Default
self.scale_dim0 = None or scale_dim0  #Default

We need something similar to above because the MXLinear layer (inherits from nn.linear) expects only one input in the forward and hence the backward also expects only one grad output. But in order to be able to perform MXFP8 AG we will need:

Quantize Input data -> MXFP8 All Gather -> Matmul

In order for correct backward pass, we need the input to be quantized in the other dimension, for this we need the quantized input in other dim or the original input. But i think saving the original input is wasteful in terms of memory, instead we can save the quantized input in other dim. If we make this an attribute of the MXTensor, then we will not have the need to pass in 2 inputs and we will be able to successfully do the above mentioned operations. The same logic can be applied to the grad_output (we can quantize grad in both dim0 and dim1 and save in the MXTensor, then in the bwd pass, we can transpose the dim0 tensor and use for matmul directly) and then used in the backward. I was able to do a POC and get the MXFP8 Allgather in fwd and bwd for a MXLinear layer as well (needed more changes as well but this was the highlight)

I took a look at Nvidia Transformer Engine, and I see a similar design there as well: https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/tensor/mxfp8_tensor.py#L226-L229

Want to know your thoughts @vkuzo

[vkuzo]

Yeah, for quantization recipes where the scale is not transposeable (rowwise scaling, 1x32 scaling, etc), I would expect that we would need tensor subclasses that store qdata+scale for both dimensions. We have thought about this before, just have not prioritized it yet due to software complexity.

If useful, we do have a set of "leave grad_weight in high precision" recipes where quantizing in both dims is not needed - it's a different tradeoff on the perf vs accuracy curve.

For actually implementing this, I think we should think whether to extend existing MXTensor or create something like MXDim0AndDim1Tensor. I'm currently leaning towards the latter to contain the complexity to only the callsites that need it - this seems to be an optimization for training + TP, so not relevant to training without TP or inference. Curious on your thoughts on this.

[avizon-aws]

I also did consider having a new tensor subclass similar to MXDim0AndDim1Tensor what you mentioned, but the tradeoff is that we will need to reimplement all of the ops again for this tensor subclass similar to MXTensor, there will be repetitive code and will increase maintainability. If we add this functionality to existing MXTensor and keep these additional functionalities encapsulated in functions which dont touch the conventional use case, I think it might be an acceptable solution as well.

What are your thoughts on above.

[vkuzo]

the tradeoff is that we will need to reimplement all of the ops again for this tensor subclass similar to MXTensor, there will be repetitive code and will increase maintainability.

do you have a list of ops / how you envision branching the logic? Having some specifics would help think about the costs of MXTensor / MXDim0Dim1Tensor. In general there doesn't seem to be a very clean solution to solving this, so we'll have to pick the least bad option along various axes such as clean abstractions, complexity, duplication vs reuse, verbosity, etc.

[danielvegamyhre]

interesting discussion, i agree with everything @vkuzo has said so far - i also want to add another possible solution that avoids the need for the bespoke MXDim0Dim1Tensor sublass, which is having a fast fused "dequant along dim0 -> requant along dim1" kernel (via inductor or handwritten) that can be used in the backward pass, so you can just send the normal MXTensor scaled along dim0 in forward all-gather.

there are different trade offs of course (need to evaluate numerics, and perf trade off of "sending 1/2 bytes over NVLink but adding this new kernel overhead", etc). it is worth thinking about though i think.

[avizon-aws]

@vkuzo , I currently dont have a list of ops i know of, but I think the following aten ops might be needed :

1. transpose
2. allgather
3. slicing
4. view
5. clone

The maintainability issues I mentioned previously an example I can think of is e.g. Support for a new aten op/collective e.g. AlltoAll is added for MXTensor, we need a corresponding update to the MXDim0Dim1Tensor.

@danielvegamyhre , I agree with your suggestion, and i had given it a thought as well, the reason I did not mention it before was the dequantization is not lossless (unless I am missing something here). But good to point it out, we can discuss it further and weigh the pros and cons.

[vkuzo]

In general I prefer 2x more LOC of simple code with clean abstractions versus 1x LOC of significantly more complicated code.

Support for a new aten op/collective e.g. AlltoAll is added for MXTensor, we need a corresponding update to the MXDim0Dim1Tensor.

Sounds better than making every callsite of MXTensor know about these extra fields!

[avizon-aws]

I see your point and agree with the point about clean abstractions. I will think again about if anything else is needed and then can discuss further. Will consider MXDim0Dim1Tensor for the further details.