wehub-resource-sync
73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
# 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
|