chore: import upstream snapshot with attribution

python/mlc_llm/support/tensor_parallel.py

@@ -0,0 +1,118 @@
+"""Sharding operators for tensor parallelism."""
+
+import dataclasses
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional  # noqa: UP035
+
+from tvm import te, tirx, topi
+from tvm.relax.frontend import nn
+
+
+@dataclasses.dataclass
+class ShardSingleDim:
+    """
+    Shard a tensor by a single dimension.
+
+
+    Parameters
+    ----------
+    name : str
+        The name of the shard func
+
+    dim : int
+        The dimension to shard
+
+    segs : Optional[List[int]]
+        The length of segments along `dim`. Default to None. If specified,
+        shard a tensor by its "segmented" dimension, where each segment has a different length
+        and sharded evenly on each worker.
+
+    """
+
+    name: str
+    dim: int
+    segs: Optional[List[int]] = None  # noqa: UP006
+
+    def gen_tir(self, shards: int, weight: nn.Tensor) -> tirx.PrimFunc:
+        """Generate a TIR function that shards the weight tensor by its rows."""
+        shape = weight.shape
+        segs = self.segs or [shape[self.dim]]
+        assert sum(segs) == shape[self.dim]
+        # NOTE: we use int64 to prevent int32 overflow
+        shape = [tirx.IntImm("int64", v) for v in shape]
+        segs = [tirx.IntImm("int64", v) for v in segs]
+        w = te.placeholder(
+            [tirx.IntImm("int64", v) for v in self._compute_in_shape(shards, weight)],
+            weight.dtype,
+            name="w",
+        )
+        ws: List[te.Tensor] = []  # noqa: UP006
+        offset = 0
+        for idx, sub_seg in enumerate(segs):
+            ws.append(
+                topi.transpose(
+                    topi.reshape(
+                        te.compute(
+                            (
+                                *shape[: self.dim],
+                                sub_seg * shards,
+                                *shape[self.dim + 1 :],
+                            ),
+                            lambda *idx: w[
+                                (
+                                    *idx[: self.dim],
+                                    idx[self.dim] + offset,
+                                    *idx[self.dim + 1 :],
+                                )
+                            ],
+                            name=f"w_{idx}",
+                        ),
+                        (
+                            *shape[: self.dim],
+                            tirx.IntImm("int64", shards),
+                            sub_seg,
+                            *shape[self.dim + 1 :],
+                        ),
+                    ),
+                    [self.dim, *range(self.dim), *range(self.dim + 1, len(shape) + 1)],
+                )
+            )
+            offset += sub_seg * shards
+        o = topi.concatenate(ws, axis=1 + self.dim)
+        func = te.create_prim_func([w, o])
+        return func
+
+    def gen_shard_info(self, shards: int, weight: nn.Tensor) -> Dict[str, Any]:  # noqa: UP006
+        """Generate shard info for this sharding strategy."""
+        return {
+            "func_name": self.name,
+            "in_shape": self._compute_in_shape(shards, weight),
+            "out_shape": (shards, *weight.shape),
+            "out_dtype": str(weight.dtype),
+        }
+
+    def _compute_in_shape(self, shards: int, weight: nn.Tensor) -> List[int]:  # noqa: UP006
+        """Compute the weight shape before sharding."""
+        shape = weight.shape
+        return [*shape[: self.dim], shape[self.dim] * shards, *shape[self.dim + 1 :]]
+
+
+@contextmanager
+def shard_bias(linear: nn.Linear, tensor_parallel_shards: int):
+    """
+    A context manager to shard the bias of a linear into `tensor_parallel_shards` shards.
+
+
+    Parameters
+    ----------
+    linear : nn.Linear
+        The linear layer whose bias would be sharded.
+
+    tensor_parallel_shards : int
+        The number of shards.
+    """
+    original_bias = linear.bias
+    if tensor_parallel_shards > 1:
+        linear.bias = linear.bias / tensor_parallel_shards
+    yield
+    linear.bias = original_bias