mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
00ffeca1b1
See #145101 for details. Pull Request resolved: https://github.com/pytorch/pytorch/pull/145164 Approved by: https://github.com/bobrenjc93
33 lines
1.0 KiB
Python
33 lines
1.0 KiB
Python
from collections.abc import Sequence
|
|
|
|
import torch
|
|
from torch.distributed._shard.metadata import ShardMetadata
|
|
|
|
|
|
DEPRECATE_MSG = "Please use DTensor instead and we are deprecating ShardedTensor."
|
|
|
|
|
|
def narrow_tensor_by_index(
|
|
tensor: torch.Tensor,
|
|
offsets: Sequence[int],
|
|
sizes: Sequence[int],
|
|
) -> torch.Tensor:
|
|
"""
|
|
Narrow the tensor according to ``offsets`` and ``sizes``.
|
|
"""
|
|
narrowed_tensor = tensor
|
|
for idx, (offset, size) in enumerate(zip(offsets, sizes)):
|
|
if size < tensor.size(idx):
|
|
# Reshape to get shard for this rank and we don't want autograd
|
|
# recording here for the narrow op and 'local_shard' should be a
|
|
# leaf variable in the autograd graph.
|
|
narrowed_tensor = narrowed_tensor.narrow(idx, offset, size)
|
|
return narrowed_tensor
|
|
|
|
|
|
def narrow_tensor(tensor: torch.Tensor, metadata: ShardMetadata) -> torch.Tensor:
|
|
"""
|
|
Narrow the tensor according to the metadata
|
|
"""
|
|
return narrow_tensor_by_index(tensor, metadata.shard_offsets, metadata.shard_sizes)
|