wehub-resource-sync
74 lines
2.1 KiB
Python
74 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, 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
|