wehub-resource-sync
76 lines
2.8 KiB
Python
76 lines
2.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
from .types import ShardingType, DEFAULT_SHARD_GRANULARITY
|
|
from .utils import shard_param, get_shard_endpoints
|
|
|
|
|
|
def shard_mlp_1_param(param: torch.Tensor,
|
|
shard_rank: int,
|
|
num_shards: int,
|
|
gated: bool = False,
|
|
is_moe: bool = False) -> torch.Tensor:
|
|
"""
|
|
Utility method for sharding an MLP 1 parameter. Both biases and weights are supported, as well
|
|
as for fused weights for MoE.
|
|
|
|
Args:
|
|
param (torch.Tensor): The parameter to shard.
|
|
shard_rank (int): Which shard of the partitioned tensor to return.
|
|
num_shards (int): The total number of shards the parameter is distributed across.
|
|
gated (bool): Whether or not the parameter is from a gated MLP.
|
|
"""
|
|
bias_dims = 2 if is_moe else 1
|
|
|
|
if gated:
|
|
return shard_param(param,
|
|
ShardingType.OUTER_DIMENSION,
|
|
shard_rank,
|
|
num_shards,
|
|
granularity=DEFAULT_SHARD_GRANULARITY * 2,
|
|
bias_dims=bias_dims)
|
|
else:
|
|
return shard_param(param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, bias_dims=bias_dims)
|
|
|
|
|
|
def shard_mlp_2_param(param: torch.Tensor,
|
|
shard_rank: int,
|
|
num_shards: int,
|
|
is_moe: bool = False) -> Optional[torch.Tensor]:
|
|
"""
|
|
Utility method for sharding an MLP 2 parameter.
|
|
|
|
Args:
|
|
param (torch.Tensor): The parameter to shard.
|
|
shard_rank (int): Which shard of the partitioned tensor to return.
|
|
num_shards (int): The total number of shards the parameter is distributed across.
|
|
is_moe (bool): Whether or not the parameter is from a MoE model.
|
|
"""
|
|
bias_dim_size = 2 if is_moe else 1
|
|
|
|
if len(param.shape) == bias_dim_size:
|
|
# We will do the bias addition on the 0th rank only rather than scale the parameter and
|
|
# implicitly reconstruct this in the distributed reduce.
|
|
return param if shard_rank == 0 else None
|
|
|
|
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards)
|
|
|
|
|
|
def sharded_intermediate_dim(intermediate_size: int, num_shards: int, shard_rank: int) -> int:
|
|
"""
|
|
Utility method for getting the size of the intermediate dimension of a sharded MLP.
|
|
|
|
Args:
|
|
intermediate_size (int): The size of the intermediate dimension.
|
|
num_shards (int): The total number of shards the parameter is distributed across.
|
|
shard_rank (int): Which shard of the partitioned tensor to return.
|
|
"""
|
|
endpoints = get_shard_endpoints(intermediate_size, shard_rank, num_shards)
|
|
return endpoints[1] - endpoints[0]
|