dtensor-uneven-shard.py

import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import DTensor, Shard

mesh_1d = init_device_mesh("cuda", (4,), mesh_dim_names=("shard",))

rank = mesh_1d.get_rank()

dtensors: list[DTensor] = []

for i in range(8):
    match i % 4 == rank:
        case True:
            dtensor = DTensor.from_local(
                torch.ones(768, 3072),
                device_mesh=mesh_1d,
                placements=[Shard(0)],
                run_check=True,
                shape=(768, 3072),
                stride=(3072, 1),
            )
        case False:
            dtensor = DTensor.from_local(
                torch.ones(0, 3072),
                device_mesh=mesh_1d,
                placements=[Shard(0)],
                run_check=True,
                shape=(768, 3072),
                stride=(3072, 1),
            )
    dtensors.append(dtensor)

print(dtensors[0].full_tensor())