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,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