Make fp8 compatible with tensor parallelism

[lw]

No description provided.

[lw]

Stack from ghstack (oldest at bottom):

• -> Make fp8 compatible with tensor parallelism #65

[tianyu-l]

I understand that you need to apply such a function in the test of pytorch/pytorch#143760 to "manually" do tiled multiplication/division to compute scaled results.
Here "if b is m x n" only appears when it's DTensor sub-row-wise scaling, in which case the local tensor of b would always have m x 1 shape. So is it correct that:

1. on L38 with local_map we can always assume no tiled multiplication is needed; and
2. on L46 if you're willing to also use local_map, tile multiplication can be avoided too.

[lw]

Thanks for taking a look!

I am learning to use DTensors and I thought it was more idiomatic to express the calculation on the "global" distributed tensor, rather than on the local shard. In order to do so, however, we need to know how many shards there are and reshape accordingly, which arguably isn't that pretty either.

I believe the "fundamental" reason for it is that we're stacking the different components in the wrong order. Here we first replace the matmuls with our custom function, and then we propagate DTensors through it (which means our function needs to know how to handle DTensors). However, I believe the ideal solution would be to first propagate DTensors through some regular matmuls, then take the resulting graph and swap the local matmuls with our function. The issue is that I don't really know how to achieve that, and our code was already written this way before we started supporting DTensors.

(There's also another open question which is how to integrate this with async-TP)

As for local_map, this is currently an unfortunate implementation detail. Ideally the scaling is supposed to be done by the _scaled_mm operator internally, which is what it does! However, because the row-wise scaled-mm is slow (when using slow accum), we use the tensor-wise (un)scaled-mm and do the scaling ourselves. If we were able to make the row-wise scaled-mm faster we could avoid local_map altogether.

[tianyu-l]

Thanks for the analysis! Makes a lot of sense to me!