xref: /external/pytorch/test/distributed/_shard/sharded_tensor/test_sharded_tensor.py

# Owner(s): ["oncall: distributed"]

import copy
import io
import itertools
import math
import pickle
import sys
from typing import List

import torch
import torch.distributed as dist
from torch.distributed import distributed_c10d, rpc
from torch.distributed._shard import sharded_tensor
from torch.distributed._shard.api import (
    _collect_local_shard,
    _reshard_output,
    _shard_tensor,
    load_with_process_group,
    shard_parameter,
)
from torch.distributed._shard.sharded_tensor import (
    custom_sharded_op_impl,
    pre_load_state_dict_hook,
    Shard,
    ShardedTensor,
    ShardedTensorBase,
    ShardedTensorMetadata,
    state_dict_hook,
)
from torch.distributed._shard.sharded_tensor.api import (
    _create_tensor_from_params,
    TensorProperties,
)
from torch.distributed._shard.sharded_tensor.utils import (
    _parse_and_validate_remote_device,
)
from torch.distributed._shard.sharding_spec import (
    ChunkShardingSpec,
    EnumerableShardingSpec,
    ShardMetadata,
)
from torch.distributed.remote_device import _remote_device
from torch.testing._internal.common_distributed import (
    requires_nccl,
    skip_if_lt_x_gpu,
    spawn_threads_and_init_comms,
    tp_transports,
)
from torch.testing._internal.common_utils import (
    run_tests,
    skip_but_pass_in_sandcastle_if,
    TEST_CUDA,
    TEST_WITH_DEV_DBG_ASAN,
    TestCase,
)
from torch.testing._internal.distributed._shard.sharded_tensor import (
    ShardedTensorTestBase,
    with_comms,
)
from torch.testing._internal.distributed._shard.sharded_tensor._test_st_common import (
    _chunk_sharding_specs_list_for_test,
    MyShardedModel1,
)


if TEST_WITH_DEV_DBG_ASAN:
    print(
        "Skip dev-asan as torch + multiprocessing spawn have known issues",
        file=sys.stderr,
    )
    sys.exit(0)


class TestShardedTensorMetadata(TestCase):
    def test_serialize_and_deserialize(self):
        shard_metadatas = [
            ShardMetadata(
                shard_offsets=[0, 0],
                shard_sizes=[5, 5],
                placement="rank:0/cuda:0",
            ),
            ShardMetadata(
                shard_offsets=[0, 5],
                shard_sizes=[5, 5],
                placement="rank:1/cuda:1",
            ),
            ShardMetadata(
                shard_offsets=[5, 0],
                shard_sizes=[5, 5],
                placement="rank:2/cuda:2",
            ),
            ShardMetadata(
                shard_offsets=[5, 5],
                shard_sizes=[5, 5],
                placement="rank:3/cuda:3",
            ),
        ]

        dtypes = [
            torch.float,
            torch.double,
            torch.cfloat,
            torch.cdouble,
            torch.half,
            torch.bfloat16,
            torch.uint8,
            torch.int8,
            torch.short,
            torch.int,
            torch.long,
            torch.bool,
        ]

        layouts = [torch.strided, torch.sparse_coo]
        requires_grads = [True, False]
        memory_formats = [
            torch.contiguous_format,
            torch.channels_last,
            torch.preserve_format,
        ]
        pin_memories = [True, False]

        for tensor_properties_input in itertools.product(
            dtypes, layouts, requires_grads, memory_formats, pin_memories
        ):
            (
                dtype,
                layout,
                requires_grad,
                memory_format,
                pin_memory,
            ) = tensor_properties_input

            expected_st_metadata = sharded_tensor.ShardedTensorMetadata(
                shard_metadatas,
                (10, 10),
                TensorProperties(
                    dtype, layout, requires_grad, memory_format, pin_memory
                ),
            )

            pickled_obj = pickle.dumps(expected_st_metadata)
            st_metadata = pickle.loads(pickled_obj)
            self.assertEqual(expected_st_metadata, st_metadata)


class TestCreateTensorFromParams(TestCase):
    @skip_but_pass_in_sandcastle_if(not TEST_CUDA, "CUDA GPU is needed")
    def test_empty(self):
        expected_dtype = torch.double
        tensor_properties = TensorProperties(
            dtype=expected_dtype,
            layout=torch.strided,
            requires_grad=False,
            pin_memory=False,
            memory_format=torch.contiguous_format,
        )
        local_device = torch.device("cuda:0")
        local_tensor = _create_tensor_from_params(
            5, 10, local_device=local_device, tensor_properties=tensor_properties
        )
        self.assertEqual(local_device, local_tensor.device)
        self.assertEqual(expected_dtype, local_tensor.dtype)
        self.assertEqual(torch.strided, local_tensor.layout)
        self.assertEqual(False, local_tensor.requires_grad)


class TestShardParameter(ShardedTensorTestBase):
    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_shard_parameter(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        fc = torch.nn.Linear(12, 12).cuda(self.rank)
        weight_og = fc.weight.clone()
        shard_parameter(fc, "weight", spec)

        # Verify.
        self.assertTrue(isinstance(fc.weight, ShardedTensor))
        local_shards = fc.weight.local_shards()
        self.assertEqual(1, len(local_shards))
        self.assertEqual(torch.Size([3, 12]), local_shards[0].tensor.size())
        self.assertEqual(3, local_shards[0].tensor.size(0))
        self.assertEqual(12, local_shards[0].tensor.size(1))
        self.assertEqual(
            torch.narrow(weight_og, 0, 3 * self.rank, 3), local_shards[0].tensor
        )

    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_shard_parameter_errors(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        fc = torch.nn.Linear(12, 12).cuda(self.rank)
        with self.assertRaisesRegex(ValueError, "does not match with src_rank"):
            shard_parameter(fc, "weight", spec, src_rank=self.rank)

        with self.assertRaisesRegex(AttributeError, "has no attribute"):
            shard_parameter(fc, "foo", spec)

        with self.assertRaisesRegex(
            ValueError, "Expected Linear.bias to be a Tensor, but found str"
        ):
            del fc.bias
            fc.bias = "foo"
            shard_parameter(fc, "bias", spec)

        with self.assertRaisesRegex(ValueError, "not a contiguous Tensor"):
            fc.bias = torch.rand(10, 10).cuda(self.rank).t()
            shard_parameter(fc, "bias", spec)

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                f"rank:{self.rank}/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        with self.assertRaisesRegex(ValueError, "does not match with sharding_spec"):
            shard_parameter(fc, "weight", spec)

        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
            ]
        )
        with self.assertRaisesRegex(NotImplementedError, "not implemented yet!"):
            shard_parameter(fc, "weight", spec)


class TestShardTensor(ShardedTensorTestBase):
    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_shard_tensor(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        tensor = torch.rand(12, 12).cuda(self.rank)
        st = _shard_tensor(tensor, spec)

        # Verify.
        self.assertTrue(isinstance(st, sharded_tensor.ShardedTensor))
        local_shard = st.local_tensor()
        self.assertEqual(1, len(st.local_shards()))
        self.assertEqual(torch.Size([3, 12]), local_shard.size())
        self.assertEqual(torch.narrow(tensor, 0, 3 * self.rank, 3), local_shard)

    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_shard_tensor_with_empty_shard(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        tensor = torch.rand(9, 12).cuda(self.rank)
        st = _shard_tensor(tensor, spec)

        # Verify.
        self.assertTrue(isinstance(st, sharded_tensor.ShardedTensor))
        sms = st.metadata().shards_metadata
        self.assertEqual(len(sms), 4)
        for sm in sms:
            self.assertTrue(sm.shard_offsets[0] + sm.shard_sizes[0] <= tensor.size(0))

        local_shard = st.local_tensor()
        self.assertEqual(1, len(st.local_shards()))
        if dist.get_rank() < 3:
            self.assertEqual(torch.Size([3, 12]), local_shard.size())
            self.assertEqual(torch.narrow(tensor, 0, 3 * self.rank, 3), local_shard)
        else:
            self.assertEqual(torch.Size([0, 12]), local_shard.size())

    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_shard_tensor_errors(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        tensor = torch.rand(12, 12).cuda(self.rank)

        with self.assertRaisesRegex(ValueError, "does not match with src_rank"):
            _shard_tensor(tensor, spec, src_rank=self.rank)

        with self.assertRaisesRegex(ValueError, "not a contiguous Tensor"):
            tensor_t = torch.rand(12, 12).cuda(self.rank).t()
            _shard_tensor(tensor_t, spec)

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                f"rank:{self.rank}/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        with self.assertRaisesRegex(ValueError, "does not match with sharding_spec"):
            _shard_tensor(tensor, spec)

        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
            ]
        )
        with self.assertRaisesRegex(NotImplementedError, "not implemented yet!"):
            _shard_tensor(tensor, spec)


class TestModuleHookApi(ShardedTensorTestBase):
    class DummyNNModule(torch.nn.Module):
        def __init__(self, spec, tensor_size):
            super().__init__()
            self.st = sharded_tensor.rand(spec, *tensor_size)

        def forward(self):
            return self.st

    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_reshard_output(self):
        specs = _chunk_sharding_specs_list_for_test([0, 1], seed=5)
        spec, reshard_spec = specs[0], specs[1]
        test_module = self.DummyNNModule(spec, [24, 12])
        st = test_module()
        local_shard = st.local_tensor()
        pg = dist.distributed_c10d._get_default_group()
        st_compare = ShardedTensor._init_from_local_shards(
            copy.deepcopy(st.local_shards()),
            st.size(),
            process_group=pg,
        )
        st_compare._sharding_spec = copy.deepcopy(spec)
        st_compare.reshard(reshard_spec)
        test_module = _reshard_output(test_module, reshard_spec)
        st = test_module()
        local_shard = st.local_tensor()
        local_shard_compare = st_compare.local_tensor()
        self.assertEqual(local_shard, local_shard_compare)
        self.assertEqual(local_shard.size(0), 24)
        self.assertEqual(local_shard.size(1), 3)

    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_collect_local_shard(self):
        specs = _chunk_sharding_specs_list_for_test([0], seed=5)
        spec = specs[0]
        test_module = self.DummyNNModule(spec, [23, 15])
        st = test_module()
        local_shard = st.local_tensor()
        test_module = _collect_local_shard(test_module)
        output = test_module()
        self.assertTrue(isinstance(output, torch.Tensor))
        self.assertEqual(local_shard, output)


class TestLocalTensor(ShardedTensorTestBase):
    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_local_tensor(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        st = sharded_tensor.rand(spec, 24, 12)
        local_shard = st.local_tensor()
        self.assertEqual(torch.Size([6, 12]), local_shard.size())
        self.assertEqual(st.local_tensor(), local_shard)

    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_local_tensor_error(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:1/cuda:1",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:2/cuda:2",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
                "rank:3/cuda:3",
            ],
        )
        st = sharded_tensor.rand(spec, 24, 12)
        with self.assertRaisesRegex(
            NotImplementedError, "Only single local shard is supported."
        ):
            local_shard = st.local_tensor()


class TestShardedTensorChunked(ShardedTensorTestBase):
    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_sharded_tensor_metadata(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        st = sharded_tensor.empty(spec, 10, 20, init_rrefs=True)
        st_metadata = st.metadata()
        self.assertEqual(torch.Size([10, 20]), st_metadata.size)
        self.assertEqual(torch.Size([10, 20]), st.size())
        self.assertEqual(torch.float, st.dtype)
        self.assertEqual(torch.strided, st.layout)
        self.assertEqual(False, st.requires_grad)
        self.assertTrue(st.is_contiguous())
        self.assertFalse(st.is_pinned())

        st = sharded_tensor.empty(spec, 10, 20, requires_grad=True, init_rrefs=True)
        self.assertEqual(True, st.requires_grad)

        st = sharded_tensor.empty(spec, 10, 20, dtype=torch.double, init_rrefs=True)
        self.assertEqual(torch.double, st.dtype)

        # Need CPU for pin_memory
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cpu",
                "rank:1/cpu",
                "rank:2/cpu",
                "rank:3/cpu",
            ],
        )

        st = sharded_tensor.empty(spec, 10, 20, pin_memory=True, init_rrefs=True)
        self.assertEqual(True, st.is_pinned())

        # test read only properties, they're read only as we can't simply change
        # the global metadata without changing the underlying shard's properties
        with self.assertRaisesRegex(RuntimeError, "torch function '__set__'"):
            st.requires_grad = True

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_complete_world_size(self):
        for dim in [0, -2]:
            spec = ChunkShardingSpec(
                dim=dim,
                placements=[
                    "rank:0/cuda:0",
                    "rank:1/cuda:1",
                    "rank:2/cuda:2",
                    "rank:3/cuda:3",
                ],
            )
            st = sharded_tensor.empty(spec, 10, 20, init_rrefs=True)

            # Validate local shard.
            local_shards = st.local_shards()
            self.assertEqual(1, len(local_shards))
            local_shard = local_shards[0].tensor
            self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
            if self.rank == 3:
                self.assertEqual((1, 20), local_shard.size())
            else:
                self.assertEqual((3, 20), local_shard.size())

            # Validate global metadata.
            st_metadata = st.metadata()
            shards_metadata = st_metadata.shards_metadata
            self.assertEqual(4, len(shards_metadata))

            for rank, shard_metadata in enumerate(shards_metadata):
                self.assertEqual([rank * 3, 0], shard_metadata.shard_offsets)
                if rank == 3:
                    self.assertEqual([1, 20], shard_metadata.shard_sizes)
                else:
                    self.assertEqual([3, 20], shard_metadata.shard_sizes)
                self.assertEqual(
                    f"rank:{rank}/cuda:{rank}", str(shard_metadata.placement)
                )

            # Validate remote shards.
            remote_shards = st.remote_shards()
            self.assertEqual(3, len(remote_shards))

            for rpc_rank, shards in remote_shards.items():
                self.assertEqual(1, len(shards))
                for remote_shard in shards:
                    self.assertEqual(rpc_rank, remote_shard.owner().id)
                    shard = remote_shard.to_here()
                    self.assertEqual(
                        f"rank:{rpc_rank}/cuda:{rpc_rank}",
                        str(shard.metadata.placement),
                    )
                    if rpc_rank == 3:
                        self.assertEqual((1, 20), shard.tensor.size())
                    else:
                        self.assertEqual((3, 20), shard.tensor.size())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_create_sharded_tensor_with_ones(self):
        """Test sharded_tensor.ones(...)"""

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 10, 20
        st = sharded_tensor.ones(spec, h, w)

        # Validate local shard is initialized with torch.ones
        local_shards = st.local_shards()
        self.assertEqual(1, len(local_shards))
        local_shard = local_shards[0].tensor
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
        # The split: for rank!=3 ceil(h/4)=3  for rank=3 1
        expected_h = 1 if self.rank == 3 else math.ceil(h / 4)
        self.assertEqual((expected_h, w), local_shard.size())
        self.assertEqual(local_shard, torch.ones(expected_h, w))

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_gather_even(self) -> None:
        """Test _sharded_tensor.gather(...) with evenly distributed._shards"""

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 10, 20
        st = sharded_tensor.ones(spec, h, w)

        full_tensor = None
        dst = 1
        if self.rank == dst:
            full_tensor = torch.zeros(
                h,
                w,
                device=torch.device(f"cuda:{dst}"),
            )
        st.gather(dst, full_tensor)

        if self.rank == dst:
            self.assertEqual(full_tensor, torch.ones(h, w))
        else:
            self.assertIsNone(full_tensor)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_gather_uneven(self) -> None:
        """Test _sharded_tensor.gather(...) with unevenly distributed._shards"""

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
            ],
        )
        h, w = 10, 20
        st = sharded_tensor.ones(spec, h, w)

        full_tensor = None
        dst = 1
        if self.rank == dst:
            full_tensor = torch.zeros(
                h,
                w,
                device=torch.device(f"cuda:{dst}"),
            )
        st.gather(dst, full_tensor)

        if self.rank == dst:
            self.assertEqual(full_tensor, torch.ones(h, w))
        else:
            self.assertIsNone(full_tensor)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_create_sharded_tensor_with_zeros(self):
        """Test sharded_tensor.zeros(...)"""

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 10, 20
        st = sharded_tensor.zeros(spec, h, w)

        # Validate local shard is initialized with torch.zeros
        local_shards = st.local_shards()
        self.assertEqual(1, len(local_shards))
        local_shard = local_shards[0].tensor
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
        # The split: for rank!=3 ceil(h/4)=3  for rank=3 1
        expected_h = 1 if self.rank == 3 else math.ceil(h / 4)
        self.assertEqual((expected_h, w), local_shard.size())
        self.assertEqual(local_shard, torch.zeros(expected_h, w))

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_create_sharded_tensor_with_rand(self):
        """Test sharded_tensor.rand(...)/randn(...)"""

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 8, 2
        seed = 1234

        expected_h = 2
        expected_device = torch.device(f"cuda:{self.rank}")
        dtype = torch.double
        torch.manual_seed(seed)
        # Test sharded_tensor.rand creation
        expected = torch.rand(expected_h, w, device=expected_device, dtype=dtype)
        # reset seed to ensure the same random numbers are generated
        torch.manual_seed(seed)
        st = sharded_tensor.rand(spec, h, w, dtype=dtype)

        # Validate local shard is initialized with torch.rand
        local_shards = st.local_shards()
        self.assertEqual(1, len(local_shards))
        local_shard = local_shards[0].tensor
        self.assertEqual(expected_device, local_shard.device)
        self.assertEqual((expected_h, w), local_shard.size())
        self.assertEqual(expected, local_shard)

        # Test sharded_tensor.randn creation
        torch.manual_seed(seed)
        expected_randn = torch.randn(expected_h, w, device=expected_device, dtype=dtype)
        # reset seed to ensure the same random numbers are generated
        torch.manual_seed(seed)
        st_randn = sharded_tensor.randn(spec, h, w, dtype=dtype)

        # Validate local shard is initialized with torch.randn
        local_shards = st_randn.local_shards()
        self.assertEqual(1, len(local_shards))
        local_shard = local_shards[0].tensor
        self.assertEqual(expected_device, local_shard.device)
        self.assertEqual((expected_h, w), local_shard.size())
        self.assertEqual(expected_randn, local_shard)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_create_sharded_tensor_with_full(self):
        """Test sharded_tensor.full(...)"""

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 10, 20
        fill_value = 1234
        st = sharded_tensor.full(
            spec, size=(h, w), fill_value=fill_value, dtype=torch.int32
        )

        # Validate local shard is initialized with torch.full
        local_shards = st.local_shards()
        self.assertEqual(1, len(local_shards))
        local_shard = local_shards[0].tensor
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
        # The split: for rank!=3 ceil(h/4)=3  for rank=3 1
        expected_h = 1 if self.rank == 3 else math.ceil(h / 4)
        self.assertEqual((expected_h, w), local_shard.size())
        self.assertEqual(
            local_shard,
            torch.full(size=(expected_h, w), fill_value=fill_value, dtype=torch.int32),
        )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_create_sharded_tensor_like(self):
        """Test tensor like methods, i.e. torch.zeros_like(...), torch.full_like, etc."""

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 8, 8
        expected_h = 2
        seed = 1234
        dtype = torch.double
        expected_device = torch.device(f"cuda:{self.rank}")
        st = sharded_tensor.rand(spec, (h, w), dtype=dtype)
        tensor_like_ops = {
            torch.zeros_like: torch.zeros,
            torch.ones_like: torch.ones,
            torch.rand_like: torch.rand,
            torch.randn_like: torch.randn,
            torch.empty_like: torch.empty,
            torch.full_like: torch.full,
        }
        for op, expect_local_op in tensor_like_ops.items():
            if op == torch.full_like:
                # special handle full/full_like as it needs to have additional fill_value arg
                expect_tensor = expect_local_op(
                    (expected_h, w), 8.8, device=expected_device, dtype=dtype
                )
                new_op_st = op(st, 8.8, dtype=dtype)
                self.assertEqual(new_op_st.local_tensor(), expect_tensor)
            elif op == torch.empty_like:
                # empty/empty_like we only compare the shape
                expect_tensor = expect_local_op(
                    expected_h, w, device=expected_device, dtype=dtype
                )
                new_op_st = op(st, dtype=dtype)
                self.assertEqual(new_op_st.local_tensor().shape, expect_tensor.shape)
            else:
                torch.manual_seed(seed)
                expect_tensor = expect_local_op(
                    expected_h, w, device=expected_device, dtype=dtype
                )
                torch.manual_seed(seed)
                new_op_st = op(st, dtype=dtype)
                self.assertEqual(new_op_st.local_tensor(), expect_tensor)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_partial_world_size(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        st = sharded_tensor.empty(spec, 10, 20, init_rrefs=True)

        # Validate local shard.
        local_shards = st.local_shards()
        if self.rank >= 2:
            self.assertEqual(1, len(local_shards))
            local_shard = local_shards[0].tensor
            self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
            self.assertEqual((5, 20), local_shard.size())
        else:
            self.assertEqual(0, len(local_shards))

        # Validate global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(2, len(shards_metadata))

        for shard_rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual([shard_rank * 5, 0], shard_metadata.shard_offsets)
            self.assertEqual([5, 20], shard_metadata.shard_sizes)
            self.assertEqual(
                f"rank:{shard_rank + 2}/cuda:{shard_rank + 2}",
                str(shard_metadata.placement),
            )

        # Validate remote shards.
        remote_shards = st.remote_shards()
        if self.rank >= 2:
            self.assertEqual(1, len(remote_shards))
        else:
            self.assertEqual(2, len(remote_shards))

        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))
            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                shard = remote_shard.to_here()
                self.assertEqual(
                    f"rank:{rpc_rank}/cuda:{rpc_rank}", str(shard.metadata.placement)
                )
                self.assertEqual((5, 20), shard.tensor.size())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_new_group(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        pg = dist.new_group(ranks=[1, 2, 3])
        st = sharded_tensor.empty(spec, 10, 20, process_group=pg, init_rrefs=True)

        # Validate local shard.
        local_shards = st.local_shards()
        if self.rank >= 2:
            self.assertEqual(1, len(local_shards))
            local_shard = local_shards[0].tensor
            self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
            self.assertEqual((5, 20), local_shard.size())
        else:
            self.assertEqual(0, len(local_shards))

        # Validate global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(2, len(shards_metadata))

        for shard_rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual([shard_rank * 5, 0], shard_metadata.shard_offsets)
            self.assertEqual([5, 20], shard_metadata.shard_sizes)
            self.assertEqual(
                f"rank:{shard_rank + 2}/cuda:{shard_rank + 2}",
                str(shard_metadata.placement),
            )

        # Validate remote shards.
        remote_shards = st.remote_shards()
        if self.rank >= 2:
            self.assertEqual(1, len(remote_shards))
        else:
            self.assertEqual(2, len(remote_shards))

        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))
            for remote_shard in shards:
                shard = remote_shard.to_here()
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                self.assertEqual(
                    f"rank:{rpc_rank}/cuda:{rpc_rank}", str(shard.metadata.placement)
                )
                self.assertEqual((5, 20), shard.tensor.size())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_multiple_local_shards(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        st = sharded_tensor.empty(spec, 16, 20, init_rrefs=True)

        # Validate local shards.
        local_shards = st.local_shards()
        self.assertEqual(2, len(local_shards))
        for local_shard in local_shards:
            self.assertEqual(
                torch.device(f"cuda:{self.rank}"), local_shard.tensor.device
            )
            self.assertEqual((2, 20), local_shard.tensor.size())

        # Validate global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(8, len(shards_metadata))

        for shard_idx, shard_metadata in enumerate(shards_metadata):
            self.assertEqual([shard_idx * 2, 0], shard_metadata.shard_offsets)
            self.assertEqual([2, 20], shard_metadata.shard_sizes)
            self.assertEqual(
                f"rank:{shard_idx % 4}/cuda:{shard_idx % 4}",
                str(shard_metadata.placement),
            )

        # Validate remote shards.
        remote_shards = st.remote_shards()
        self.assertEqual(3, len(remote_shards))
        owners = {}
        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(2, len(shards))
            for remote_shard in shards:
                shard = remote_shard.to_here()
                self.assertEqual((2, 20), shard.tensor.size())
                self.assertEqual(rpc_rank, remote_shard.owner().id)

    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_sharding_columns(self):
        self.init_pg()

        for dim in [1, -1]:
            spec = ChunkShardingSpec(
                dim=dim,
                placements=[
                    "rank:0/cuda:0",
                    "rank:1/cuda:1",
                    "rank:2/cuda:2",
                    "rank:3/cuda:3",
                ],
            )

            st = sharded_tensor.empty(spec, 10, 32)

            # Validate local shard.
            local_shards = st.local_shards()
            self.assertEqual(1, len(local_shards))
            local_shard = local_shards[0].tensor
            self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
            self.assertEqual((10, 8), local_shard.size())

            # Validate global metadata.
            st_metadata = st.metadata()
            shards_metadata = st_metadata.shards_metadata
            self.assertEqual(4, len(shards_metadata))

            for rank, shard_metadata in enumerate(shards_metadata):
                self.assertEqual([0, rank * 8], shard_metadata.shard_offsets)
                self.assertEqual([10, 8], shard_metadata.shard_sizes)
                self.assertEqual(
                    f"rank:{rank}/cuda:{rank}", str(shard_metadata.placement)
                )

    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_invalid_sharding(self):
        self.init_pg()

        with self.assertRaisesRegex(
            NotImplementedError, "does not support named dimension"
        ):
            spec = ChunkShardingSpec(dim="H", placements=["rank:1/cuda:1"])
            sharded_tensor.empty(spec, 10, 20)

        for dim in [2, 3, 4, -3, -4, -5]:
            spec = ChunkShardingSpec(dim=dim, placements=["rank:1/cuda:1"])
            with self.assertRaisesRegex(ValueError, "Invalid sharding dim"):
                sharded_tensor.empty(spec, 10, 20)

        spec = ChunkShardingSpec(dim=0, placements=["rank:5/cuda:1"])
        with self.assertRaisesRegex(
            ValueError, "Global rank 5 does not exist in input process group"
        ):
            sharded_tensor.empty(spec, 10, 20)

        spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"])
        st = sharded_tensor.empty(spec, 10, 20)
        tensor = torch.empty(10, 20)
        with self.assertRaisesRegex(
            RuntimeError, r".*not supported for ShardedTensor!$"
        ):
            torch.add(st, tensor)

        spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"])
        with self.assertRaisesRegex(
            ValueError, "Only torch.strided layout is currently supported"
        ):
            sharded_tensor.empty(spec, 10, 20, layout=torch.sparse_coo)

        spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"])
        with self.assertRaisesRegex(
            ValueError,
            "Only torch.contiguous_format memory_format is currently supported",
        ):
            sharded_tensor.empty(spec, 10, 20, memory_format=torch.channels_last)

        spec = ChunkShardingSpec(dim=0, placements=["worker0/cuda:1"])
        with self.assertRaisesRegex(
            RuntimeError, "RPC framework needs to be initialized"
        ):
            sharded_tensor.empty(spec, 10, 20)

        spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"])
        with self.assertRaisesRegex(
            RuntimeError, "RPC Framework needs to be initialized"
        ):
            st = sharded_tensor.empty(spec, 10, 20, init_rrefs=True)

        with self.assertRaisesRegex(
            RuntimeError, "ShardedTensor created with init_rrefs=False"
        ):
            st = sharded_tensor.empty(spec, 10, 20)
            st.remote_shards()

        self.init_rpc()
        spec = ChunkShardingSpec(dim=0, placements=["workerfoo/cuda:1"])
        with self.assertRaisesRegex(ValueError, "Invalid worker name"):
            sharded_tensor.empty(spec, 10, 20, init_rrefs=True)

    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_invalid_pg_rpc_ranks(self):
        self.init_pg()

        # Init RPC with different ranks.
        rpc_backend_options = rpc.TensorPipeRpcBackendOptions(
            _transports=tp_transports()
        )
        rpc_backend_options.init_method = f"file://{self.file_name}"
        rank = (self.rank + 1) % self.world_size
        rpc.init_rpc(
            name=f"worker{rank}",
            rank=rank,
            world_size=self.world_size,
            rpc_backend_options=rpc_backend_options,
        )

        spec = ChunkShardingSpec(dim=0, placements=["rank:1/cuda:1"])
        with self.assertRaisesRegex(
            ValueError, "Default ProcessGroup and RPC ranks must be the same"
        ):
            sharded_tensor.empty(spec, 10, 20, init_rrefs=True)

    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_insufficient_sharding_dims(self):
        self.init_pg()

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        st = sharded_tensor.empty(spec, 2, 20)

        # Validate local shard.
        local_shards = st.local_shards()
        if self.rank <= 1:
            self.assertEqual(1, len(local_shards))
            local_shard = local_shards[0].tensor
            self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
            self.assertEqual((1, 20), local_shard.size())
        else:
            self.assertEqual(1, len(local_shards))
            local_shard = local_shards[0].tensor
            self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.device)
            self.assertEqual(local_shard.numel(), 0)

        # Validate global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(4, len(shards_metadata))

        for shard_rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual([shard_rank, 0], shard_metadata.shard_offsets)
            self.assertEqual(
                f"rank:{shard_rank}/cuda:{shard_rank}", str(shard_metadata.placement)
            )
            if shard_rank <= 1:
                self.assertEqual([1, 20], shard_metadata.shard_sizes)
            else:
                self.assertEqual([0, 20], shard_metadata.shard_sizes)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_sharded_tensor_sizes(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        # Test with *args
        st = sharded_tensor.empty(spec, 10, 20, init_rrefs=True)
        self.assertEqual(torch.Size([10, 20]), st.size())

        # Test with single *args
        st = sharded_tensor.empty(spec, 10, init_rrefs=True)
        self.assertEqual(torch.Size([10]), st.size())

        # Test with list
        st = sharded_tensor.empty(spec, [10, 20], init_rrefs=True)
        self.assertEqual(torch.Size([10, 20]), st.size())

        # Test with tuple
        st = sharded_tensor.empty(spec, (10, 20), init_rrefs=True)
        self.assertEqual(torch.Size([10, 20]), st.size())

        # Test with row size
        st = sharded_tensor.empty(spec, (10, 20), init_rrefs=True)
        self.assertEqual(st.size(0), 10)

        # Test with col size
        st = sharded_tensor.empty(spec, (10, 20), init_rrefs=True)
        self.assertEqual(st.size(1), 20)

        # Test with negative indexed size
        st = sharded_tensor.empty(spec, (10, 20), init_rrefs=True)
        self.assertEqual(st.size(-1), 20)

        # Test with dim/ndim
        self.assertEqual(st.dim(), 2)
        self.assertEqual(st.ndim, 2)
        # Test with invalid input
        st = sharded_tensor.empty(spec, (10, 20), init_rrefs=True)
        with self.assertRaisesRegex(IndexError, "Dimension out of range"):
            st.size(-3)
        with self.assertRaisesRegex(IndexError, "Dimension out of range"):
            st.size(2)

        with self.assertRaises(TypeError):
            st = sharded_tensor.empty(spec, "foo")

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_state_dict(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        m = MyShardedModel1(spec)

        # Test save
        m._register_state_dict_hook(state_dict_hook)
        buffer = io.BytesIO()
        mod_state_dict = m.state_dict()
        mod_state_keys = mod_state_dict.keys()
        self.assertTrue("sharded_tensor1" in mod_state_keys)
        self.assertTrue("submodule.sharded_tensor2" in mod_state_keys)
        torch.save(mod_state_dict, buffer)

        # Test load.
        module_load = MyShardedModel1()
        module_load._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True)

        buffer.seek(0)
        state_dict_deser = torch.load(buffer)
        module_load.load_state_dict(state_dict_deser, strict=False)

        module_load._register_state_dict_hook(state_dict_hook)
        loaded_dict_keys = module_load.state_dict().keys()
        self.assertTrue("sharded_tensor1" in loaded_dict_keys)
        self.assertTrue("submodule.sharded_tensor2" in loaded_dict_keys)
        # Verify after load.
        self.assertTrue(torch.equal(m.sharded_tensor1, module_load.sharded_tensor1))
        self.assertTrue(
            torch.equal(
                m.submodule.sharded_tensor2, module_load.submodule.sharded_tensor2
            )
        )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_state_dict_new_group(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:2/cuda:0",
                "rank:3/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        pg = dist.new_group([2, 3])

        m = MyShardedModel1(spec, pg)

        # Test save
        m._register_state_dict_hook(state_dict_hook)
        buffer = io.BytesIO()
        torch.save(m.state_dict(), buffer)

        # Test load.
        module_load = MyShardedModel1(spec=None, group=pg)
        module_load._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True)

        buffer.seek(0)
        with load_with_process_group(pg):
            state_dict_deser = torch.load(buffer)
            module_load.load_state_dict(state_dict_deser, strict=False)

        # Verify after load.
        self.assertTrue(torch.equal(m.sharded_tensor1, module_load.sharded_tensor1))
        self.assertTrue(
            torch.equal(
                m.submodule.sharded_tensor2, module_load.submodule.sharded_tensor2
            )
        )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_state_dict_no_sharded_tensors(self):
        # Verify hooks don't affect modules with no ShardedTensors.
        m = torch.nn.Linear(10, 10)

        # Test save
        state_dict_before = m.state_dict()
        m._register_state_dict_hook(state_dict_hook)
        buffer = io.BytesIO()
        torch.save(m.state_dict(), buffer)
        self.assertEqual(state_dict_before, m.state_dict())

        # Test load.
        module_load = torch.nn.Linear(10, 10)
        module_load._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True)

        buffer.seek(0)
        state_dict_deser = torch.load(buffer)
        module_load.load_state_dict(state_dict_deser, strict=False)

        # Verify after load.
        self.assertEqual(m.weight, module_load.weight)
        self.assertEqual(m.bias, module_load.bias)

    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_load_state_dict_errors(self):
        self.init_rpc()

        dist.init_process_group(
            backend="nccl",
            world_size=self.world_size,
            rank=self.rank,
            init_method=f"file://{self.file_name}",
        )

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        m = MyShardedModel1(spec)

        # Test save
        m._register_state_dict_hook(state_dict_hook)
        buffer = io.BytesIO()
        torch.save(m.state_dict(), buffer)

        pg = dist.new_group(ranks=[0, 2, 3])

        buffer.seek(0)
        if self.rank != 0:
            with self.assertRaisesRegex(RuntimeError, "Local rank at save time was"):
                with load_with_process_group(pg):
                    state_dict_deser = torch.load(buffer)
        else:
            with self.assertRaisesRegex(
                RuntimeError, "Local world size at save time was"
            ):
                with load_with_process_group(pg):
                    state_dict_deser = torch.load(buffer)

        dist.destroy_process_group()
        buffer.seek(0)
        with self.assertRaisesRegex(
            RuntimeError, "Need to initialize default process group"
        ):
            state_dict_deser = torch.load(buffer)
        rpc.shutdown()

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_cleanup(self):
        def create_tensors():
            spec = ChunkShardingSpec(
                dim=0,
                placements=[
                    "rank:0/cuda:0",
                    "rank:1/cuda:1",
                    "rank:2/cuda:2",
                    "rank:3/cuda:3",
                ],
            )
            st1 = sharded_tensor.empty(spec, 10, 20, init_rrefs=True)
            st2 = sharded_tensor.empty(spec, 10, 20)

        create_tensors()
        self.assertEqual(0, len(sharded_tensor.api._sharded_tensor_map))


class TestShardedTensorEnumerable(ShardedTensorTestBase):
    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_sharded_tensor_metadata(self):
        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[0, 5],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:2/cuda:2",
                ),
                ShardMetadata(
                    shard_offsets=[5, 5],
                    shard_sizes=[5, 5],
                    placement="rank:3/cuda:3",
                ),
            ]
        )

        st = sharded_tensor.empty(spec, 10, 10, init_rrefs=True)
        st_metadata = st.metadata()
        self.assertEqual(torch.Size([10, 10]), st_metadata.size)
        self.assertEqual(torch.float, st.dtype)
        self.assertEqual(torch.strided, st.layout)
        self.assertEqual(False, st.requires_grad)
        self.assertTrue(st.is_contiguous())
        self.assertFalse(st.is_pinned())

        st = sharded_tensor.empty(spec, 10, 10, requires_grad=True, init_rrefs=True)
        self.assertEqual(True, st.requires_grad)

        st = sharded_tensor.empty(spec, 10, 10, dtype=torch.double, init_rrefs=True)
        self.assertEqual(torch.double, st.dtype)

        # Need CPU for pin_memory
        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cpu",
                ),
                ShardMetadata(
                    shard_offsets=[0, 5],
                    shard_sizes=[5, 5],
                    placement="rank:1/cpu",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:2/cpu",
                ),
                ShardMetadata(
                    shard_offsets=[5, 5],
                    shard_sizes=[5, 5],
                    placement="rank:3/cpu",
                ),
            ]
        )

        st = sharded_tensor.empty(spec, 10, 10, pin_memory=True, init_rrefs=True)
        self.assertTrue(st.is_pinned())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_grid_sharding(self):
        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[0, 5],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:2/cuda:2",
                ),
                ShardMetadata(
                    shard_offsets=[5, 5],
                    shard_sizes=[5, 5],
                    placement="rank:3/cuda:3",
                ),
            ]
        )

        st = sharded_tensor.empty(spec, 10, 10, init_rrefs=True)
        self.assertEqual((10, 10), st.size())
        self.assertEqual(1, len(st.local_shards()))

        # Verify local shard.
        local_shard = st.local_shards()[0]
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.tensor.device)
        self.assertEqual((5, 5), local_shard.tensor.size())

        # Verify local shard metadata.
        self.assertEqual(
            (self.rank // 2 * 5, (self.rank % 2) * 5),
            local_shard.metadata.shard_offsets,
        )
        self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
        self.assertEqual(
            f"rank:{self.rank}/cuda:{self.rank}", str(local_shard.metadata.placement)
        )

        # Verify global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(4, len(shards_metadata))
        for rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual(
                (rank // 2 * 5, (rank % 2) * 5), shard_metadata.shard_offsets
            )
            self.assertEqual((5, 5), shard_metadata.shard_sizes)
            self.assertEqual(f"rank:{rank}/cuda:{rank}", str(shard_metadata.placement))

        # Validate remote shards.
        remote_shards = st.remote_shards()
        self.assertEqual(3, len(remote_shards))

        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))
            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                shard = remote_shard.to_here()
                self.assertEqual((5, 5), shard.tensor.size())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_create_sharded_tensor_with_ones(self):
        """Test sharded_tensor.ones(...)"""

        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[0, 5],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:2/cuda:2",
                ),
                ShardMetadata(
                    shard_offsets=[5, 5],
                    shard_sizes=[5, 5],
                    placement="rank:3/cuda:3",
                ),
            ]
        )

        st = sharded_tensor.ones(spec, 10, 10, init_rrefs=True)
        self.assertEqual((10, 10), st.size())
        self.assertEqual(1, len(st.local_shards()))

        # Verify local shard is initialized with torch.ones
        local_shard = st.local_shards()[0]
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.tensor.device)
        self.assertEqual((5, 5), local_shard.tensor.size())
        self.assertEqual(local_shard.tensor, torch.ones(5, 5))

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_gather_even(self) -> None:
        """Test _sharded_tensor.gather(...) with evenly distributed._shards"""

        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[0, 5],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:2/cuda:2",
                ),
                ShardMetadata(
                    shard_offsets=[5, 5],
                    shard_sizes=[5, 5],
                    placement="rank:3/cuda:3",
                ),
            ]
        )

        h, w = 10, 10
        st = sharded_tensor.ones(spec, h, w, init_rrefs=True)

        full_tensor = None
        dst = 0
        if self.rank == dst:
            full_tensor = torch.zeros(h, w, device=torch.device(f"cuda:{dst}"))
        st.gather(dst, full_tensor)

        if self.rank == dst:
            self.assertEqual(full_tensor, torch.ones(h, w))
        else:
            self.assertIsNone(full_tensor)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_gather_uneven(self) -> None:
        """Test _sharded_tensor.gather(...) with unevenly distributed._shards"""

        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[0, 5],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[5, 5],
                    shard_sizes=[5, 5],
                    placement="rank:3/cuda:3",
                ),
            ]
        )

        h, w = 10, 10
        st = sharded_tensor.ones(spec, h, w, init_rrefs=True)

        full_tensor = None
        dst = 0
        if self.rank == dst:
            full_tensor = torch.zeros(h, w, device=torch.device(f"cuda:{dst}"))
        st.gather(dst, full_tensor)

        if self.rank == dst:
            self.assertEqual(full_tensor, torch.ones(h, w))
        else:
            self.assertIsNone(full_tensor)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_sharded_tensor_to_cpu(self):
        cpu_spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cpu",
                "rank:1/cpu",
                "rank:2/cpu",
                "rank:3/cpu",
            ],
        )
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 10, 20
        gloo_pg = dist.new_group(backend="gloo")

        # CPU sharded tensor should return the same instance (no copy)
        st_cpu = sharded_tensor.zeros(cpu_spec, h, w, process_group=gloo_pg)
        new_st_cpu = st_cpu.cpu()
        self.assertTrue(st_cpu is new_st_cpu)

        # GPU sharded tensor to cpu
        st = sharded_tensor.zeros(spec, h, w)
        # test ability to move st to CPU
        spec_before_move = st.sharding_spec()
        new_st = st.cpu(process_group=gloo_pg)
        # return a copy of original st
        self.assertFalse(st is new_st)
        # check the spec is still ChunkShardingSpec
        spec_after_move = new_st.sharding_spec()
        self.assertIsInstance(spec_after_move, ChunkShardingSpec)
        self.assertIsInstance(new_st._process_group, distributed_c10d.ProcessGroup)
        # test specs before and after the move almost the same except placement device
        self.assertEqual(spec_before_move.dim, spec_after_move.dim)
        self.assertEqual(
            len(spec_before_move.placements), len(spec_after_move.placements)
        )
        for i, remote_device_after in enumerate(spec_after_move.placements):
            remote_device_before = spec_before_move.placements[i]
            self.assertEqual(remote_device_before.rank(), remote_device_after.rank())
            self.assertEqual(str(remote_device_after.device()), "cpu")

        # ensure metdata also get changed to CPU
        metas = new_st.metadata().shards_metadata
        for meta in metas:
            self.assertEqual(str(meta.placement.device()), "cpu")

        # Test if a mixed sharded tensor (ShardedTensor with different devices) to cpu
        mixed_spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cpu",
                "rank:1/cpu",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        st = sharded_tensor.zeros(mixed_spec, h, w, process_group=gloo_pg)
        new_st = st.cpu()
        # return a copy of original st
        self.assertFalse(st is new_st)
        # check the spec is still ChunkShardingSpec
        spec_after_move = new_st.sharding_spec()
        self.assertIsInstance(spec_after_move, ChunkShardingSpec)
        # test specs before and after the move almost the same except placement device
        self.assertEqual(mixed_spec.dim, spec_after_move.dim)
        self.assertEqual(len(mixed_spec.placements), len(spec_after_move.placements))
        for i, remote_device_after in enumerate(spec_after_move.placements):
            remote_device_before = mixed_spec.placements[i]
            self.assertEqual(remote_device_before.rank(), remote_device_after.rank())
            self.assertEqual(str(remote_device_after.device()), "cpu")

        # ensure metdata also get changed to CPU
        metas = new_st.metadata().shards_metadata
        for meta in metas:
            self.assertEqual(str(meta.placement.device()), "cpu")

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_sharded_tensor_to_cuda(self):
        cpu_spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cpu",
                "rank:1/cpu",
                "rank:2/cpu",
                "rank:3/cpu",
            ],
        )
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 10, 20
        # CUDA sharded tensor should return a new ShardedTensor, but same
        # local shards(no movements)
        st_cuda = sharded_tensor.zeros(spec, h, w)
        new_st_cuda = st_cuda.cuda()
        self.assertTrue(st_cuda is not new_st_cuda)
        self.assertTrue(st_cuda.local_tensor() is new_st_cuda.local_tensor())

        gloo_pg = dist.new_group(backend="gloo")

        # CPU sharded tensor to GPU
        st_cpu = sharded_tensor.zeros(cpu_spec, h, w, process_group=gloo_pg)
        # test ability to move st to GPU
        spec_before_move = st_cpu.sharding_spec()
        new_st_gpu = st_cpu.cuda()
        # check the spec is still ChunkShardingSpec
        spec_after_move = new_st_gpu.sharding_spec()
        self.assertIsInstance(spec_after_move, ChunkShardingSpec)
        # test specs before and after the move almost the same except placement device
        self.assertEqual(spec_before_move.dim, spec_after_move.dim)
        self.assertEqual(
            len(spec_before_move.placements), len(spec_after_move.placements)
        )
        for i, remote_device_after in enumerate(spec_after_move.placements):
            remote_device_before = spec_before_move.placements[i]
            self.assertEqual(remote_device_before.rank(), remote_device_after.rank())
            self.assertEqual(str(remote_device_before.device().type), "cpu")
            self.assertEqual(str(remote_device_after.device().type), "cuda")

        # ensure metdata also get changed to GPU
        metas = new_st_gpu.metadata().shards_metadata
        for meta in metas:
            self.assertEqual(str(meta.placement.device().type), "cuda")

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_sharded_tensor_to_test(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 10, 20
        # CUDA sharded tensor should return a new ShardedTensor, but same
        # local shards(no movements)
        st = sharded_tensor.zeros(spec, h, w)
        # test same dtype, device return itself
        st_self = st.to(dtype=st.dtype, device="cuda")
        self.assertTrue(st_self is st)

        # test dtype to
        st_16 = st.to(torch.float16)
        self.assertFalse(st_16 is st)
        self.assertEqual(st_16.dtype, torch.float16)
        # test device to
        st_cpu = st.to(device=torch.device("cpu"))
        self.assertFalse(st_cpu is st)
        self.assertEqual(st_cpu.local_tensor().device.type, "cpu")
        st_cuda = st_cpu.to(device=torch.device("cuda"))
        self.assertEqual(st_cuda.local_tensor().device.type, "cuda")
        # non-kwarg device to
        st_cuda = st_cpu.to(torch.device("cuda"))
        self.assertEqual(st_cuda.local_tensor().device.type, "cuda")
        st_cpu = st_cuda.to(torch.device("cpu"))
        self.assertEqual(st_cpu.local_tensor().device.type, "cpu")
        # with string like device conversion
        st_cpu = st_cuda.to("cpu")
        self.assertEqual(st_cpu.local_tensor().device.type, "cpu")
        st_cuda = st_cpu.to("cuda")
        self.assertEqual(st_cuda.local_tensor().device.type, "cuda")
        # with int like device conversion
        st_cpu = st_cuda.to("cpu")
        self.assertEqual(st_cpu.local_tensor().device.type, "cpu")
        st_cuda = st_cpu.to(self.rank)
        self.assertEqual(st_cuda.local_tensor().device.type, "cuda")

        # test tensor to
        cuda_tensor = torch.randn(3, 4, dtype=torch.float16, device="cuda")
        st_cuda = st.to(cuda_tensor)
        self.assertFalse(st_cuda is st)
        self.assertEqual(st_cuda.dtype, torch.float16)

        cuda_tensor = torch.randn(3, 4, dtype=torch.float16, device="cuda:2")
        st_cuda = st.to(cuda_tensor)
        self.assertEqual(st_cuda.dtype, torch.float16)

        # test dtype and device together
        st_cpu_16 = st.to("cpu", torch.float16)
        self.assertEqual(st_cpu_16.dtype, torch.float16)
        self.assertEqual(st_cpu_16.local_tensor().device.type, "cpu")

        st_cuda_32 = st_cpu_16.to("cuda", torch.float32)
        self.assertEqual(st_cuda_32.dtype, torch.float32)
        self.assertEqual(st_cuda_32.local_tensor().device.type, "cuda")

        # test pass additional process group
        gloo_pg = dist.new_group(backend="gloo")
        st_gloo = st.to(device="cpu", process_group=gloo_pg)
        self.assertFalse(st_gloo is st)
        self.assertEqual(st_gloo.local_tensor().device.type, "cpu")
        self.assertEqual(st_gloo._process_group, gloo_pg)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_sharded_tensor_device(self):
        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        h, w = 10, 20
        # CUDA sharded tensor should return a new ShardedTensor, but same
        # local shards(no movements)
        st = sharded_tensor.zeros(spec, h, w)
        current_device = torch.device(torch.cuda.current_device())
        self.assertEqual(current_device, st.device)

        # test after to cpu, device get changed
        cpu_device = torch.device("cpu")
        st_cpu = st.to(device=cpu_device)
        self.assertEqual(st_cpu.device, cpu_device)

    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_uneven_shards(self):
        self.init_pg()

        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[2, 4],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[0, 4],
                    shard_sizes=[4, 2],
                    placement="rank:1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[2, 0],
                    shard_sizes=[4, 4],
                    placement="rank:2/cuda:2",
                ),
                ShardMetadata(
                    shard_offsets=[4, 4],
                    shard_sizes=[2, 2],
                    placement="rank:3/cuda:3",
                ),
            ]
        )

        st = sharded_tensor.empty(spec, 6, 6)
        self.assertEqual((6, 6), st.size())
        self.assertEqual(1, len(st.local_shards()))

        def verify_size(rank, tensor_dims):
            if rank == 0:
                self.assertEqual((2, 4), tensor_dims)
            elif rank == 1:
                self.assertEqual((4, 2), tensor_dims)
            elif rank == 2:
                self.assertEqual((4, 4), tensor_dims)
            elif rank == 3:
                self.assertEqual((2, 2), tensor_dims)

        def verify_offsets(rank, offsets):
            if rank == 0:
                self.assertEqual((0, 0), offsets)
            elif rank == 1:
                self.assertEqual((0, 4), offsets)
            elif rank == 2:
                self.assertEqual((2, 0), offsets)
            elif rank == 3:
                self.assertEqual((4, 4), offsets)

        # Verify local shard.
        local_shard = st.local_shards()[0]
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.tensor.device)
        verify_size(self.rank, local_shard.tensor.size())

        # Verify local shard metadata.
        verify_offsets(self.rank, local_shard.metadata.shard_offsets)
        verify_size(self.rank, local_shard.metadata.shard_sizes)
        self.assertEqual(
            f"rank:{self.rank}/cuda:{self.rank}", str(local_shard.metadata.placement)
        )

        # Verify global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(4, len(shards_metadata))
        for rank, shard_metadata in enumerate(shards_metadata):
            verify_offsets(rank, shard_metadata.shard_offsets)
            verify_size(rank, shard_metadata.shard_sizes)
            self.assertEqual(f"rank:{rank}/cuda:{rank}", str(shard_metadata.placement))

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_partial_world_size(self):
        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
            ]
        )

        st = sharded_tensor.empty(spec, 10, 5, init_rrefs=True)
        self.assertEqual((10, 5), st.size())
        if self.rank <= 1:
            self.assertEqual(1, len(st.local_shards()))
        else:
            self.assertEqual(0, len(st.local_shards()))

        if self.rank <= 1:
            # Verify local shard.
            local_shard = st.local_shards()[0]
            self.assertEqual(
                torch.device(f"cuda:{self.rank}"), local_shard.tensor.device
            )
            self.assertEqual((5, 5), local_shard.tensor.size())

            # Verify local shard metadata.
            self.assertEqual((self.rank * 5, 0), local_shard.metadata.shard_offsets)
            self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
            self.assertEqual(
                f"rank:{self.rank}/cuda:{self.rank}",
                str(local_shard.metadata.placement),
            )

        # Verify global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(2, len(shards_metadata))
        for rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual((rank * 5, 0), shard_metadata.shard_offsets)
            self.assertEqual((5, 5), shard_metadata.shard_sizes)
            self.assertEqual(f"rank:{rank}/cuda:{rank}", str(shard_metadata.placement))

        # Validate remote shards.
        remote_shards = st.remote_shards()
        if self.rank <= 1:
            self.assertEqual(1, len(remote_shards))
        else:
            self.assertEqual(2, len(remote_shards))

        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))

            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                shard = remote_shard.to_here()
                self.assertEqual((5, 5), shard.tensor.size())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_new_group(self):
        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:3/cuda:3",
                ),
            ]
        )

        pg = dist.new_group(ranks=[1, 2, 3])

        st = sharded_tensor.empty(spec, 10, 5, process_group=pg, init_rrefs=True)
        self.assertEqual((10, 5), st.size())
        if self.rank == 1 or self.rank == 3:
            # Verify local shard.
            local_shard = st.local_shards()[0]
            self.assertEqual(
                torch.device(f"cuda:{self.rank}"), local_shard.tensor.device
            )
            self.assertEqual((5, 5), local_shard.tensor.size())

            # Verify local shard metadata.
            self.assertEqual(
                (self.rank // 2 * 5, 0), local_shard.metadata.shard_offsets
            )
            self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
            self.assertEqual(
                f"rank:{self.rank}/cuda:{self.rank}",
                str(local_shard.metadata.placement),
            )

        # Verify global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(2, len(shards_metadata))
        for rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual((rank * 5, 0), shard_metadata.shard_offsets)
            self.assertEqual((5, 5), shard_metadata.shard_sizes)
            self.assertEqual(
                f"rank:{rank * 2 + 1}/cuda:{rank * 2 + 1}",
                str(shard_metadata.placement),
            )

        # Validate remote shards.
        remote_shards = st.remote_shards()
        if self.rank == 1 or self.rank == 3:
            self.assertEqual(1, len(remote_shards))
        else:
            self.assertEqual(2, len(remote_shards))

        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))

            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                shard = remote_shard.to_here()
                self.assertEqual((5, 5), shard.tensor.size())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_multiple_local_shards(self):
        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[0, 5],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[5, 5],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
            ]
        )

        st = sharded_tensor.empty(spec, 10, 10, init_rrefs=True)
        self.assertEqual((10, 10), st.size())

        if self.rank <= 1:
            self.assertEqual(2, len(st.local_shards()))

            # Verify local shards.
            for idx, local_shard in enumerate(st.local_shards()):
                self.assertEqual(
                    torch.device(f"cuda:{self.rank}"), local_shard.tensor.device
                )
                self.assertEqual((5, 5), local_shard.tensor.size())

                # Verify local shard metadata.
                self.assertEqual(
                    (idx * 5, self.rank * 5), local_shard.metadata.shard_offsets
                )
                self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
                self.assertEqual(
                    f"rank:{self.rank}/cuda:{self.rank}",
                    str(local_shard.metadata.placement),
                )
        else:
            self.assertEqual(0, len(st.local_shards()))

        # Verify global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(4, len(shards_metadata))
        for shard_rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual(
                (shard_rank // 2 * 5, (shard_rank % 2) * 5),
                shard_metadata.shard_offsets,
            )
            self.assertEqual((5, 5), shard_metadata.shard_sizes)
            self.assertEqual(
                f"rank:{shard_rank % 2}/cuda:{shard_rank % 2}",
                str(shard_metadata.placement),
            )

        # Validate remote shards.
        remote_shards = st.remote_shards()
        if self.rank <= 1:
            self.assertEqual(1, len(remote_shards))
        else:
            self.assertEqual(2, len(remote_shards))

        owners = {}
        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(2, len(shards))
            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                shard = remote_shard.to_here()
                self.assertEqual((5, 5), shard.tensor.size())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_with_rpc_names(self):
        spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="worker0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[0, 5],
                    shard_sizes=[5, 5],
                    placement="worker1/cuda:1",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="worker2/cuda:2",
                ),
                ShardMetadata(
                    shard_offsets=[5, 5],
                    shard_sizes=[5, 5],
                    placement="worker3/cuda:3",
                ),
            ]
        )

        st = sharded_tensor.empty(spec, 10, 10, init_rrefs=True)
        self.assertEqual((10, 10), st.size())
        self.assertEqual(1, len(st.local_shards()))

        # Verify local shard.
        local_shard = st.local_shards()[0]
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.tensor.device)
        self.assertEqual((5, 5), local_shard.tensor.size())

        # Verify local shard metadata.
        self.assertEqual(
            (self.rank // 2 * 5, (self.rank % 2) * 5),
            local_shard.metadata.shard_offsets,
        )
        self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
        self.assertEqual(
            f"worker{self.rank}/cuda:{self.rank}", str(local_shard.metadata.placement)
        )

        # Verify global metadata.
        st_metadata = st.metadata()
        shards_metadata = st_metadata.shards_metadata
        self.assertEqual(4, len(shards_metadata))
        for rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual(
                (rank // 2 * 5, (rank % 2) * 5), shard_metadata.shard_offsets
            )
            self.assertEqual((5, 5), shard_metadata.shard_sizes)
            self.assertEqual(f"worker{rank}/cuda:{rank}", str(shard_metadata.placement))

        # Validate remote shards.
        remote_shards = st.remote_shards()
        self.assertEqual(3, len(remote_shards))

        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))
            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                shard = remote_shard.to_here()
                self.assertEqual((5, 5), shard.tensor.size())


class TestShardedTensorFromLocalTensor(ShardedTensorTestBase):
    def _generate_st_from_chunk_local_tensor(self, st_size, sharding_spec):
        tensor_meta = sharding_spec.build_metadata(st_size, TensorProperties())
        pg = dist.distributed_c10d._get_default_group()

        local_tensor = None
        local_shard_metadata = None
        rank_to_metadata = {}
        for shard_metadata in tensor_meta.shards_metadata:
            rank, device = _parse_and_validate_remote_device(
                pg, shard_metadata.placement
            )
            rank_to_metadata[rank] = shard_metadata
            if rank == self.rank:
                local_tensor = torch.rand(shard_metadata.shard_sizes).cuda(device)
                local_shard_metadata = shard_metadata

        # TODO: figure out what the API should behave when some rank have no shard
        # see https://github.com/pytorch/pytorch/issues/73133
        assert local_tensor is not None
        st = ShardedTensor._init_from_local_tensor(
            local_tensor,
            sharding_spec,
            st_size,
            init_rrefs=True,
        )
        self.assertEqual(tuple(st_size), st.size())
        self.assertEqual(1, len(st.local_shards()))

        # Verify local shard.
        local_shard = st.local_shards()[0]
        self.assertEqual(st.local_tensor(), local_tensor)
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.tensor.device)

        # Verify local shard metadata.
        self.assertEqual(
            local_shard_metadata.shard_offsets, local_shard.metadata.shard_offsets
        )
        self.assertEqual(
            local_shard_metadata.shard_sizes, local_shard.metadata.shard_sizes
        )
        self.assertEqual(local_shard_metadata.placement, local_shard.metadata.placement)

        # Verify global metadata.
        st_shards_metadata = st.metadata().shards_metadata
        self.assertEqual(self.world_size, len(st_shards_metadata))
        self.assertEqual(tensor_meta.shards_metadata, st_shards_metadata)

        # Validate remote shards.
        remote_shards = st.remote_shards()
        self.assertEqual(self.world_size - 1, len(remote_shards))
        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))
            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                # If remote shard does not exist, to_here() will throw exception.
                if tensor_meta.shards_metadata[rpc_rank]:
                    shard = remote_shard.to_here()
                    self.assertEqual(
                        rank_to_metadata[rpc_rank].shard_sizes, shard.tensor.size()
                    )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_tensor(self):
        chunk_specs = _chunk_sharding_specs_list_for_test([0, 1, 1, 0], seed=31)
        for spec in chunk_specs:
            self._generate_st_from_chunk_local_tensor([20, 10], spec)
            self._generate_st_from_chunk_local_tensor([21, 11], spec)
            self._generate_st_from_chunk_local_tensor([23, 16], spec)
            self._generate_st_from_chunk_local_tensor([44, 16, 8], spec)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_tensor_errors(self):
        enumerable_sharding_spec = EnumerableShardingSpec(
            [
                ShardMetadata(
                    shard_offsets=[0, 0],
                    shard_sizes=[5, 5],
                    placement="rank:0/cuda:0",
                ),
                ShardMetadata(
                    shard_offsets=[5, 0],
                    shard_sizes=[5, 5],
                    placement="rank:1/cuda:1",
                ),
            ]
        )
        st_size = [24, 12]
        local_tensor = torch.rand(*st_size).cuda(self.rank)
        with self.assertRaisesRegex(ValueError, "do not cover the entire tensor"):
            ShardedTensor._init_from_local_tensor(
                local_tensor,
                enumerable_sharding_spec,
                st_size,
            )
        chunk_specs = _chunk_sharding_specs_list_for_test([0], seed=31)
        with self.assertRaisesRegex(
            ValueError, "local_tensor is not a contiguous Tensor."
        ):
            ShardedTensor._init_from_local_tensor(
                local_tensor.t(),
                chunk_specs[0],
                st_size,
            )


class TestShardedTensorFromLocalShards(ShardedTensorTestBase):
    @with_comms(init_rpc=False)
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_local_shards(self):
        shard_offsets = [(self.rank // 2) * 5, (self.rank % 2) * 5]
        local_shard_metadata = ShardMetadata(
            shard_offsets=shard_offsets,
            shard_sizes=[5, 5],
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )

        local_tensor = torch.randn(5, 5, device=f"cuda:{self.rank}")
        local_shard = sharded_tensor.Shard(local_tensor, local_shard_metadata)
        local_shard_from_offsets = sharded_tensor.Shard.from_tensor_and_offsets(
            local_tensor, shard_offsets=shard_offsets, rank=self.rank
        )
        self.assertEqual(local_shard.metadata, local_shard_from_offsets.metadata)

        wrong_local_shard_metadata = ShardMetadata(
            shard_offsets=shard_offsets,
            shard_sizes=[6, 5],
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )
        with self.assertRaisesRegex(ValueError, "Shard tensor size does not match"):
            local_shard_from_wrong_meta = sharded_tensor.Shard(
                local_tensor,
                metadata=wrong_local_shard_metadata,
            )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_shards(self):
        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=[5, 5],
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )

        local_shards = [
            sharded_tensor.Shard(
                torch.randn(5, 5, device=f"cuda:{self.rank}"), local_shard_metadata
            )
        ]

        st = sharded_tensor.init_from_local_shards(
            local_shards, [10, 10], init_rrefs=True
        )
        self.assertEqual((10, 10), st.size())
        self.assertEqual(1, len(st.local_shards()))

        # Verify local shard.
        local_shard = st.local_shards()[0]
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.tensor.device)
        self.assertEqual((5, 5), local_shard.tensor.size())

        # Verify local shard metadata.
        self.assertEqual(
            (self.rank // 2 * 5, (self.rank % 2) * 5),
            local_shard.metadata.shard_offsets,
        )
        self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
        self.assertEqual(
            f"rank:{self.rank}/cuda:{self.rank}", str(local_shard.metadata.placement)
        )

        # Verify global metadata.
        shards_metadata = st.metadata().shards_metadata
        self.assertEqual(4, len(shards_metadata))
        for rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual(
                (rank // 2 * 5, (rank % 2) * 5), shard_metadata.shard_offsets
            )
            self.assertEqual((5, 5), shard_metadata.shard_sizes)
            self.assertEqual(f"rank:{rank}/cuda:{rank}", str(shard_metadata.placement))

        # Validate remote shards.
        remote_shards = st.remote_shards()
        self.assertEqual(3, len(remote_shards))

        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))
            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                shard = remote_shard.to_here()
                self.assertEqual((5, 5), shard.tensor.size())

    @skip_if_lt_x_gpu(4)
    def test_st_base_init_from_local_shards_and_global_metadata(self):
        world_size = 4
        shards_metadata = []
        shards = []
        for rank in range(world_size):
            local_shard_metadata = ShardMetadata(
                shard_offsets=[(rank // 2) * 5, (rank % 2) * 5],
                shard_sizes=[5, 5],
                placement=f"rank:{rank}/cuda:{rank}",
            )
            shards_metadata.append(local_shard_metadata)
            shards.append(
                sharded_tensor.Shard(
                    torch.randn(5, 5, device=f"cuda:{rank}"), local_shard_metadata
                )
            )

        tensor_properties = TensorProperties(
            dtype=torch.get_default_dtype(),
            layout=torch.strided,
            requires_grad=False,
            memory_format=torch.contiguous_format,
            pin_memory=False,
        )

        sharded_tensor_metadata = sharded_tensor.ShardedTensorMetadata(
            shards_metadata=shards_metadata,
            size=torch.Size([10, 10]),
            tensor_properties=tensor_properties,
        )

        st_base = sharded_tensor.ShardedTensorBase._init_from_local_shards_and_global_metadata(
            shards, sharded_tensor_metadata=sharded_tensor_metadata
        )
        self.assertEqual(4, len(st_base.local_shards()))

        # Verify local shard of st_base
        local_shard = st_base.local_shards()[0]
        self.assertEqual(torch.device("cuda:0"), local_shard.tensor.device)
        self.assertEqual((5, 5), local_shard.tensor.size())

        # Verify local shard metadata.
        self.assertEqual(
            (0, 0),
            local_shard.metadata.shard_offsets,
        )
        self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
        self.assertEqual("rank:0/cuda:0", str(local_shard.metadata.placement))

        # Verify global metadata.
        shards_metadata = st_base.metadata().shards_metadata
        self.assertEqual(4, len(shards_metadata))
        for rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual(
                (rank // 2 * 5, (rank % 2) * 5), shard_metadata.shard_offsets
            )
            self.assertEqual((5, 5), shard_metadata.shard_sizes)
            self.assertEqual(f"rank:{rank}/cuda:{rank}", str(shard_metadata.placement))

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_shards_and_global_metadata(self):
        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=[5, 5],
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )

        shards_metadata = []
        for r in range(self.world_size):
            if r == self.rank:
                shards_metadata.append(local_shard_metadata)
            else:
                shards_metadata.append(
                    ShardMetadata(
                        shard_offsets=[(r // 2) * 5, (r % 2) * 5],
                        shard_sizes=[5, 5],
                        placement=f"rank:{r}/cuda:{r}",
                    )
                )

        local_shards = [
            sharded_tensor.Shard(
                torch.randn(5, 5, device=f"cuda:{self.rank}"), local_shard_metadata
            )
        ]

        tensor_properties = TensorProperties(
            dtype=torch.get_default_dtype(),
            layout=torch.strided,
            requires_grad=False,
            memory_format=torch.contiguous_format,
            pin_memory=False,
        )

        sharded_tensor_metadata = sharded_tensor.ShardedTensorMetadata(
            shards_metadata=shards_metadata,
            size=torch.Size([10, 10]),
            tensor_properties=tensor_properties,
        )

        st = ShardedTensor._init_from_local_shards_and_global_metadata(
            local_shards,
            sharded_tensor_metadata,
            init_rrefs=True,
        )
        self.assertEqual((10, 10), st.size())
        self.assertEqual(1, len(st.local_shards()))

        # Verify local shard.
        local_shard = st.local_shards()[0]
        self.assertEqual(torch.device(f"cuda:{self.rank}"), local_shard.tensor.device)
        self.assertEqual((5, 5), local_shard.tensor.size())

        # Verify local shard metadata.
        self.assertEqual(
            (self.rank // 2 * 5, (self.rank % 2) * 5),
            local_shard.metadata.shard_offsets,
        )
        self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
        self.assertEqual(
            f"rank:{self.rank}/cuda:{self.rank}", str(local_shard.metadata.placement)
        )

        # Verify global metadata.
        shards_metadata = st.metadata().shards_metadata
        self.assertEqual(4, len(shards_metadata))
        for rank, shard_metadata in enumerate(shards_metadata):
            self.assertEqual(
                (rank // 2 * 5, (rank % 2) * 5), shard_metadata.shard_offsets
            )
            self.assertEqual((5, 5), shard_metadata.shard_sizes)
            self.assertEqual(f"rank:{rank}/cuda:{rank}", str(shard_metadata.placement))

        # Validate remote shards.
        remote_shards = st.remote_shards()
        self.assertEqual(3, len(remote_shards))

        for rpc_rank, shards in remote_shards.items():
            self.assertEqual(1, len(shards))
            for remote_shard in shards:
                self.assertEqual(rpc_rank, remote_shard.owner().id)
                shard = remote_shard.to_here()
                self.assertEqual((5, 5), shard.tensor.size())

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_shards_new_group(self):
        new_pg = dist.new_group(ranks=[1, 2, 3])

        if self.rank != 0:
            local_shard_metadata = ShardMetadata(
                shard_offsets=[5 * (self.rank - 1), 0],
                shard_sizes=[5, 5],
                placement=f"rank:{self.rank}/cuda:{self.rank}",
            )
            local_shards = [
                sharded_tensor.Shard(
                    torch.randn(5, 5, device=f"cuda:{self.rank}"), local_shard_metadata
                )
            ]

            st = sharded_tensor.init_from_local_shards(
                local_shards, [15, 5], process_group=new_pg
            )

            # Verify local shard.
            local_shard = st.local_shards()[0]
            self.assertEqual(
                torch.device(f"cuda:{self.rank}"), local_shard.tensor.device
            )
            self.assertEqual((5, 5), local_shard.tensor.size())

            # Verify local shard metadata.
            self.assertEqual(
                ((self.rank - 1) * 5, 0), local_shard.metadata.shard_offsets
            )
            self.assertEqual((5, 5), local_shard.metadata.shard_sizes)
            self.assertEqual(
                f"rank:{self.rank}/cuda:{self.rank}",
                str(local_shard.metadata.placement),
            )

            # Verify global metadata.
            st_metadata = st.metadata()
            shards_metadata = st_metadata.shards_metadata
            self.assertEqual(3, len(shards_metadata))
            for rank, shard_metadata in enumerate(shards_metadata):
                self.assertEqual((rank * 5, 0), shard_metadata.shard_offsets)
                self.assertEqual((5, 5), shard_metadata.shard_sizes)
                self.assertEqual(
                    f"rank:{rank + 1}/cuda:{rank + 1}", str(shard_metadata.placement)
                )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_shards_invalid_local_shards(self):
        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=[5, 5],
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )

        indices = [[0, 1, 1], [2, 0, 2]]
        values = [3.2, 4.5, 5.8]
        sparse_tensor = torch.sparse_coo_tensor(
            indices, values, (5, 5), device=f"cuda:{self.rank}"
        )

        empty_local_shards = []
        with self.assertRaisesRegex(ValueError, "have no local shards on all ranks"):
            st = sharded_tensor.init_from_local_shards(
                empty_local_shards, [10, 10], init_rrefs=True
            )

        wrong_layout_shards = [
            sharded_tensor.Shard(sparse_tensor, local_shard_metadata)
        ]
        with self.assertRaisesRegex(
            ValueError, "Only torch.strided layout is currently supported"
        ):
            st = sharded_tensor.init_from_local_shards(
                wrong_layout_shards, [10, 10], init_rrefs=True
            )

        wrong_memory_format_shards = [
            sharded_tensor.Shard(
                torch.randn(5, 5, device=f"cuda:{self.rank}").t(), local_shard_metadata
            )
        ]
        with self.assertRaisesRegex(
            ValueError,
            "Only torch.contiguous_format memory_format is currently supported",
        ):
            st = sharded_tensor.init_from_local_shards(
                wrong_memory_format_shards, [10, 10], init_rrefs=True
            )

        with self.assertRaisesRegex(ValueError, "Shard tensor size does not match"):
            wrong_size_shards = [
                sharded_tensor.Shard(
                    torch.randn(2, 3, device=f"cuda:{self.rank}"), local_shard_metadata
                )
            ]

        with self.assertRaisesRegex(
            ValueError, "Local shard tensor device does not match"
        ):
            wrong_device_shards = [
                sharded_tensor.Shard(torch.randn(5, 5), local_shard_metadata)
            ]

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_shards_invalid_property_cross_ranks(self):
        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=[5, 5],
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )
        tensor_overall_size = [10, 10] if self.rank == 0 else [10, 5]
        wrong_dtype_shards = [
            sharded_tensor.Shard(
                torch.ones(5, 5, device=f"cuda:{self.rank}"), local_shard_metadata
            )
        ]
        with self.assertRaisesRegex(
            ValueError,
            "ShardedTensor global_size property does not match from different ranks!",
        ):
            st = sharded_tensor.init_from_local_shards(
                wrong_dtype_shards, tensor_overall_size, init_rrefs=True
            )

        tensor_dtype = torch.int if self.rank == 0 else torch.float32
        wrong_dtype_shards = [
            sharded_tensor.Shard(
                torch.ones(5, 5, device=f"cuda:{self.rank}", dtype=tensor_dtype),
                local_shard_metadata,
            )
        ]
        with self.assertRaisesRegex(
            ValueError,
            "ShardedTensor dtype property does not match from different ranks!",
        ):
            st = sharded_tensor.init_from_local_shards(
                wrong_dtype_shards, [10, 10], init_rrefs=True
            )

        tensor_requires_grad = True if self.rank == 0 else False
        wrong_requires_grad_shards = [
            sharded_tensor.Shard(
                torch.randn(
                    5, 5, device=f"cuda:{self.rank}", requires_grad=tensor_requires_grad
                ),
                local_shard_metadata,
            )
        ]
        with self.assertRaisesRegex(
            ValueError,
            "ShardedTensor requires_grad property does not match from different ranks!",
        ):
            st = sharded_tensor.init_from_local_shards(
                wrong_requires_grad_shards, [10, 10], init_rrefs=True
            )

        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=[5, 5],
            placement=f"rank:{self.rank}/cpu",
        )

    @with_comms(init_rpc=False, backend="gloo")
    @skip_if_lt_x_gpu(4)
    def test_init_from_local_shards_invalid_pin_memory(self):
        # pin memory can only be on dense cpu
        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=[5, 5],
            placement=f"rank:{self.rank}/cpu",
        )
        wrong_pin_memory_local_shards = [
            sharded_tensor.Shard(
                torch.randn(5, 5, pin_memory=True), local_shard_metadata
            ),
            sharded_tensor.Shard(
                torch.randn(5, 5, pin_memory=False), local_shard_metadata
            ),
        ]
        with self.assertRaisesRegex(
            ValueError, "Local shards' tensor pin_memory property need to be the same"
        ):
            st = sharded_tensor.init_from_local_shards(
                wrong_pin_memory_local_shards, [10, 10], init_rrefs=True
            )

        tensor_pin_memory = True if self.rank == 0 else False
        wrong_pin_memory_shards_cross_ranks = [
            sharded_tensor.Shard(
                torch.randn(5, 5, pin_memory=tensor_pin_memory), local_shard_metadata
            )
        ]
        with self.assertRaisesRegex(
            ValueError,
            "ShardedTensor pin_memory property does not match from different ranks!",
        ):
            st = sharded_tensor.init_from_local_shards(
                wrong_pin_memory_shards_cross_ranks, [10, 10], init_rrefs=True
            )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_shards_invalid_shards_overlap(self):
        local_shard_size = [5, 5] if self.rank != 0 else [6, 6]
        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=local_shard_size,
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )

        local_shards = [
            sharded_tensor.Shard(
                torch.randn(local_shard_size, device=f"cuda:{self.rank}"),
                local_shard_metadata,
            )
        ]

        with self.assertRaisesRegex(ValueError, "overlap"):
            sharded_tensor.init_from_local_shards(
                local_shards, [10, 10], init_rrefs=True
            )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_shards_invalid_shards_gaps(self):
        local_shard_size = [5, 5] if self.rank != 0 else [4, 4]
        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=local_shard_size,
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )

        local_shards = [
            sharded_tensor.Shard(
                torch.randn(local_shard_size, device=f"cuda:{self.rank}"),
                local_shard_metadata,
            )
        ]

        with self.assertRaisesRegex(ValueError, "does not match tensor volume"):
            sharded_tensor.init_from_local_shards(
                local_shards, [10, 10], init_rrefs=True
            )

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_init_from_local_shards_and_global_metadata_invalid_shards(self):
        local_shard_metadata = ShardMetadata(
            shard_offsets=[(self.rank // 2) * 5, (self.rank % 2) * 5],
            shard_sizes=[5, 5],
            placement=f"rank:{self.rank}/cuda:{self.rank}",
        )

        shards_metadata = []
        for r in range(self.world_size):
            if r == self.rank:
                shards_metadata.append(local_shard_metadata)
            else:
                shards_metadata.append(
                    ShardMetadata(
                        shard_offsets=[(r // 2) * 5, (r % 2) * 5],
                        shard_sizes=[5, 5],
                        placement=f"rank:{r}/cuda:{r}",
                    )
                )

        tensor_properties = TensorProperties(
            dtype=torch.get_default_dtype(),
            layout=torch.strided,
            requires_grad=False,
            memory_format=torch.contiguous_format,
            pin_memory=False,
        )

        sharded_tensor_metadata = sharded_tensor.ShardedTensorMetadata(
            shards_metadata=shards_metadata,
            size=torch.Size([10, 10]),
            tensor_properties=tensor_properties,
        )

        empty_local_shards = []
        with self.assertRaisesRegex(
            RuntimeError, "does not match number of local shards metadata"
        ):
            ShardedTensor._init_from_local_shards_and_global_metadata(
                empty_local_shards, sharded_tensor_metadata
            )

        wrong_num_shards = [
            sharded_tensor.Shard(
                torch.randn(5, 5, device=f"cuda:{self.rank}"), local_shard_metadata
            ),
            sharded_tensor.Shard(
                torch.randn(5, 5, device=f"cuda:{self.rank}"), local_shard_metadata
            ),
        ]
        with self.assertRaisesRegex(
            RuntimeError, "does not match number of local shards metadata"
        ):
            ShardedTensor._init_from_local_shards_and_global_metadata(
                wrong_num_shards, sharded_tensor_metadata
            )

        with self.assertRaisesRegex(
            ValueError, "Shard tensor size does not match with metadata.shard_lengths"
        ):
            wrong_size_shards = [
                sharded_tensor.Shard(
                    torch.randn(2, 3, device=f"cuda:{self.rank}"), local_shard_metadata
                )
            ]

        with self.assertRaisesRegex(
            ValueError,
            "Local shard tensor device does not match with local Shard's placement",
        ):
            wrong_device_shards = [
                sharded_tensor.Shard(torch.randn(5, 5), local_shard_metadata)
            ]

        wrong_dtype_shards = [
            sharded_tensor.Shard(
                torch.ones(5, 5, device=f"cuda:{self.rank}", dtype=torch.int),
                local_shard_metadata,
            )
        ]
        with self.assertRaisesRegex(
            ValueError, "Local shards' tensor dtype property is incompatible with"
        ):
            ShardedTensor._init_from_local_shards_and_global_metadata(
                wrong_dtype_shards, sharded_tensor_metadata
            )

        indices = [[0, 1, 1], [2, 0, 2]]
        values = [3.2, 4.5, 5.8]
        sparse_tensor = torch.sparse_coo_tensor(
            indices, values, (5, 5), device=f"cuda:{self.rank}"
        )

        wrong_layout_shards = [
            sharded_tensor.Shard(sparse_tensor, local_shard_metadata)
        ]
        with self.assertRaisesRegex(
            ValueError, "Local shards' tensor layout property is incompatible with"
        ):
            ShardedTensor._init_from_local_shards_and_global_metadata(
                wrong_layout_shards, sharded_tensor_metadata
            )

        wrong_requires_grad_shards = [
            sharded_tensor.Shard(
                torch.randn(5, 5, device=f"cuda:{self.rank}", requires_grad=True),
                local_shard_metadata,
            )
        ]
        with self.assertRaisesRegex(
            ValueError,
            "Local shards' tensor requires_grad property is incompatible with",
        ):
            ShardedTensor._init_from_local_shards_and_global_metadata(
                wrong_requires_grad_shards, sharded_tensor_metadata
            )

        wrong_memory_format_shards = [
            sharded_tensor.Shard(
                torch.randn(5, 5, device=f"cuda:{self.rank}").t(), local_shard_metadata
            )
        ]
        with self.assertRaisesRegex(
            ValueError,
            "Only torch.contiguous_format memory_format is currently supported",
        ):
            ShardedTensor._init_from_local_shards_and_global_metadata(
                wrong_memory_format_shards, sharded_tensor_metadata
            )
        # pin_memory can only be on CPU
        local_shard_metadata.placement = _remote_device(f"rank:{self.rank}/cpu")
        wrong_pin_memory_shards = [
            sharded_tensor.Shard(
                torch.randn(5, 5, pin_memory=True), local_shard_metadata
            )
        ]
        with self.assertRaisesRegex(
            ValueError, "Local shards' tensor pin_memory property is incompatible with"
        ):
            ShardedTensor._init_from_local_shards_and_global_metadata(
                wrong_pin_memory_shards, sharded_tensor_metadata
            )


class TestShardedTensorCustomOps(ShardedTensorTestBase):
    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_custom_op(self):
        @custom_sharded_op_impl(torch.asin)
        def my_sharded_asin(types, args, kwargs, process_group):
            return torch.asin(args[0].local_shards()[0].tensor)

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )

        st = sharded_tensor.rand(spec, 10, 10)
        res = torch.asin(st)
        self.assertEqual(res, torch.asin(st.local_shards()[0].tensor))

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_custom_op_override(self):
        t = torch.rand(10, 10).cuda(self.rank)

        from torch.distributed._shard.sharding_spec.api import custom_sharding_spec_op

        @custom_sharding_spec_op(ChunkShardingSpec, torch.nn.functional.linear)
        def my_sharded_linear(types, args, kwargs, process_group):
            return t

        spec = ChunkShardingSpec(
            dim=0,
            placements=[
                "rank:0/cuda:0",
                "rank:1/cuda:1",
                "rank:2/cuda:2",
                "rank:3/cuda:3",
            ],
        )
        m = torch.nn.Linear(32, 16).cuda(self.rank)
        shard_parameter(m, "weight", spec)

        result = m(torch.rand(15, 32).cuda(self.rank))
        self.assertEqual(t, result)

    @with_comms
    @skip_if_lt_x_gpu(4)
    @requires_nccl()
    def test_custom_op_errors(self):
        with self.assertRaisesRegex(TypeError, "expects signature"):

            @custom_sharded_op_impl(torch.nn.functional.linear)
            def my_op1(types, args, kwargs, process_group, random_param):
                pass

        with self.assertRaisesRegex(TypeError, "expects signature"):

            @custom_sharded_op_impl(torch.nn.functional.linear)
            def my_op2(types):
                pass


class TestShardMetadata(ShardedTensorTestBase):
    @with_comms
    @requires_nccl()
    def test_shard_metadata_init(self):
        pg = dist.distributed_c10d._get_default_group()

        md = ShardMetadata([10], [0])
        self.assertIsNone(md.placement)
        with self.assertRaisesRegex(ValueError, "remote device is None"):
            _parse_and_validate_remote_device(pg, md.placement)

        # String placement gets converted by ctor
        md = ShardMetadata([10], [0], "rank:0/cpu")
        self.assertEqual(md.placement, _remote_device("rank:0/cpu"))
        rank, device = _parse_and_validate_remote_device(pg, md.placement)
        self.assertEqual(0, rank)
        self.assertEqual(device, torch.device("cpu"))

    @with_comms
    @requires_nccl()
    def test_create_shard_with_no_placement(self):
        md = ShardMetadata([0], [10])
        shard = Shard(torch.zeros(10), md)
        self.assertIsNone(shard.metadata.placement)


class TestShardedTensorSubGroupInit(TestCase):
    @spawn_threads_and_init_comms(world_size=4)
    def test_sub_process_group_sharded_tensor_init(self):
        world_pg = dist.GroupMember.WORLD
        rank = dist.get_rank()

        sub_group_sz = 2
        sub_pg_ranks = [r for r in range(4) if r % sub_group_sz == rank % sub_group_sz]
        sub_pg = dist.new_group(
            sub_pg_ranks,
            backend=dist.get_backend(world_pg),
            use_local_synchronization=True,
        )
        dist.barrier(sub_pg)

        ShardedTensor._init_from_local_shards(
            [
                Shard(
                    tensor=torch.tensor([1, 2, 3], device="meta"),
                    metadata=ShardMetadata(
                        shard_offsets=[3 * (rank // sub_group_sz)],
                        shard_sizes=[3],
                        placement=f"rank:{rank}/meta",
                    ),
                )
            ],
            6,
            process_group=sub_pg,
        )

    @spawn_threads_and_init_comms(world_size=4)
    def test_sub_process_group_placement_validation(self):
        world_pg = dist.GroupMember.WORLD
        self.assertIsNotNone(world_pg)
        rank = dist.get_rank()

        sub_group_sz = 2
        sub_pg_ranks = [r for r in range(4) if r % sub_group_sz == rank % sub_group_sz]
        sub_pg = dist.new_group(
            sub_pg_ranks,
            backend=dist.get_backend(world_pg),
            use_local_synchronization=True,
        )
        dist.barrier(sub_pg)

        for r in sub_pg_ranks:
            _parse_and_validate_remote_device(
                sub_pg, _remote_device(f"rank:{r}/cuda:{r % sub_group_sz}")
            )


class TestCreateTensorNoProcessGroupMode(TestCase):
    def test_init_from_local_shards_and_global_metadata(self):
        st_metadata: ShardedTensorMetadata = ShardedTensorMetadata(
            shards_metadata=[
                ShardMetadata(
                    shard_offsets=[0, 0], shard_sizes=[2, 2], placement="rank:0/cpu"
                ),
                ShardMetadata(
                    shard_offsets=[2, 0], shard_sizes=[2, 2], placement="rank:1/cpu"
                ),
            ],
            size=torch.Size([4, 2]),
        )
        st_local_shards: List[Shard] = []
        for shard_metadata in st_metadata.shards_metadata:
            st_local_shards.append(
                Shard(
                    tensor=torch.zeros(
                        shard_metadata.shard_sizes,
                        device=shard_metadata.placement.device(),
                    ),
                    metadata=shard_metadata,
                )
            )

        ShardedTensorBase._init_from_local_shards_and_global_metadata(
            local_shards=st_local_shards,
            sharded_tensor_metadata=st_metadata,
        )

    def test_non_contiguous_local_shards(self):
        st_metadata: ShardedTensorMetadata = ShardedTensorMetadata(
            shards_metadata=[
                ShardMetadata(
                    shard_offsets=[0, 0], shard_sizes=[2, 2], placement="rank:0/cpu"
                ),
                ShardMetadata(
                    shard_offsets=[2, 0], shard_sizes=[2, 2], placement="rank:1/cpu"
                ),
            ],
            size=torch.Size([4, 2]),
        )
        st_local_shards: List[Shard] = []
        src = torch.randn(4, 2)
        for shard_metadata in st_metadata.shards_metadata:
            offsets = shard_metadata.shard_offsets
            sizes = shard_metadata.shard_sizes
            st_local_shards.append(
                Shard(
                    tensor=src[
                        offsets[0] : offsets[0] + sizes[0],
                        offsets[1] : offsets[1] + sizes[1],
                    ],
                    metadata=shard_metadata,
                )
            )

        ShardedTensorBase._init_from_local_shards_and_global_metadata(
            local_shards=st_local_shards,
            sharded_tensor_metadata=st_metadata,
        )


if __name__ == "__main__":
    run_tests()