chore: import upstream snapshot with attribution

deepspeed/inference/v2/modules/interfaces/linear_base.py

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