wehub-resource-sync
62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
# 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
|