wehub-resource-sync
71 lines
1.8 KiB
Python
71 lines
1.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import Optional
|
|
|
|
from ...inference_utils import DtypeEnum, NormTypeEnum
|
|
from ...modules.ds_module import DSModuleConfig
|
|
"""
|
|
Trying to define the space we need to support here right now:
|
|
|
|
Types of embeddings I've found so far:
|
|
1. Token embedding
|
|
2. Position embedding
|
|
3. Token type embedding
|
|
4. LN
|
|
|
|
GPTNeo: 1, 2, 3 (shared with 1)
|
|
GPTNeoX: 1
|
|
GPTJ: 1, 3
|
|
LLaMA: 1
|
|
BERT: 1, 2, 3, 4
|
|
GPT2: 1, 2, 3 (shared with 1)
|
|
|
|
Sidebar for OPT:
|
|
OPT: 1, 2
|
|
1 may not actually project to the actual hidden dimension according to the raw
|
|
code, but for the model configs we care about it does.
|
|
2 has a weird offset associated with it that the others do not.
|
|
"""
|
|
|
|
|
|
class DSEmbeddingsConfig(DSModuleConfig):
|
|
"""
|
|
Config class for DSEmbeddings.
|
|
"""
|
|
|
|
residual_dtype: DtypeEnum = DtypeEnum.fp16
|
|
"""
|
|
Data type the module should use for its output.
|
|
"""
|
|
|
|
embedding_dim: int
|
|
"""
|
|
Dimensionality of the embedding projections.
|
|
"""
|
|
|
|
positional_embedding: bool = False
|
|
"""
|
|
Whether the module should expect a positional embedding matrix. The shape of this
|
|
matrix should be of shape [max_seq_len + positional_offset, embedding_dim]
|
|
"""
|
|
|
|
positional_offset: int = 0
|
|
"""
|
|
Whether the linearized token IDs should be offset by a certain amount. For an example
|
|
of this, see the OPT model implementation.
|
|
"""
|
|
|
|
use_token_type: bool = False
|
|
"""
|
|
Whether the module should expect a token type embedding matrix.
|
|
"""
|
|
|
|
output_normalization: Optional[NormTypeEnum] = None
|
|
"""
|
|
If a the output of the embedding module should be normalized, specify here. See
|
|
``inference.inference_utils.NormTypeEnum`` for supported values.
|
|
"""
|