wehub-resource-sync
86 lines
3.0 KiB
Python
86 lines
3.0 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 ...ragged import RaggedBatchWrapper
|
|
from ..ds_module import DSModuleBase
|
|
from ..module_registry import DSModuleRegistryBase
|
|
from ..configs import DSEmbeddingsConfig
|
|
from ...inference_parameter import InferenceParameter
|
|
|
|
|
|
class DSEmbeddingBase(DSModuleBase):
|
|
"""
|
|
Base mixin for embedding modules. The interface represented by this module is:
|
|
|
|
hidden_out = embedding(input_ids) +
|
|
position_embedding(position_ids) +
|
|
token_type_embedding(token_type_ids)
|
|
with optional normalization.
|
|
"""
|
|
|
|
@staticmethod
|
|
def config_class() -> Type[DeepSpeedConfigModel]:
|
|
return DSEmbeddingsConfig
|
|
|
|
def __init__(self, config: DSEmbeddingsConfig, implementation_config: Dict[str, Any]) -> None:
|
|
super().__init__(config, implementation_config)
|
|
|
|
def transform_param(self, embed_param: torch.Tensor) -> InferenceParameter:
|
|
"""
|
|
Perform any necessary transformations on an embedding parameter. This module assumes
|
|
that all embedding parameters would require the same set of transformations.
|
|
|
|
Parameters:
|
|
embed_param (torch.Tensor): Embedding parameter. Shape is of [vocab_size, hidden_size]
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
@property
|
|
@abstractmethod
|
|
def output(self) -> torch.Tensor:
|
|
"""
|
|
Pre-allocated output Tensor. This currently needs to be exposed for gather operations
|
|
on the output.
|
|
|
|
TODO(cmikeh2): This is not ideal. We need a better abstraction for this, such as giving
|
|
access to the inference comm object to the DSModule.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def forward(self,
|
|
ragged_batch: RaggedBatchWrapper,
|
|
word_embeddings: torch.Tensor,
|
|
position_embeddings: Optional[torch.Tensor] = None,
|
|
token_type_ids: Optional[torch.Tensor] = None,
|
|
token_type_embeddings: Optional[torch.Tensor] = None) -> InferenceParameter:
|
|
"""
|
|
Parameters:
|
|
ragged_batch (torch.Tensor): Ragged batch of token ids + associated metadata.
|
|
word_embeddings (torch.Tensor): Word embeddings.
|
|
position_embeddings (torch.Tensor): Position embeddings. If passed, IDs will be
|
|
inferred from the ragged batch itself.
|
|
token_type_ids (torch.Tensor): Token type ids.
|
|
token_type_embeddings (torch.Tensor): Token type embeddings.
|
|
|
|
Returns:
|
|
torch.Tensor: Hidden states. This should be the sum of the relevant
|
|
encodings for the model.
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
|
|
class DSEmbeddingRegistry(DSModuleRegistryBase):
|
|
registry: Dict = {}
|
|
|
|
@staticmethod
|
|
def associated_class() -> Type[DSModuleBase]:
|
|
return DSEmbeddingBase
|