test/distributed/fsdp/test_shard_utils.py

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

import torch
from torch.distributed._shard.sharded_tensor import (
    init_from_local_shards,
    Shard,
    ShardMetadata,
)
from torch.distributed._shard.sharding_spec import (
    ChunkShardingSpec,
    EnumerableShardingSpec,
)
from torch.distributed.distributed_c10d import _get_default_group
from torch.distributed.fsdp._shard_utils import (
    _create_chunk_sharded_tensor,
    _offsets_to_split_sizes,
    _reshard_flatten_tensor,
)
from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
from torch.testing._internal.common_fsdp import FSDPTest
from torch.testing._internal.common_utils import TestCase


class TestShardUtils(TestCase):
    def test_offsets_to_split_sizes(self):
        tensor_numel = 40

        def _get_and_check_split_sizes(
            world_size,
            in_offsets,
            out_offsets,
            in_split_sizes,
        ):

            for my_rank in range(world_size):
                _in_split_sizes = in_split_sizes[my_rank]
                _out_split_sizes = [
                    in_split_sizes[i][my_rank] for i in range(world_size)
                ]
                res_in_split_sizes, res_out_split_sizes = _offsets_to_split_sizes(
                    in_offsets, out_offsets, tensor_numel, world_size, my_rank
                )
                self.assertEqual(_in_split_sizes, res_in_split_sizes)
                self.assertEqual(_out_split_sizes, res_out_split_sizes)

        # The tensor size can be evenly divided by the world size.
        world_size = 4
        in_offsets = [0, 10, 20, 30]
        out_offsets = [0, 10, 20, 30]
        in_split_sizes = [
            [10, 0, 0, 0],
            [0, 10, 0, 0],
            [0, 0, 10, 0],
            [0, 0, 0, 10],
        ]
        _get_and_check_split_sizes(world_size, in_offsets, out_offsets, in_split_sizes)

        world_size = 4
        in_offsets = [0, 3, 17, 18]
        out_offsets = [0, 10, 20, 30]
        in_split_sizes = [
            [3, 0, 0, 0],
            [7, 7, 0, 0],
            [0, 1, 0, 0],
            [0, 2, 10, 10],
        ]
        _get_and_check_split_sizes(world_size, in_offsets, out_offsets, in_split_sizes)

        world_size = 4
        in_offsets = [0, 10, 20, 30]
        out_offsets = [0, 3, 17, 18]
        in_split_sizes = [
            [3, 7, 0, 0],
            [0, 7, 1, 2],
            [0, 0, 0, 10],
            [0, 0, 0, 10],
        ]
        _get_and_check_split_sizes(world_size, in_offsets, out_offsets, in_split_sizes)

        world_size = 4
        in_offsets = [0, 7, 11, 25]
        out_offsets = [0, 10, 17, 18]
        in_split_sizes = [
            [7, 0, 0, 0],
            [3, 1, 0, 0],
            [0, 6, 1, 7],
            [0, 0, 0, 15],
        ]
        _get_and_check_split_sizes(world_size, in_offsets, out_offsets, in_split_sizes)

        # The tensor size cannot be evenly divided by the world size.
        world_size = 6
        in_offsets = [0, 7, 14, 21, 28, 35]
        out_offsets = [0, 7, 14, 21, 28, 35]
        in_split_sizes = [
            [7, 0, 0, 0, 0, 0],
            [0, 7, 0, 0, 0, 0],
            [0, 0, 7, 0, 0, 0],
            [0, 0, 0, 7, 0, 0],
            [0, 0, 0, 0, 7, 0],
            [0, 0, 0, 0, 0, 5],
        ]
        _get_and_check_split_sizes(world_size, in_offsets, out_offsets, in_split_sizes)

        world_size = 6
        in_offsets = [0, 0, 10, 11, 28, 40]
        out_offsets = [0, 7, 14, 21, 28, 35]
        in_split_sizes = [
            [0, 0, 0, 0, 0, 0],
            [7, 3, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0],
            [0, 3, 7, 7, 0, 0],
            [0, 0, 0, 0, 7, 5],
            [0, 0, 0, 0, 0, 0],
        ]
        _get_and_check_split_sizes(world_size, in_offsets, out_offsets, in_split_sizes)


class TestShardUtilsDistributed(FSDPTest):
    @property
    def world_size(self):
        return 2

    def _create_local_chunk(self, tensor):
        chunk = tensor.chunk(2)[self.rank]
        offsets = [0] if self.rank == 0 else [tensor.shape[0] - chunk.shape[0]]
        shard = Shard.from_tensor_and_offsets(chunk, offsets, self.rank)
        return init_from_local_shards([shard], tensor.numel())

    def _create_enumerate_spec(self, tensor):
        # Since placement is not used, always set placement to rank0 to mimic
        # the actual usage.
        metadata = [
            ShardMetadata([0], [101], placement="rank0/cuda:0"),
            ShardMetadata([101], [900], placement="rank0/cuda:0"),
        ]
        return EnumerableShardingSpec(metadata)

    def _create_chunk_spec(self):
        return ChunkShardingSpec(dim=0, placements=["rank0/cuda:0"])

    def _create_tensor(self, *size):
        # Keep everything deterministic.
        torch.manual_seed(0)
        return torch.rand(*size).cuda()

    @skip_if_lt_x_gpu(2)
    def test_reshard_flatten_tensor(self):
        def get_offsets(tensor, shard):
            if self.rank == 0:
                return [0]
            else:
                return [tensor.shape[0] - shard.shape[0]]

        tensor = self._create_tensor(1001)

        shard = _reshard_flatten_tensor(
            self._create_local_chunk(tensor),
            self._create_enumerate_spec(tensor),
            self.world_size,
            self.rank,
            tensor.device,
            _get_default_group(),
        )
        offsets = [0] if self.rank == 0 else [tensor.shape[0] - shard.shape[0]]
        shard = Shard.from_tensor_and_offsets(shard, offsets, self.rank)
        uneven_sharded_tensor = init_from_local_shards([shard], tensor.numel())

        shard = _reshard_flatten_tensor(
            uneven_sharded_tensor,
            self._create_chunk_spec(),
            self.world_size,
            self.rank,
            tensor.device,
            _get_default_group(),
        )
        offsets = [0] if self.rank == 0 else [tensor.shape[0] - shard.shape[0]]
        shard = Shard.from_tensor_and_offsets(shard, offsets, self.rank)
        even_sharded_tensor = init_from_local_shards([shard], tensor.numel())

        output = torch.empty(tensor.shape).cuda() if self.rank == 0 else None
        even_sharded_tensor.gather(0, output)
        if self.rank == 0:
            self.assertEqual(tensor, output)
        output = torch.empty(tensor.shape).cuda() if self.rank == 0 else None
        uneven_sharded_tensor.gather(0, output)
        if self.rank == 0:
            self.assertEqual(tensor, output)

    @skip_if_lt_x_gpu(2)
    def test_create_chunk_sharded_tensor(self):
        for size in ((1,), (1, 6), (12,), (12, 6), (25,), (25, 6)):
            tensor = self._create_tensor(*size)

            sharded_tensor = _create_chunk_sharded_tensor(
                tensor,
                self.rank,
                self.world_size,
                torch.cuda.device_count(),
                _get_default_group(),
            )
            output = torch.empty(*size).cuda() if self.rank == 0 else None
            sharded_tensor.gather(0, output)
            if self.rank == 0:
                self.assertEqual(tensor, output)