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,13 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .attn_output_parameters import *
+from .embedding_parameters import *
+from .mlp_parameters import *
+from .moe_parameters import *
+from .norm_parameters import *
+from .qkv_parameters import *
+from .unembed_parameters import *
+from .invfreq_parameters import *
@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase
+"""
+Common Attention Output Parameter Patterns
+"""
+
+
+class AttentionOutputParameter(ParameterBase):
+    """
+    Attention output parameter container.
+
+    Note: The differentiation for something like GQA for this matrix is primarily
+    encompassed in the sharding logic, which is currently expected to be performed by
+    the model implementation.
+    """
+
+    params: torch.Tensor
+    """
+    Unsharded attention output parameter of shape [model_dim, model_dim]
+    """
+
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform_attn_out_param(self.params)
@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase
+"""
+Embedding containers.
+"""
+
+
+class EmbeddingParameter(ParameterBase):
+    """
+    Embedding container. This should be safe to use for all types of embeddings (i.e. word, position,
+    and token type).
+    """
+
+    params: torch.Tensor
+    """
+    Vocabulary parameter of shape [vocab_size, model_dim].
+    """
+
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform_embedding_param(self.params)
@@ -0,0 +1,19 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase
+"""
+Common InvFreq Parameter Patterns
+"""
+
+
+class InvFreqParameter(ParameterBase):
+
+    params: torch.Tensor
+
+    def finalize(self) -> torch.Tensor:
+        return self.params.to(self.inference_model.activation_dtype.value)
@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase
+"""
+MLP Parameter Containers
+"""
+
+
+class MLP1Parameter(ParameterBase):
+    """
+    First MLP projection weight container. This performs a straight pass-through to the
+    model implementation for transformation.
+    """
+    params: torch.Tensor
+
+    def finalize(self) -> torch.Tensor:
+        # NOTE(cmikeh2): If we are gated but not in the format specified below, we should trigger a permutation here.
+        # I am not currently aware of any models that use this format (or how we should even detect it; probably should
+        # just be a different param entirely, but until then we'll just assume the format is correct).
+        return self.inference_model.transform_mlp_1_param(self.params)
+
+
+class GatedMLPParameter(ParameterBase):
+    """
+    Gated MLP projection container.
+    """
+
+    gate_params: torch.Tensor
+    """
+    Weight parameter for the gating matrix.
+    """
+
+    up_params: torch.Tensor
+    """
+    For lack of a better name, the non-gating weight parameters.
+    """
+
+    def finalize(self) -> torch.Tensor:
+        """
+        Our gated format (this is different from InferenceV1!) is to have the gate and activated neurons
+        interleaved. So if we have 4 output neurons (two effective neurons) with 4 input neurons, the finalized
+        parameter will look like:
+        [g0_0, g0_1, g0_2, g0_3]
+        [a0_0, a0_1, a0_2, a0_3]
+        [g1_0, g1_1, g1_2, g1_3]
+        [a1_0, a1_1, a1_2, a1_3]
+
+        As a reference, in inference v1, the format is:
+        [g0_0, g0_1, g0_2, g0_3]
+        [g1_0, g1_1, g1_2, g1_3]
+        [a0_0, a0_1, a0_2, a0_3]
+        [a1_0, a1_1, a1_2, a1_3]
+        """
+        assert self.gate_params.shape[0] == self.up_params.shape[
+            0], "Gated MLP parameters must have the same number of neurons."
+        total_neurons = self.gate_params.shape[0] + self.up_params.shape[0]
+
+        # flip the order if even with the correct tokenizer we get wrong output
+        #fused_param = torch.cat([self.up_params, self.gate_params], dim=-1).reshape(total_neurons, -1)
+        fused_param = torch.cat([self.gate_params, self.up_params], dim=-1).reshape(total_neurons, -1)
+        return self.inference_model.transform_mlp_1_param(fused_param)
+
+
+class FusedGatedMLPParameter(ParameterBase):
+    """
+    Gated MLP projection container.
+    """
+
+    params: torch.Tensor
+    """
+    Weight parameter for the fused gating and non-gating weight parameters.
+    """
+
+    def finalize(self) -> torch.Tensor:
+        gate_params = self.params[:self.params.shape[0] // 2]
+        up_params = self.params[self.params.shape[0] // 2:]
+        total_neurons = gate_params.shape[0] + up_params.shape[0]
+        fused_param = torch.cat([gate_params, up_params], dim=-1).reshape(total_neurons, -1)
+        return self.inference_model.transform_mlp_1_param(fused_param)
+
+
+class MLP2Parameter(ParameterBase):
+    """
+    Second MLP projection weight container. This performs a straight pass-through to the
+    model implementation for transformation.
+    """
+
+    params: torch.Tensor
+    """
+    Full weight parameter.
+    """
+
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform_mlp_2_param(self.params)
@@ -0,0 +1,78 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase, ParamList
+"""
+Moe Parameters
+
+These parameters are compatible with any model inheriting from ``DSMoETransformerModelBase``.
+"""
+
+
+class MoEGatingWeightParameter(ParameterBase):
+    """
+    Gating weight matrix.
+    """
+
+    params: torch.Tensor
+    """
+    Projection matrix from the input activations to the gate logits.
+    """
+
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform_moe_gate_param(self.params)
+
+
+class UnfusedMoEMLP1Parameter(ParameterBase):
+    """
+    This container should be used when the experts are held in separate parameters
+    and need to be joined into a single group.
+    """
+
+    experts: ParamList("n_experts")  # noqa: F821
+
+    def finalize(self) -> torch.Tensor:
+        stacked_experts = torch.stack([p for p in self.experts], dim=0)
+        return self.inference_model.transform_moe_mlp_1_param(stacked_experts)
+
+
+class UnfusedMoEMLP2Parameter(ParameterBase):
+    """
+    This container should be used when the experts are held in separate parameters
+    and need to be joined into a single group.
+    """
+
+    experts: ParamList("n_experts")  # noqa: F821
+
+    def finalize(self) -> torch.Tensor:
+        stacked_experts = torch.stack([p for p in self.experts], dim=0)
+        return self.inference_model.transform_moe_mlp_2_param(stacked_experts)
+
+
+class UnfusedMoEGatedMLPParameter(ParameterBase):
+    """
+    MoE Parameter for a gated activation function in which the gating matrix is not
+    fused in the same parameter as the non-gating matrix.
+
+    This is a stacked version of the ``GatedMLPParameter``. Please see that class for more
+    documentation on the layout of the parameters.
+    """
+
+    gating_experts: ParamList("n_experts")  # noqa: F821
+
+    up_experts: ParamList("n_experts")  # noqa: F821
+
+    def finalize(self) -> torch.Tensor:
+        transposed_experts = []
+        for gate, up in zip(self.gating_experts, self.up_experts):
+            assert gate.shape[0] == up.shape[0], "Gated MLP parameters must have the same number of neurons."
+            total_neurons = gate.shape[0] + up.shape[0]
+            fused_expert = torch.cat([gate, up], dim=-1).reshape(total_neurons, -1)
+            transposed_experts.append(fused_expert)
+
+        stacked_experts = torch.stack(transposed_experts, dim=0)
+        return self.inference_model.transform_moe_mlp_1_param(stacked_experts)
@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase
+"""
+Common Attention Output Parameter Patterns
+"""
+
+
+class NormParameter(ParameterBase):
+    """
+    Simple normalization container.
+    """
+
+    params: torch.Tensor
+
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform_norm_param(self.params)
@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase
+"""
+Common QKV Parameter Patterns
+"""
+
+
+class FusedQKVParameter(ParameterBase):
+    """
+    Traditional fused QKV parameters for QKV projection. This is functionally
+    a direct copy.
+
+    src_qkv_w shape: [3 * out_features, in_features]
+    qkv_w shape: [3 * out_features, in_features]
+    """
+
+    params: torch.Tensor
+
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform_qkv_param(self.params)
+
+
+class UnfusedQKVParameter(ParameterBase):
+    """
+    QKV parameter container for unfused QKV projection.
+
+    src_param shapes: 3 x [out_features, in_features]
+    dst_param shape: [3 x out_features, in_features]
+    """
+
+    q_params: torch.Tensor
+
+    k_params: torch.Tensor
+
+    v_params: torch.Tensor
+
+    def finalize(self):
+        fused_param = torch.cat([self.q_params, self.k_params, self.v_params], dim=0)
+        return self.inference_model.transform_qkv_param(fused_param)
+
+
+def megatron_qkv_reshape(param: torch.Tensor, head_size: int, n_heads: int) -> torch.Tensor:
+    assert param.shape[0] == 3 * n_heads * head_size
+
+    all_heads = torch.chunk(param, chunks=3 * n_heads, dim=0)
+    q_heads = all_heads[::3]
+    k_heads = all_heads[1::3]
+    v_heads = all_heads[2::3]
+    return torch.cat([q_heads, k_heads, v_heads], dim=0)
+
+
+class MegatronQKVParameter(ParameterBase):
+    """
+    QKV parameter container for Megatron-style QKV projection. Megatron stores the parameter
+    as [n_heads, 3, head_size, in_features] whereas our inference system is built around
+    [3, n_heads, head_size, in_features]. This container handles the conversion.
+
+    Note: this container expects the model implementation to implement properties for
+    `head_size` and `n_heads`.
+
+    src_qkv_w shape: [3 * out_features, in_features]
+    qkv_w shape: [3 * out_features, in_features]
+    """
+
+    params: torch.Tensor
+
+    def finalize(self) -> torch.Tensor:
+        head_size = self.inference_model.head_size
+        n_heads = self.inference_model.n_heads
+
+        transposed_param = megatron_qkv_reshape(self.params, head_size, n_heads)
+        return self.inference_model.transform_qkv_param(transposed_param)
+
+
+def transform_gqa_megatron(src_param: torch.Tensor, head_size: int, n_q_heads: int, n_kv_heads: int) -> torch.Tensor:
+    assert src_param.shape[0] == (2 * n_kv_heads + n_q_heads) * head_size
+
+    head_ratio = n_q_heads // n_kv_heads
+
+    # Reshape to get the groups as the leading dimension
+    groups_leading_view = src_param.reshape(n_kv_heads, 2 + head_ratio, head_size, -1)
+    q_heads = groups_leading_view[:, :head_ratio, :, :].reshape(-1, groups_leading_view.shape[-1])
+    k_heads = groups_leading_view[:, head_ratio, :, :].reshape(-1, groups_leading_view.shape[-1])
+    v_heads = groups_leading_view[:, head_ratio + 1, :, :].reshape(-1, groups_leading_view.shape[-1])
+    # Squeeze will remove extra dimension for bias
+    return torch.cat([q_heads, k_heads, v_heads], dim=0).squeeze()
+
+
+class GQAMegatronQKVParameter(ParameterBase):
+    """
+    QKV parameter for Megatron-style QKV projection with GQA-style QKV projection. In this
+    storage format each of the groups is stored consecutively, so there will be multiple q_heads,
+    then one k head, and one v head.
+
+    Note: this container expects the model implementation to implement properties for
+    `head_size`, `n_q_heads`, and `n_kv_heads`.
+
+    src_qkv_w shape: [(2 * n_kv_heads + n_q_heads) * head_size, in_features]
+    qkv_w shape: [(2 * n_kv_heads + n_q_heads) * head_size, in_features]
+    """
+
+    params: torch.Tensor
+
+    def finalize(self) -> torch.Tensor:
+        head_size = self.inference_model.head_size
+        n_q_heads = self.inference_model.n_heads_q
+        n_kv_heads = self.inference_model.n_heads_kv
+        transposed_param = transform_gqa_megatron(self.params, head_size, n_q_heads, n_kv_heads)
+        return self.inference_model.transform_qkv_param(transposed_param)
@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ...model_implementations.parameter_base import ParameterBase
+"""
+Unembedding containers.
+"""
+
+
+class UnembedParameter(ParameterBase):
+    """
+    Unembedding parameter. This will likely be mapped to the same original weight in the model as the
+    embedding, but we have a different preferred sharding approach.
+    """
+
+    params: torch.Tensor
+    """
+    Unembedding parameter of shape [vocab_size, model_dim].
+    """
+
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform_unembed_param(self.params)