def __init__(
self,
sharding_spec: shard_spec.ShardingSpec,
*size,
dtype=None,
layout=torch.strided,
requires_grad=False,
pin_memory=False,
memory_format=torch.contiguous_format,
process_group=None,
init_rrefs=False,
):
# prepare initialization, initialize fields like
# _process_group, _local_shards, etc.
self._prepare_init(process_group=process_group, init_rrefs=init_rrefs)
tensor_properties = TensorProperties(dtype, layout, requires_grad, memory_format, pin_memory)
if tensor_properties is None:
raise ValueError('tensor_properties must not be None.')
if tensor_properties.dtype is None:
tensor_properties.dtype = torch.get_default_dtype()
if tensor_properties.layout != torch.strided:
raise ValueError('Only torch.strided layout is currently supported')
if tensor_properties.memory_format != torch.contiguous_format:
raise ValueError('Only torch.contiguous_format memory_format is currently supported')
dims = _flatten_tensor_size(size)
if not isinstance(sharding_spec, shard_spec.ShardingSpec):
raise ValueError(f'Expecting ShardingSpec but got: {type(sharding_spec)}')
self._sharding_spec = sharding_spec
sharded_tensor_metadata = sharding_spec.build_metadata(
dims, tensor_properties=tensor_properties)
current_rank = dist.get_rank(self._process_group)
for shard_metadata in sharded_tensor_metadata.shards_metadata:
rank, device = _parse_and_validate_remote_device(self._process_group, shard_metadata.placement)
if rank == current_rank:
local_tensor = _create_tensor_from_params(
shard_metadata.shard_sizes,
local_device=device,
tensor_properties=sharded_tensor_metadata.tensor_properties
)
self._local_shards.append(Shard(local_tensor, shard_metadata))
self._metadata = sharded_tensor_metadata
# do post initialization (i.e. register sharded_tensor_id, initialize_rpc)
self._post_init()
def __new__(cls, sharding_spec: shard_spec.ShardingSpec, *size, **kwargs):
# Use __new__ to construct a wrapper tensor, for recording tensor
# properties and logging purposes.
torch._C._log_api_usage_once("torch.distributed._shard.sharded_tensor")
# check sharding spec and build sharded tensor metadata
if not isinstance(sharding_spec, shard_spec.ShardingSpec):
raise ValueError(
f"Expecting ShardingSpec but got: {type(sharding_spec)}")
sizes = _flatten_tensor_size(size)
dtype = kwargs["dtype"]
layout = kwargs["layout"]
pin_memory = kwargs["pin_memory"]
requires_grad = kwargs["requires_grad"]
if dtype is None:
dtype = torch.get_default_dtype()
tensor_properties = TensorProperties(dtype,
layout,
requires_grad,
pin_memory=pin_memory)
sharded_tensor_metadata = sharding_spec.build_metadata(
sizes, tensor_properties=tensor_properties)
r = torch.Tensor._make_wrapper_subclass( # type: ignore[attr-defined]
cls,
sizes,
dtype=dtype,
layout=layout,
pin_memory=pin_memory,
requires_grad=requires_grad,
)
# set sharding spec
r._sharding_spec = sharding_spec
# set metadata
r._metadata = sharded_tensor_metadata
# set local shards
r._local_shards = []
return r
def _init_from_local_tensor(
cls,
local_tensor: torch.Tensor,
sharding_spec: shard_spec.ShardingSpec,
*global_size: Sequence[int],
process_group: dist.ProcessGroup = None,
init_rrefs=False,
) -> "ShardedTensor":
"""
Initialize a ShardedTensor given only one local tensor, global sharded tensor
size and sharding spec on each rank.
Args:
local_tensor (Tensor): Single tensor of local shard stored in each rank.
sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`):
The specification describing how to shard the Tensor.
global_size (Sequence[int]): Size of the sharded tensor.
process_group (ProcessGroup, optional): The process group to aggregate on.
Default: None
init_rrefs (bool, optional): Whether or not to initialize
:class:`torch.distributed.rpc.RRef`s pointing to remote shards.
Need to initialize the RPC Framework if specified as ``True``.
Default: ``False``.
Returns:
A :class:`ShardedTensor` sharded based on the given sharding_spec with local
tensor stored in the current rank.
Examples:
>>> # All tensors below are of torch.int64 type.
>>> # We have 2 process groups, 2 ranks.
>>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
>>> local_tensor = torch.unsqueeze(torch.cat([tensor, tensor + 2]))
>>> local_tensor
tensor([[1, 2, 3, 4]]) # Rank 0
tensor([[3, 4, 5, 6]]) # Rank 1
>>> sharding_dim = 0
>>> sharding_spec = ChunkShardingSpec(
dim=sharding_dim,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
],
)
>>> st = ShardedTensor._init_from_local_tensor(local_tensor, sharding_spec, [2, 4])
>>> st
ShardedTensor(
ShardedTensorMetadata(
shards_metadata=[
ShardMetadata(shard_offsets=[0, 0], shard_sizes=[1, 4], placement=rank:0/cuda:0),
ShardMetadata(shard_offsets=[1, 0], shard_sizes=[1, 4], placement=rank:1/cuda:1),
],
size=torch.Size([2, 4])
)
>>> st.local_tensor()
tensor([1, 2, 3, 4]) # Rank 0
tensor([3, 4, 5, 6]) # Rank 1
Warning: This API is experimental and subject to change. It lacks of a fully across
rank validations, and we only validate the local shard on the current rank.
We fully rely on the user to ensure local tensor is sharded based on the
sharding spec.
"""
if not local_tensor.is_contiguous():
raise ValueError('local_tensor is not a contiguous Tensor.')
global_tensor_size = _flatten_tensor_size(global_size)
tensor_properties = TensorProperties(
dtype=local_tensor.dtype,
layout=local_tensor.layout,
requires_grad=local_tensor.requires_grad,
memory_format=torch.contiguous_format,
pin_memory=local_tensor.is_pinned())
sharded_tensor_metadata = sharding_spec.build_metadata(
global_tensor_size,
tensor_properties
)
process_group = (
process_group
if process_group is not None
else distributed_c10d._get_default_group()
)
current_rank = dist.get_rank(process_group)
local_shards: List[Shard] = []
for shard_metadata in sharded_tensor_metadata.shards_metadata:
rank, device = _parse_and_validate_remote_device(process_group, shard_metadata.placement)
if rank == current_rank:
local_shards.append(Shard(local_tensor, shard_metadata))
# TODO: figure out what the API should behave when some rank have no shard
# see https://github.com/pytorch/pytorch/issues/7313
return ShardedTensor._init_from_local_shards_and_global_metadata(
local_shards,
sharded_tensor_metadata,
process_group=process_group,
init_rrefs=init_rrefs,
sharding_spec=sharding_spec,
)
