[Float8] Unable to run asyncTP + Float8 row with 'full' AC active, leading dims mismatch

[lessw2020]

Bug description

With the latest PR for Float8 rowwise to support TP, I hit the following error when full checkpointing is active. It is not an issue using no checkpointing.

_fused_scaled_matmul_reduce_scatter_fallback
[rank5]:[E0216](https://www.internalfb.com/servicelab/experiment/0216) 16:37:13.835000 2683263 site-packages/torch/_subclasses/fake_tensor.py:2391] [0/0]     raise ValueError(
[rank5]:[E0216](https://l.workplace.com/l.php?u=https%3A%2F%2Fwww.internalfb.com%2Fservicelab%2Fexperiment%2F0216&h=AT1TX4lglt8Gt6VweeYU3P0oFRLUvsJsvML5Y_yZJMPKCdVzFRdqqmTdvdtqmpXcIGhySXy-kct0PRVIYHpBY3OOFCLjLhFoaBNkUI3dt8eeFEvhGwmkZIm8ArpSAEFcFvqXXHIAHklO_ZuYQ9dt40L_8qfAErr7nWWJgg) 16:37:13.835000 2683263 site-packages/torch/_subclasses/fake_tensor.py:2391] [0/0] ValueError: For row-wise scaling, the leading dims of A_scale must match the leading dims of A (A shape: torch.Size([5, 8192, 4096]), A_scale shape: torch.Size([40960, 1]))
....
rank5]: ValueError: For row-wise scaling, the leading dims of A_scale must match the leading dims of A (A shape: torch.Size([5, 8192, 4096]), A_scale shape: torch.Size([40960, 1]))
[P1734043454]

Versions

Latest TorchAO nightly + TP PR + PyTorch nightly (Feb 10)
TP PR = https://github.com/pytorch/torchtitan/pull/808/files +
pytorch/ao#1718
This was at 256 scale, with full AC
TP=2 + FP8 rowwise

[danielvegamyhre]

@lessw2020 quick question, by "latest PR for rowwise TP" are you referring to the torchtitan PR here #808 or the torchao PR here pytorch/ao#1718? Were you using both?

[danielvegamyhre]

Also feel free to assign this to me

[lessw2020]

Hi @danielvegamyhre - yes, good questions - I was using both and have updated the original issue.
Assigning over to you, thanks for taking a look!

[danielvegamyhre]

Different error occurs at pytorch/ao@32a51ec:

aten.view.default with axiswise scaling and t.shape torch.Size([1, 8192, 4096]) t._scale.shape torch.Size([1, 8192, 1]) t._axiswise_dim -1 new_shape [1, 8192, 4096] is not supported yet.

I believe this was fixed later in pytorch/ao@988c5c9

[danielvegamyhre]

Testing at pytorch/ao@988c5c9 the scale dim mismatch issue repros:

  ValueError: For row-wise scaling, the leading dims of A_scale must match the leading dims of A (A shape: torch.Size([1, 8192, 2048]), A_scale shape: torch.Size([8192, 1]))

[danielvegamyhre]

Tried disabling power of 2 scale factors, issue still occurs.

[danielvegamyhre]

Error: torch._inductor.exc.InductorError: ValueError: For row-wise scaling, the leading dims of A_scale must match the leading dims of A (A shape: torch.Size([1, 8192, 2048]), A_scale shape: torch.Size([8192, 1]))

Error is thrown here when A shape is (1,8192,2048) and A_scale shape is (8192, 1): https://github.com/pytorch/pytorch/blob/fb1f7f6a09576e833647a19b0328d599c754af2a/torch/distributed/_symmetric_memory/__init__.py#L1151

Dumping logs from torchao scale calculations: https://www.internalfb.com/phabricator/paste/view/P1736882650

We can see that we have a tensor shape (1,8192,4096) and scale shape of (1,8192,1). It follows naturally that with async TP degree = 2, this tensor is split along the final dim into 2x of (1,8192,2048), and we keep the same scale of (1,8192,1).

So we would expect to see an A shape of (1,8192,2048) and A_scale shape of (1,8192,1). In the error, we see the A shape is correct but the A scale shape is not.

Looking around the code where the error is thrown, it seems inductor inserts a reduce scatter op (needed for backward presumably) here, the A_scale has had the leading dim squeezed off, and it is (8192,1), which is incorrect, leading to the mismatch between the leading dims of A and A_scale.

I believe at this point in the code the node has the wrong shape for the A_scale: https://github.com/pytorch/pytorch/blob/863ac20659adf63c998a272f7e32e33a38dcea91/torch/_inductor/fx_passes/micro_pipeline_tp.py#L393

Need to install pytorch from source with print statements to debug further.

[danielvegamyhre]

In torch/_inductor/fx_passes i see squeezing going on in a couple places:

• in split_cat.py https://github.com/pytorch/pytorch/blob/fb1f7f6a09576e833647a19b0328d599c754af2a/torch/_inductor/fx_passes/split_cat.py#L412
• in fused SDPA: https://github.com/pytorch/pytorch/blob/fb1f7f6a09576e833647a19b0328d599c754af2a/torch/_inductor/fx_passes/fuse_attention.py#L327

If "split_cat" has to do with concatenation, this seems like it could be related since presumably in TP we need to all-gather and concat at certain points. Wondering if the scaling factor tensor could be affected by this.

[danielvegamyhre]

I think this is where the graph is parsed to find reduce scatters and creates a ReduceScatterMatch for reduce scatters with non-zero dim tensors: https://github.com/pytorch/pytorch/blob/863ac20659adf63c998a272f7e32e33a38dcea91/torch/_inductor/fx_passes/micro_pipeline_tp.py#L265

Presumably we would go down this code path, since the scales tensor is not 0 dim.

I also see in the non zero reduce scatter pattern they're looking for in the graph, the aten.cat.default op is what it starts with, which potentially relates to the squeeze() found in split_cat in #864 (comment)

I'm beginning to wonder if async TP is splitting the scales tensor instead of broadcasting it?

[danielvegamyhre]

Dumped TORCH_LOGS=+aot,+inductor,+dynamo logs (link) and found some relevant lines:

[rank0]:I0220 12:15:11.288000 3845728 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:402] [0/0] [__aot_joint_graph]         exp2: "f32[1, 8192, 1][8192, 1, 1]cuda:0" = torch.ops.aten.exp2.default(floor);  floor = None
...
[rank0]:I0220 12:15:11.288000 3845728 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:402] [0/0] [__aot_joint_graph]         view_4: "f32[8192, 1][1, 1]cuda:0" = torch.ops.aten.view.default(exp2, [-1, 1]);  exp2 = None

Looks like there is a view op which changes the shape of exp2 output from (1,8192,1) to (8192,1)?

[danielvegamyhre]

Confirmed issue still occurs with no AC.

[danielvegamyhre]

This line of code will result in the A_scale shape going from (1,8192,1) -> (8192,1), which would explain the error.

However, the error is actually raised on the line just above it, so it would seem this isn't actually reached?