def test_get_chunk_sharding_params(self):
ranks = [
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
]
spec = ChunkShardingSpec(
dim=0,
placements=ranks,
)
result = get_chunk_sharding_params(21, 4, spec, 1)
self.assertEqual(6, result[0])
self.assertEqual(6, result[1])
result = get_chunk_sharding_params(21, 4, spec, 3)
self.assertEqual(18, result[0])
self.assertEqual(3, result[1])
ranks[1], ranks[2] = ranks[2], ranks[1]
ranks[0], ranks[3] = ranks[3], ranks[0]
spec.placements = ranks
result = get_chunk_sharding_params(21, 4, spec, 1)
self.assertEqual(12, result[0])
self.assertEqual(6, result[1])
result = get_chunk_sharding_params(21, 4, spec, 3)
self.assertEqual(0, result[0])
self.assertEqual(6, result[1])
def _handle_col_wise_sharding(input, world_size, weight, rank, local_shard_t,
bias, pg):
"""
Entry-point function to handle the logic of col-wise sharding of weight
for Linear. (Detailed explanations of the logic can be found in the
comment for sharded_linear.)
When the local tensor only has one dimension, we increase one more dimension
for reshard. We need to do squeeze manually to reduce the dimension later-on.
For example, if we have:
input: size[15]
weight: size[15, 16]
world_size: 4
In each rank, we will have 4 * [4] tensors. We then stack them into a [4, 4]
tensor and generate a sharded tenor sharded by dim 1.
For the rest situations, we just simply concatenate local tensors. No more actions
are needed afterward.
Args:
input: matrix to be multiplied with the sharded weight.
world_size: number of ranks.
weight: shareded weight tensor.
rank: # of cuda process.
local_shard_t: row-wise shared local weight used for lookup.
bias: bias term of linear op.
pg: process group.
Returns:
A :class:`ShardedTensor` object which filled with local intermediate results.
"""
# allgather the inputs first.
gathered_inputs = all_gather(input, group=pg)
(start_pos,
chunk_size) = get_chunk_sharding_params(bias.size(0), world_size,
weight._sharding_spec, rank)
local_bias = _BiasTensorNarrow.apply(world_size, start_pos, chunk_size,
weight, pg, bias)
results = [None] * world_size
indices = {}
for idx, placement in enumerate(weight._sharding_spec.placements):
indices[placement.rank()] = idx
for i, inp in enumerate(gathered_inputs):
results[indices[i]] = inp.matmul(local_shard_t) + local_bias
# When the local result only has one dimension, we need to make sure
# it does not shard by dim 0. So reshard can work properly.
if results[0].dim() == 1: # type: ignore[attr-defined]
result = torch.stack(results) # type: ignore[arg-type]
else:
result = torch.cat(results) # type: ignore[arg-type]
return _init_sharded_tensor_from_local_result(
weight,
result,
0,
-1,
world_size,
pg # type: ignore[arg-type]
)
def _handle_row_wise_mask(gather_inp, padding_idx, weight, world_size, rank):
"""
Mask the input for embedding look-up for IDs which are not stored
on the current rank. This function also adjust the ``padding_idx``
so that it is only used on the rank where the corresponding row is
stored.
Note that, with ``max_norm`` flag on, only weights of rows being
looked up will be re-normed. So we need an extra row for masked ID
so that it does not affect the final result and ``max_norm``.
Args:
gather_inp: tensor to be applied op on gathered from all ranks.
padding_idx: If specified, the entries at padding_idx do
not contribute to the gradient; therefore, the embedding
vector at padding_idx is not updated during training,
i.e. it remains as a fixed “pad”.
Note that the embedding vector at padding_idx is
excluded from the reduction.
weight: weight tensor of Embedding look-up table.
world_size: number of ranks.
rank: # of cuda process.
Returns:
lookup_input: Tensor of masked input.
padding_idx: adjusted padding_idx.
padding_row: The extra row we used during lookup so that
looking up does not affect ``max_norm``.
"""
(start_pos,
chunk_size) = get_chunk_sharding_params(weight.size(0), world_size,
weight._sharding_spec, rank)
mask = (gather_inp < start_pos) | (gather_inp >= start_pos + chunk_size)
lookup_input = gather_inp.clone() - start_pos
lookup_input[mask] = chunk_size
if (padding_idx is not None and padding_idx >= start_pos and padding_idx <
(start_pos + chunk_size)):
padding_idx = padding_idx - start_pos
else:
padding_idx = None
# When max_norm is set, it will only re-norm the row being looked up.
padding_row = torch.zeros(1,
weight.size(1),
device=gather_inp.device,
dtype=weight.dtype)
return lookup_input, padding_idx, padding_row
def _handle_col_wise_sharding(input, world_size, weight, rank, local_shard_t,
bias, pg):
"""
Entry-point function to handle the logic of col-wise sharding of weight
for Linear. (Detailed explanations of the logic can be found in the
comment for sharded_linear.)
When the local tensor only has one dimension, we increase one more dimension
for reshard. We need to do squeeze manually to reduce the dimension later-on.
For example, if we have:
input: size[15]
weight: size[15, 16]
world_size: 4
In each rank, we will have 4 * [4] tensors. We then stack them into a [4, 4]
tensor and generate a sharded tenor sharded by dim 1.
For the rest situations, we just simply concatenate local tensors. No more actions
are needed afterward.
Args:
input: matrix to be multiplied with the sharded weight.
world_size: number of ranks.
weight: shareded weight tensor.
rank: # of cuda process.
local_shard_t: row-wise shared local weight used for lookup.
bias: bias term of linear op.
pg: process group.
Returns:
A :class:`ShardedTensor` object which filled with local intermediate results.
"""
# allgather the inputs first.
out_size = list(input.size())
out_size[0] = input.size(0) * dist.get_world_size(pg)
output = torch.empty(out_size, device=input.device)
output = _all_gather_base(output, input, group=pg)
# Adjust bias and perform local matmul.
(start_pos,
chunk_size) = get_chunk_sharding_params(bias.size(0), world_size,
weight._sharding_spec, rank)
local_bias = _BiasTensorNarrow.apply(world_size, start_pos, chunk_size,
weight, pg, bias)
if output.dim() == 1:
output = output.view(dist.get_world_size(pg), -1)
if output.dim() <= 2:
# Use fused version if possible.
result = torch.addmm(local_bias, output, local_shard_t)
else:
result = output.matmul(local_shard_t) + local_bias
# Build ShardedTensor as result.
st_size = list(result.size())
st_size[-1] = weight.size(0)
new_sharding_spec = ChunkShardingSpec(
dim=-1, placements=weight.sharding_spec().placements)
return ShardedTensor._init_from_local_tensor(
result,
new_sharding_spec,
*st_size, # type: ignore[arg-type]
process_group=pg,
)
def sharded_layer_norm(args, kwargs, pg):
"""
Handles ``__torch_function__`` dispatch for the ``torch.nn.LayerNorm`` op.
We gather all shards from local shards and perform a global normalization.
We then scatter the result back to each rank.
Args: same as ``torch.nn.LayerNorm``.
Return:
local_tensor (Tensor): New local tensor to build the sharded tensor.
sharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`):
sharding spec of the new sharded tensor.
new_st_size (torch.Size): Size of the new sharded tensor.
"""
st = args[0]
normalized_shape = args[1]
sharding_dim = st.sharding_spec().dim # type: ignore[attr-defined]
sharding_dim = sharding_dim if sharding_dim >= 0 else st.dim(
) + sharding_dim
local_tensor = st.local_tensor()
# If sharding dim is smaller than shape start, we just perform a local norm.
shape_start = st.dim() - len(normalized_shape)
if shape_start > sharding_dim:
args = (local_tensor, *args[1:])
local_tensor = torch.nn.functional.layer_norm(*args, **kwargs)
return local_tensor, st.sharding_spec(), st.size()
elementwise_affine = kwargs.get("elementwise_affine", False)
eps = kwargs.get("eps", 1e-05)
norm_dims = tuple(i for i in range(-1, -len(normalized_shape) - 1, -1))
local_size = math.prod(
local_tensor.size()[shape_start:]) # type: ignore[attr-defined]
st_size = math.prod(st.size()[shape_start:]) # type: ignore[attr-defined]
local_mean = torch.mul(local_tensor.mean(norm_dims, keepdim=True),
local_size)
global_mean = torch.div(all_reduce(local_mean), st_size)
local_variant_sq = torch.square(local_tensor - global_mean).sum(
norm_dims, keepdim=True)
global_variant = torch.div(all_reduce(local_variant_sq), st_size)
denom = torch.rsqrt(global_variant + eps)
local_tensor = torch.mul(local_tensor - global_mean, denom)
if elementwise_affine:
weight = kwargs["weight"]
bias = kwargs["bias"]
current_rank = dist.get_rank(pg) # type: ignore[attr-defined]
world_size = dist.get_world_size(pg)
(start_pos,
chunk_size) = get_chunk_sharding_params(bias.size(0), world_size,
st.sharding_spec(),
current_rank)
local_tensor = torch.addmm(
torch.narrow(bias, 0, start_pos, chunk_size),
local_tensor,
torch.narrow(weight, sharding_dim - shape_start, start_pos,
chunk_size),
)
return local_tensor, st.sharding_spec(), st.size()