chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,104 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional, Tuple
+
+import torch
+
+from .types import ShardingType, DEFAULT_SHARD_GRANULARITY
+
+
+def get_shard_endpoints(dim_size: int,
+                        shard_rank: int,
+                        num_shards: int,
+                        granularity: int = DEFAULT_SHARD_GRANULARITY) -> Tuple[int, int]:
+    """
+    Given a dimension to shard with size dim_size, return the start and end indices of the slice
+    that belong to the given rank.
+
+    The typical use of this is as an internal helper function, so see if there is a higher level
+    API that better suits the application.
+
+    Args:
+        dim_size (int): The size of the dimension to shard.
+        shard_rank (int): The rank of the shard to return.
+        num_shards (int): Total number of shards the dimension will be distributed across.
+        granularity (int): The minimum alignment of the shard endpoints. This is used to support
+            non-even head counts as well as align dimensions to cleaner GEMM boundaries.
+    """
+    assert dim_size % granularity == 0, "Dimension size must be divisible by granularity"
+
+    total_chunks = dim_size // granularity
+    base_chunks_per_rank = total_chunks // num_shards
+    remainder_chunks = total_chunks % num_shards
+
+    start_chunk_id = shard_rank * base_chunks_per_rank + min(shard_rank, remainder_chunks)
+    end_chunk_id = start_chunk_id + base_chunks_per_rank + (1 if shard_rank < remainder_chunks else 0)
+
+    return start_chunk_id * granularity, end_chunk_id * granularity
+
+
+def shard_param(param: Optional[torch.Tensor],
+                shard_mode: ShardingType,
+                shard_rank: int,
+                num_shards: int,
+                num_concatenated_matrices: int = 1,
+                granularity: int = 32,
+                bias_dims: int = 1) -> torch.Tensor:
+    """
+    Utility for sharding a parameter. This will return the slice of the parameter that should
+    exist on the given shard_rank given the sharding configuration. The workflow here is
+    to find the minimum bounded Tensor to shard, get the slicing endpoints, and then concatenate
+    as needed.
+
+    The typical use of this is as an internal helper function, so see if there is a higher level
+    API that better suits the application.
+
+    Args:
+        param (torch.Tensor): The parameter to shard.
+        shard_mode (ShardingType): The type of sharding to apply. See ShardingType for more context.
+        shard_rank (int): The rank of the shard to return.
+        num_shards (int): Total number of shards the parameter will be distrbuted across.
+        num_concatenated_matrices (int): The number of matrices that have been concatenated together in the original
+            parameter. An example of this is a fused QKV projection matrix, where the `num_concatenated_matrices`
+            argument would be 3.
+        granularity (int): The minimum alignment of the shard endpoints. For attention projection matrices, this
+            should be set to the head size to support non-even sharding.
+        bias_dims (int): The number of dimensions that are considered bias dimensions. This is used to support
+            sharding of MoE and non-MoE biases on the same codepath.
+    """
+    assert shard_rank < num_shards, "Shard rank must be less than num_shards"
+
+    # Easier to hide this inside of the sharding logic than to add checks in every model
+    # implementation.
+    if param is None:
+        return None
+
+    if num_shards == 1:
+        # Trivial case of no sharding.
+        return param
+
+    if shard_mode == ShardingType.OUTER_DIMENSION:
+
+        def get_matrices(dim_idx: int) -> torch.Tensor:
+            dim_size = param.size(dim_idx) // num_concatenated_matrices
+            start_channel_id, end_channel_id = get_shard_endpoints(dim_size, shard_rank, num_shards, granularity)
+            return torch.chunk(param, num_concatenated_matrices, dim=dim_idx), start_channel_id, end_channel_id
+
+        if param.ndim == bias_dims:
+            # Special case for bias parameters.
+            matrices, start_channel_id, end_channel_id = get_matrices(dim_idx=-1)
+            return torch.cat([mat[..., start_channel_id:end_channel_id] for mat in matrices], dim=-1)
+        else:
+            # General case for weight parameters. This assumes MoE parameters are stored in the format of
+            # [num_experts, out_features, in_features]
+            matrices, start_channel_id, end_channel_id = get_matrices(dim_idx=-2)
+            return torch.cat([mat[..., start_channel_id:end_channel_id, :] for mat in matrices], dim=-2)
+
+    elif shard_mode == ShardingType.INNER_DIMENSION:
+        dim_size = param.size(-1) // num_concatenated_matrices
+        start_channel_id, end_channel_id = get_shard_endpoints(dim_size, shard_rank, num_shards, granularity)
+        matrices = torch.chunk(param, num_concatenated_matrices, dim=-1)
+        return torch.cat([mat[..., start_channel_id:end_channel_id] for mat in matrices], dim=-1)