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