Matmul sharding rule incorrectly allows for contraction dim with inconsistent shard order

@H-Huang

🐛 Describe the bug

These two matmuls should be the same but they are not:

import torch
from torch.distributed.tensor import (
    DTensor,
    DeviceMesh,
    distribute_tensor,
    init_device_mesh,
    Partial,
    Replicate,
    Shard,
)
from torch.distributed.tensor.placement_types import _StridedShard
from torch.distributed._local_tensor import (
    local_tensor_mode,
    LocalTensor,
    LocalTensorMode,
)

S = Shard
_SS = _StridedShard

def product(it):
    x = 1
    for i in it:
        x *= i
    return x

def arange_nd(*sizes):
    if len(sizes) == 1 and isinstance(sizes[0], (list, tuple)):
        sizes = sizes[0]
    return torch.arange(product(sizes)).float().view(sizes)

mesh = init_device_mesh("cuda", [2, 2])
A = distribute_tensor(arange_nd(4, 4), mesh, [S(1), S(1)])
B = distribute_tensor(arange_nd(4, 4), mesh, [S(0), S(0)])
print((A @ B).full_tensor())

B2 = distribute_tensor(arange_nd(4, 4), mesh, [_SS(0, split_factor=2), S(0)])
print((A @ B2).full_tensor())

Output:

tensor([[ 56.,  62.,  68.,  74.],
        [152., 174., 196., 218.],
        [248., 286., 324., 362.],
        [344., 398., 452., 506.]], device='cuda:0')
tensor([[ 56.,  62.,  68.,  74.],
        [152., 174., 196., 218.],
        [248., 286., 324., 362.],
        [344., 398., 452., 506.]], device='cuda:1')
tensor([[ 56.,  62.,  68.,  74.],
        [152., 174., 196., 218.],
        [248., 286., 324., 362.],
        [344., 398., 452., 506.]], device='cuda:2')
tensor([[ 56.,  62.,  68.,  74.],
        [152., 174., 196., 218.],
        [248., 286., 324., 362.],
        [344., 398., 452., 506.]], device='cuda:3')
tensor([[ 52.,  58.,  64.,  70.],
        [148., 170., 192., 214.],
        [244., 282., 320., 358.],
        [340., 394., 448., 502.]], device='cuda:1')tensor([[ 52.,  58.,  64.,  70.],
        [148., 170., 192., 214.],
        [244., 282., 320., 358.],
        [340., 394., 448., 502.]], device='cuda:2')

tensor([[ 52.,  58.,  64.,  70.],
        [148., 170., 192., 214.],
        [244., 282., 320., 358.],
        [340., 394., 448., 502.]], device='cuda:3')
tensor([[ 52.,  58.,  64.,  70.],
        [148., 170., 192., 214.],
        [244., 282., 320., 358.],
        [340., 394., 448., 502.]], device='cuda:0')

Versions

main

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @tianyu-l @XilunWu @SherlockNoMad

Concretely, I think we should first try making StridedShard NOT report is_shard() and see what the blast radius is.

[DTensor] Fix matmul operations with _StridedShard placement (pytorch…

@ezyang

Hi @ezyang , I was looking into this to better understand the placement strategies. I propose a fix here #166999 which separates out the StridedShard ("exotic placement") and generates a decomposition strategy which has an explicit cost. I consider it a kind of proof of concept and so am very open to changes.

I believe it to be in alignment with #156217