wehub-resource-sync
42 lines
1.6 KiB
Python
42 lines
1.6 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
|
|
from .types import ShardingType
|
|
from .utils import shard_param, get_shard_endpoints
|
|
|
|
|
|
def shard_unembed_param(param: torch.Tensor, shard_rank: int, num_shards: int) -> torch.Tensor:
|
|
"""
|
|
Utility method for sharding an unembed parameter. We shard unembeddings on the vocab dimension
|
|
with the expectation of an all-gather to produce the full results.
|
|
|
|
TODO(cmikeh2): Really ideal would be if MII could have access to the comm and we would do
|
|
an A2A and sharded sampling.
|
|
|
|
Args:
|
|
param (torch.Tensor): The parameter to shard. Should be of shape [vocab_size, model_dim]
|
|
shard_rank (int): Which shard of the partitioned tensor to return.
|
|
num_shards (int): The total number of shards the parameter is distributed across.
|
|
|
|
Returns:
|
|
torch.Tensor: The sharded parameter of shape [sharded_vocab_size, model_dim]
|
|
"""
|
|
return shard_param(param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, granularity=1)
|
|
|
|
|
|
def sharded_unembed_dim(vocab_size: int, shard_rank: int, num_shards: int) -> int:
|
|
"""
|
|
Utility method for determining the sharded vocab size of a sharded unembed parameter.
|
|
|
|
Args:
|
|
vocab_size (int): The size of the vocabulary.
|
|
shard_rank (int): Which shard of the partitioned tensor to return.
|
|
num_shards (int): The total number of shards the parameter is distributed across.
|
|
"""
|
|
start_idx, end_idx = get_shard_endpoints(vocab_size, shard_rank, num_shards, granularity=1)
|
|
return end_idx - start_idx
|