FM / _sharding

/usr/local/lib64/python3.6/site-packages/torch/distributed/_sharding_spec/

/usr/local/lib64/python3.6/site-packages/torch/distributed/_sharding_spec

Name	Size	Mode	Actions
__pycache__/	-	0755	rm
api.py	3788	0644	edit dl rm
_internals.py	5711	0644	edit dl rm
__init__.py	153	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/distributed/_sharding_spec/_internals.py (5711B)

from typing import List, Union
from dataclasses import dataclass
from torch.distributed.remote_device import _remote_device

import torch

@dataclass
class ShardMetadata(object):
    """
    Represents a shard of the overall Tensor including its
    offsets, lengths and device placement.

Args:
        shard_offsets(List[int]): Offsets in the orignal tensor indicating
            the start offsets for this shard. Should have the same rank as
            the original tensor.
        shard_lengths(List[int]): Lengths indicating the length of each
            dimension for this shard. Should have the same rank as the
            original tensor.
        placement(:class:`torch.distributed._remote_device`):
            Specifies the placement of this shard.
    """

__slots__ = ['shard_offsets', 'shard_lengths', 'placement']

shard_offsets: List[int]
    shard_lengths: List[int]
    placement: Union[str, _remote_device]

def __post_init__(self):
        if isinstance(self.placement, str):
            self.placement = torch.distributed._remote_device(self.placement)

if len(self.shard_offsets) != len(self.shard_lengths):
            raise ValueError(
                f'shard_offsets and shard_lengths should have '
                f'the same number of elements, found {len(self.shard_offsets)} '
                f'and {self.shard_lengths} respectively')

for i in range(len(self.shard_offsets)):
            if self.shard_offsets[i] < 0:
                raise ValueError('shard_offsets should be >=0')
            if self.shard_lengths[i] < 0:
                raise ValueError('shard_lengths should be >= 0')

def _check_shard_metadata_pair_overlap(shard1: ShardMetadata, shard2: ShardMetadata):
    """
    Checks if two shards overlap.
    """

# For each dim of each shard, check if one shard resides on the other
    # end of second shard with respect to that dim. As an example for a 2D
    # shard, we would check if one shard is above or on the left of the
    # other shard.
    ndims = len(shard1.shard_offsets)
    for i in range(ndims):
        if shard1.shard_offsets[i] >= shard2.shard_offsets[i] + shard2.shard_lengths[i]:
            return False
        if shard2.shard_offsets[i] >= shard1.shard_offsets[i] + shard1.shard_lengths[i]:
            return False

return True

def validate_non_overlapping_shards_metadata(shards: List[ShardMetadata]):
    """
    Ensures none of the shards overlap with each other.

Args:
        shards(List[ShardMetadata]): List of :class:`ShardMetadata` objects representing
            each shard.
    Raises:
        ``ValueError`` if there's overlap in any two shards.
    """
    # TODO: evaluate optimizing this if needed.
    for i in range(len(shards)):
        for j in range(i + 1, len(shards)):
            if _check_shard_metadata_pair_overlap(shards[i], shards[j]):
                raise ValueError(f'Shards {shards[i]} and {shards[j]} overlap')

def check_tensor(shards_metadata, tensor_dims) -> None:
    """
    Checks if the shards_metadata is compatible with the provided tensor dims.

Args:
        shards_metadata(List[ShardMetadata]): List of :class:`ShardMetadata`
            objects representing each shard of the tensor.
        tensor_dims(Sequence of int): Dimensions of tensor to verify
    Raises:
        ``ValueError`` if not compatible.
    """

# If the tensor's volume matches the total volume of all shards and
    # all shard boundaries are within tensor dims, we have a compatible
    # sharding spec for this tensor. Note that we have already verified
    # we don't have overlapping shards.
    tensor_rank = len(tensor_dims)
    shards_rank = len(shards_metadata[0].shard_offsets)
    if tensor_rank != shards_rank:
        raise ValueError(f'Rank of tensor is {tensor_rank}, but shards rank is {shards_rank}')

total_shard_volume = 0
    for shard in shards_metadata:
        shard_volume = 1
        for i, shard_length in enumerate(shard.shard_lengths):
            shard_volume *= shard_length
            if shard.shard_offsets[i] + shard.shard_lengths[i] > tensor_dims[i]:
                raise ValueError(
                    f'Shard offset {shard.shard_offsets[i]} and length '
                    f'{shard.shard_lengths[i]} exceeds tensor dim: {tensor_dims[i]} for shard {shard}')
        total_shard_volume += shard_volume

tensor_volume = 1
    for size in tensor_dims:
        tensor_volume *= size

if total_shard_volume != tensor_volume:
        # TODO: Can we improve this error message to point out the gaps?
        raise ValueError(
            f'Total volume of shards: {total_shard_volume} '
            f'does not match tensor volume: {tensor_volume}, in other words '
            f'all the individual shards do not cover the entire tensor')

def get_split_size(dim_size, chunks):
    """
    Computes the split size inline with ``torch.chunk``

Args:
        dim_size(int): Size of the dimension being chunked.
        chunks(int): Number of chunks to create for ``dim_size``.

Returns:
        An int indicating the split size to use.
    """
    return (dim_size + chunks - 1) // chunks

def get_chunked_dim_size(dim_size, split_size, idx):
    """
    Computes the dim size of the chunk for provided ``idx`` given ``dim_size``
    and ``split_size``.

Args:
        dim_size(int): Size of the dimension being chunked.
        split_size(int): The chunk size for each chunk of ``dim_size``.
        idx(int): The index of chunk whose dim size is being requested.

Returns:
        An int indicating the dim size of the chunk.
    """
    return min(dim_size, split_size * (idx + 1)) - split_size * idx

cmd: