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,12 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .attention_base import DSSelfAttentionRegistry, DSSelfAttentionBase
+from .embedding_base import DSEmbeddingRegistry, DSEmbeddingBase
+from .linear_base import DSLinearRegistry, DSLinearBase
+from .moe_base import DSMoERegistry, DSMoEBase
+from .post_norm_base import DSPostNormRegistry, DSPostNormBase
+from .pre_norm_base import DSPreNormRegistry, DSPreNormBase
+from .unembed_base import DSUnembedRegistry, DSUnembedBase
@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional, Tuple, Type
+
+import torch
+
+from ...ragged import RaggedBatchWrapper
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSSelfAttentionConfig
+
+
+class DSSelfAttentionBase(DSModuleBase):
+    """
+    Base mixin for all attention modules. The interface represented by this module
+    is broadly:
+
+    output = attention(query_key_value,
+                       Optional[kv_cache],
+                       Optional[attention_mask],
+                       Optional[attention_bias])
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSSelfAttentionConfig
+
+    def __init__(self, config: DSSelfAttentionConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    @property
+    def kv_block_size(self) -> int:
+        """
+        Return preferred granulatity for blocked KV-cache implementation.
+        """
+        raise NotImplementedError()
+
+    @property
+    def q_block_size(self) -> int:
+        """
+        Property to calculate blocking granularity for the query dimension.
+        This has no impact on the KV-cache structure, but will  affect the
+        number of attention atoms associated with a batch.
+        """
+        raise NotImplementedError()
+
+    def build_atoms(self, ragged_batch: RaggedBatchWrapper) -> None:
+        """
+        Build the atoms for this module. This is not a strict requirement for the class,
+        so this method is a no-op by default rather than abstract.
+        """
+        pass
+
+    def forward(self,
+                q_k_v: torch.Tensor,
+                kv_cache: torch.Tensor,
+                batch: RaggedBatchWrapper,
+                attention_mask: Optional[torch.Tensor] = None,
+                attention_bias: Optional[torch.Tensor] = None,
+                inv_freqs: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            q_k_v (torch.Tensor): Query, key, and value tensors. Expected shape is:
+                [
+                    batch,
+                    seq_len,
+                    2 * self._config.n_heads_kv + self._config.n_heads_q,
+                    self._config.head_size
+                ].
+            kv_cache (Optional[torch.Tensor]): Key and value cache tensor. Expected shape is
+                [
+                    2,
+                    batch,
+                    kv_cache_len,
+                    self._config.n_heads_kv,
+                    self._config.head_size
+                ]. If None, cache is disabled. The `kv_cache_len` dimension does not need to
+                be contiguous (it should expand stride by `max_out_tokens`).
+            batch (RaggedBatchWrapper): Ragged batch metadata.
+            attention_mask (Optional[torch.Tensor]): Attention mask tensor. If None, masking is
+                disabled. This will defer to the config in the case of conflicting information.
+                This means if the config class is implying causal attention, the mask will be ignored.
+            attention_bias (Optional[torch.Tensor]): Attention bias tensor. If None, bias is disabled.
+        """
+        raise NotImplementedError()
+
+
+class DSSelfAttentionRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSSelfAttentionBase
@@ -0,0 +1,85 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ...ragged import RaggedBatchWrapper
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSEmbeddingsConfig
+from ...inference_parameter import InferenceParameter
+
+
+class DSEmbeddingBase(DSModuleBase):
+    """
+    Base mixin for embedding modules. The interface represented by this module is:
+
+    hidden_out = embedding(input_ids) +
+                 position_embedding(position_ids) +
+                 token_type_embedding(token_type_ids)
+    with optional normalization.
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSEmbeddingsConfig
+
+    def __init__(self, config: DSEmbeddingsConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    def transform_param(self, embed_param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations on an embedding parameter. This module assumes
+        that all embedding parameters would require the same set of transformations.
+
+        Parameters:
+            embed_param (torch.Tensor): Embedding parameter. Shape is of [vocab_size, hidden_size]
+        """
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def output(self) -> torch.Tensor:
+        """
+        Pre-allocated output Tensor. This currently needs to be exposed for gather operations
+        on the output.
+
+        TODO(cmikeh2): This is not ideal. We need a better abstraction for this, such as giving
+        access to the inference comm object to the DSModule.
+        """
+        raise NotImplementedError()
+
+    def forward(self,
+                ragged_batch: RaggedBatchWrapper,
+                word_embeddings: torch.Tensor,
+                position_embeddings: Optional[torch.Tensor] = None,
+                token_type_ids: Optional[torch.Tensor] = None,
+                token_type_embeddings: Optional[torch.Tensor] = None) -> InferenceParameter:
+        """
+        Parameters:
+            ragged_batch (torch.Tensor): Ragged batch of token ids + associated metadata.
+            word_embeddings (torch.Tensor): Word embeddings.
+            position_embeddings (torch.Tensor): Position embeddings. If passed, IDs will be
+                inferred from the ragged batch itself.
+            token_type_ids (torch.Tensor): Token type ids.
+            token_type_embeddings (torch.Tensor): Token type embeddings.
+
+        Returns:
+            torch.Tensor: Hidden states. This should be the sum of the relevant
+                encodings for the model.
+        """
+        raise NotImplementedError()
+
+
+class DSEmbeddingRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSEmbeddingBase
@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSLinearConfig
+from ...inference_parameter import InferenceParameter
+
+
+class DSLinearBase(DSModuleBase):
+    """
+    Base mixin for all Linear modules. The interface represented by this module
+    is:
+
+    hidden_out = activation(hidden_in * weight + bias)
+
+    The format and dtype of the weight and bias tensors are not defined and implementations
+    may compress as necessary. Must support a bias.
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSLinearConfig
+
+    def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    @abstractmethod
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations of the parameters of this module.
+
+        Parameters:
+            param (torch.Tensor): Weight or bias tensor.
+        """
+        ...
+
+    def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Parameters:
+            hidden_states (torch.Tensor): Hidden states tensor. Expected shape is either
+                [batch, seq_len, in_channels] or [batch, in_channels].
+
+        Returns:
+            torch.Tensor: Output tensor. Tensor should have same number of dimensions as
+                input tensor.
+        """
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def output(self) -> torch.Tensor:
+        """
+        Return the padded, pre-allocated output Tensor.
+        """
+        ...
+
+
+class DSLinearRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSLinearBase
@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSMoEConfig
+from ...inference_parameter import InferenceParameter
+
+
+class DSMoEBase(DSModuleBase):
+    """
+    Base mixing for MoE modules. The interface represented by this module is:
+
+    expert_assignments = gate(hidden_states)
+    intermediate = ragged_linear(hidden_states, expert_assignments)
+    output = ragged_linear(intermediate, expert_assignments)
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSMoEConfig
+
+    def __init__(self, config: DSMoEConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    @abstractmethod
+    def transform_gate_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations of the gate parameter.
+
+        Args:
+            param (torch.Tensor): gate_w (shape: [num_experts, model_dim])
+        """
+        ...
+
+    @abstractmethod
+    def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations of the parameter. The specific component
+        being transformed should be inferred from the shape of the parameter.
+
+        Args:
+            param (torch.Tensor): One of either mlp_1_w, mlp_1_b
+        """
+        ...
+
+    @abstractmethod
+    def transform_moe_mlp_2_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations of the parameter. The specified component being
+        transformed should be inferred from the shape of the parameter. This interface is
+        separate from transform_moe_1_param because the two components may have identical
+        shapes.
+
+        Args:
+            param (torch.Tensor): One of either mlp_2_w or mlp_2_b
+        """
+        ...
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                gate_w: torch.Tensor,
+                mlp_1_w: torch.Tensor,
+                mlp_2_w: torch.Tensor,
+                mlp_1_b: Optional[torch.Tensor] = None,
+                mlp_2_b: Optional[torch.Tensor] = None) -> torch.Tensor:
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def output(self) -> torch.Tensor:
+        """
+        Returns the pre-allocated, padded output Tensor.
+        """
+        ...
+
+
+class DSMoERegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSMoEBase
@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Tuple, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..configs.norm_config import DSNormConfig
+from ..module_registry import DSModuleRegistryBase
+from ...inference_parameter import InferenceParameter
+
+
+class DSPostNormBase(DSModuleBase):
+    """
+    Base MixIn for all Post-Normalization modules. The interface represented by this
+    module is:
+
+    residual, hidden_out = norm(residual + hidden_in)
+
+    If residual and hidden_out are the same data type, then they may alias each other.
+    Furthermore, residual should be updated in-place.
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSNormConfig
+
+    def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    @abstractmethod
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Transform a gamma/beta parameter. It is assumed that both transformations are
+        the same.
+
+        Parameters:
+            param (torch.Tensor): Gamma or beta parameter.
+        """
+        ...
+
+    def forward(self,
+                residual: torch.Tensor,
+                hidden_states: torch.Tensor,
+                gamma: torch.Tensor,
+                beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            residual (torch.Tensor): Residual tensor.
+            hidden_states (torch.Tensor): Hidden states tensor.
+
+        Returns:
+            (torch.Tensor, torch.Tensor): Tuple of residual and hidden states.
+                Hidden states may alias with residual.
+        """
+        raise NotImplementedError()
+
+
+class DSPostNormRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSPostNormBase
@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Tuple, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..configs.norm_config import DSNormConfig
+from ..module_registry import DSModuleRegistryBase
+from ...inference_parameter import InferenceParameter
+
+
+class DSPreNormBase(DSModuleBase):
+    """
+    Base mixin for all Pre-Normalization modules. The interface represented by this module
+    is:
+
+    if hidden_in is not None:
+        residual_out = residual + hidden_in
+    else:
+        residual_out = residual
+
+    hidden_out = normalize(residual_out)
+    return residual_out, hidden_out
+
+    Residual should be updated in-place.
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSNormConfig
+
+    def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]):
+        super().__init__(config, implementation_config)
+
+    @abstractmethod
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Transform a gamma/beta parameter. It is assumed that both transformations are
+        the same.
+
+        Parameters:
+            param (torch.Tensor): Gamma or beta parameter.
+        """
+        ...
+
+    def forward(self,
+                residual: torch.Tensor,
+                hidden_states: Optional[torch.Tensor],
+                gamma: torch.Tensor,
+                beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            residual (torch.Tensor): Residual tensor.
+            hidden_states (torch.Tensor): Hidden states tensor.
+
+        Returns:
+            (torch.Tensor, torch.Tensor): Tuple of residual and hidden states.
+        """
+        raise NotImplementedError()
+
+
+class DSPreNormRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSPreNormBase
@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ...ragged import RaggedBatchWrapper
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSUnembedConfig
+
+
+class DSUnembedBase(DSModuleBase):
+    """
+    Base mixin for unmebedding modules. The interface represented by this module is:
+
+    if config.do_normalization
+        hidden = layer_norm(hidden)
+    logits = hidden @ projection
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSUnembedConfig
+
+    def __init__(self, config: DSUnembedConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                vocab_embedding: torch.Tensor,
+                ragged_metadata: RaggedBatchWrapper,
+                gamma: Optional[torch.Tensor] = None,
+                beta: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Forward interface. Gamma and beta are optional parameters passed depending on
+        `self.config.do_normalization`.
+
+        Args:
+            hidden_states (torch.Tensor): Hidden states of shape [tokens, model_dim]
+            vocab_embedding (torch.Tensor): Embedding matrix of shape [vocab_size, model_dim]
+            ragged_metadata (RaggedBatchWrapper): Metadata for the ragged batch.
+            gamma (Optional[torch.Tensor]): Gamma parameter for layer norm.
+            beta (Optional[torch.Tensor]): Beta parameter for layer norm.
+
+        Returns:
+            torch.Tensor: Unembedded hidden states of shape [n_seqs, model_dim]
+        """
+        raise NotImplementedError()
+
+
+class DSUnembedRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSUnembedBase