Class ShardedTensorLoader

accessory/util/tensor_parallel.py:526–575

r"""A helper class to load a tensor parallel sharded tensor and track the loading status (i.e., check if a tensor is loaded with consistent tensor parallel size, check that each shard is loaded only once, and check that all shards are loaded).

class ShardedTensorLoader:
    r"""A helper class to load a tensor parallel sharded tensor and track the
    loading status (i.e., check if a tensor is loaded with consistent tensor
    parallel size, check that each shard is loaded only once, and check that
    all shards are loaded).
    """

    def __init__(
        self,
        target: torch.Tensor,
        num_shards: int,
        shard_dim: int,
        mode: str = "set",
    ) -> None:
        r"""Initialize a ShardedTensorLoader.

        Args:
            target (``torch.Tensor``): The target tensor where value shards are
                loaded into.
            num_shards (int): Number of expected shards.
            shard_dim (int): The dimension along which the tensor is sharded.
            mode (str): Supported options are ``set`` and ``add``. If ``set``,
                the old value in the target tensor is overrided with the new
                value. If ``add``, the new value is added to the old value.
        """
        self._target = target
        self._num_shards = num_shards
        self._shard_dim = shard_dim
        self._mode = mode

        self._loaded_shards = set()

    def load_shard(self, shard_id: int, value: torch.Tensor) -> None:
        r"""Load a shard into the target tensor.

        Args:
            shard_id (int): The shard id of the current value.
            value (``torch.Tensor``): The value to be loaded.
        """
        assert shard_id not in self._loaded_shards
        assert shard_id >= 0 and shard_id < self._num_shards
        self._loaded_shards.add(shard_id)
        tensor_load_shard(self._target, self._shard_dim, self._num_shards,
                          shard_id, value, self._mode)

    def is_complete(self) -> bool:
        r"""Check if all the shards are loaded to the target tensor."""
        assert all(x >= 0 and x < self._num_shards
                   for x in self._loaded_shards)
        return len(self._loaded_shards) == self._num_shards