mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-08-11 04:41:54 +00:00
26b7aac0be
* [zero] refactor low-level zero folder structure * [zero] fix legacy zero import path * [zero] fix legacy zero import path * [zero] remove useless import * [zero] refactor gemini folder structure * [zero] refactor gemini folder structure * [zero] refactor legacy zero import path * [zero] refactor gemini folder structure * [zero] refactor gemini folder structure * [zero] refactor gemini folder structure * [zero] refactor legacy zero import path * [zero] fix test import path * [zero] fix test * [zero] fix circular import * [zero] update import
23 lines
706 B
Python
23 lines
706 B
Python
from typing import Tuple
|
|
|
|
import torch
|
|
|
|
|
|
def get_shard(tensor: torch.Tensor, rank: int, world_size: int) -> Tuple[torch.Tensor, int]:
|
|
"""Return the local shard of a full tensor."""
|
|
# Shard using torch.chunk to match all-gather/reduce-scatter.
|
|
chunks = list(torch.flatten(tensor).chunk(world_size))
|
|
while len(chunks) < world_size:
|
|
chunks.append(chunks[0].new_empty(0))
|
|
|
|
# Determine number of padding elements.
|
|
num_to_pad = chunks[0].numel() - chunks[rank].numel()
|
|
assert num_to_pad >= 0, num_to_pad
|
|
|
|
shard = torch.zeros_like(chunks[0])
|
|
length = chunks[rank].size(0)
|
|
shard_temp = shard[:length]
|
|
shard_temp.copy_(chunks[rank])
|
|
|
|
return shard, num_to_pad
|