chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,123 @@
+"""Functions for pre-sharding weights"""
+
+import logging
+from collections.abc import Sequence
+from typing import Any, Callable, Dict, Tuple  # noqa: UP035
+
+from tvm import IRModule, relax
+from tvm.relax.frontend import nn
+from tvm.runtime import Device, Tensor
+from tvm.s_tir import dlight as dl
+from tvm.target import Target
+
+logger = logging.getLogger("preshard")
+
+
+def _sharded_param_name(param_name, worker_id):
+    return f"{param_name}_shard-{worker_id}"
+
+
+def _create_shard_func(bb: relax.BlockBuilder, param: nn.Parameter, tensor_parallel_shards: int):
+    shard_strategy = param.attrs.get("shard_strategy", None)
+    # generate tirx shard function
+    tir_func = shard_strategy.gen_tir(shards=tensor_parallel_shards, weight=param)
+    tir_func = tir_func.with_attr("global_symbol", f"{shard_strategy.name}_tir")
+    # add tirx shard function to the IRModule
+    tir_gvar = bb.add_func(tir_func, func_name=f"{shard_strategy.name}_tir")
+    # create relax function that
+    #     1. shard weight with tirx shard function, result: [num_shards, *sharded_weight_shape]
+    #     2. split the sharded weight along dim 0, result: num_shards * [1, *sharded_weight_shape]
+    #     3. squeeze the 0th-dim of all shards, result: num_shards * [*sharded_weight_shape]
+    weight_shape = param.shape
+    weight_shape[shard_strategy.dim] = weight_shape[shard_strategy.dim] * tensor_parallel_shards
+    sharded_weight_shape = [tensor_parallel_shards, *param.shape]
+    weight_var = relax.Var("weight", relax.TensorType(weight_shape, param.dtype))
+    with bb.function(name=shard_strategy.name, params=[weight_var]):
+        with bb.dataflow():
+            lv0 = bb.emit(
+                relax.call_tir(
+                    tir_gvar,
+                    weight_var,
+                    out_ty=relax.TensorType(sharded_weight_shape, param.dtype),
+                )
+            )
+            lv1 = bb.emit(relax.op.split(lv0, indices_or_sections=tensor_parallel_shards, axis=0))
+            output_vars = []
+            for i in range(tensor_parallel_shards):
+                lvi = bb.emit(relax.TupleGetItem(lv1, i))
+                squeezed_lvi = bb.emit(relax.op.squeeze(lvi, 0))
+                output_vars.append(squeezed_lvi)
+            gv = bb.emit_output(output_vars)
+        bb.emit_func_output(gv)
+
+
+def _compile_shard_funcs(mod: IRModule, device: Device):
+    target = Target.from_device(device)
+    with target:
+        mod = relax.transform.LegalizeOps()(mod)
+        mod = dl.ApplyDefaultSchedule(
+            dl.gpu.Matmul(),
+            dl.gpu.GEMV(),
+            dl.gpu.Reduction(),
+            dl.gpu.GeneralReduction(),
+            dl.gpu.Fallback(),
+        )(mod)
+    ex = relax.build(mod, target=target)
+    vm = relax.VirtualMachine(ex, device)
+    return vm
+
+
+def apply_preshard(
+    named_params: Dict[str, nn.Parameter],  # noqa: UP006
+    tensor_parallel_shards: int,
+    args: Any,
+) -> Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]]]:  # noqa: UP006
+    """Apply pre-sharding to the named parameters.
+
+    Parameters
+    ----------
+    named_params : Dict[str, nn.Parameter]
+        The named parameters of the model. If the model is quantized, the named parameters should
+        the state dictionary of the quantized model.
+    tensor_parallel_shards : int
+        The number of tensor parallel shards.
+    args : Any
+        The parsed arguments of weight conversion.
+
+    Returns
+    -------
+    Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]]
+        The updated named parameters and the mapping from parameter name to the shard function.
+    """
+    bb = relax.BlockBuilder()
+    param_to_shard_func = {}
+    shard_func_names = set()
+    new_named_params: Dict[str, nn.Parameter] = {}  # noqa: UP006
+    has_shard_strategy = False
+    for name, param in named_params.items():
+        shard_strategy = param.attrs.get("shard_strategy", None)
+        if shard_strategy is not None:
+            has_shard_strategy = True
+            for i in range(tensor_parallel_shards):
+                new_named_params[_sharded_param_name(name, i)] = param
+            # create shard functions
+            param_to_shard_func[name] = shard_strategy.name
+            if shard_strategy.name not in shard_func_names:
+                _create_shard_func(bb, param, tensor_parallel_shards)
+                shard_func_names.add(shard_strategy.name)
+        else:
+            new_named_params[name] = param
+
+    if not has_shard_strategy:
+        logger.warning(
+            "No parameters with 'shard_strategy' found."
+            "At least one parameter must have a 'shard_strategy' for presharding. "
+            "The model will continue to convert weights in a non-presharded manner."
+        )
+
+    mod = bb.finalize()
+    vm = _compile_shard_funcs(mod, args.device)
+
+    for name in param_to_shard_func:
+        param_to_shard_func[name] = vm[param_to_shard_func[name]]
+    return new_named_params, param_to_shard_func