File size: 9,168 Bytes
204da06 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 |
# Copyright 2024 **AUTHORS_TODO**
# License: Apache-2.0
# Copyright 2022 MosaicML Examples authors
# SPDX-License-Identifier: Apache-2.0
# Copyright 2023 MosaicML Examples authors
# SPDX-License-Identifier: Apache-2.0
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.
# Copyright (c) 2023, Tri Dao.
import torch
import torch.nn as nn
from typing import Optional
from .configuration_bert import FlexBertConfig
from .normalization import get_norm_layer
from .initialization import ModuleType, init_weights
class BertAlibiEmbeddings(nn.Module):
"""Construct the embeddings for words, ignoring position.
There are no positional embeddings since we use ALiBi and token_type
embeddings.
This module is modeled after the Hugging Face BERT's
:class:`~transformers.model.bert.modeling_bert.BertEmbeddings`, but is
modified as part of Mosaic BERT's ALiBi implementation. The key change is
that position embeddings are removed. Position information instead comes
from attention biases that scale linearly with the position distance
between query and key tokens.
This module ignores the `position_ids` input to the `forward` method.
"""
def __init__(self, config):
super().__init__()
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
# ALiBi doesn't use position embeddings
if getattr(config, "token_type_embeddings", True):
self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
self.use_token_type_embeddings = True
else:
self.use_token_type_embeddings = False
self.LayerNorm = get_norm_layer(config)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
if self.use_token_type_embeddings:
self.register_buffer(
"token_type_ids", torch.zeros(config.max_position_embeddings, dtype=torch.long), persistent=False
)
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
token_type_ids: Optional[torch.LongTensor] = None,
position_ids: Optional[torch.LongTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
past_key_values_length: int = 0,
) -> torch.Tensor:
if (input_ids is not None) == (inputs_embeds is not None):
raise ValueError("Must specify either input_ids or input_embeds!")
if input_ids is not None:
input_shape = input_ids.size()
else:
assert inputs_embeds is not None # just for type checking
input_shape = inputs_embeds.size()[:-1]
seq_length = input_shape[1]
if position_ids is None:
# great! ALiBi
pass
# Setting the token_type_ids to the registered buffer in constructor
# where it is all zeros, which usually occurs when it's auto-generated;
# registered buffer helps users when tracing the model without passing
# token_type_ids, solves issue #5664
if self.use_token_type_embeddings and token_type_ids is None:
if hasattr(self, "token_type_ids"):
buffered_token_type_ids = self.token_type_ids[:, :seq_length]
buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
token_type_ids = buffered_token_type_ids_expanded
else:
token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=input_ids.device)
if inputs_embeds is None:
inputs_embeds = self.word_embeddings(input_ids)
if self.use_token_type_embeddings:
token_type_embeddings = self.token_type_embeddings(token_type_ids)
embeddings = inputs_embeds + token_type_embeddings
else:
embeddings = inputs_embeds
# no position embeddings! ALiBi
embeddings = self.LayerNorm(embeddings)
embeddings = self.dropout(embeddings)
return embeddings
class FlexBertEmbeddingsBase(nn.Module):
"""A FlexBERT embeddings base class for type hints."""
def __init__(self, config: FlexBertConfig):
super().__init__()
self.config = config
def _init_weights(self, reset_params: bool = False):
raise NotImplementedError("This is a base class and should not be used directly.")
def reset_parameters(self):
self._init_weights(reset_params=True)
def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
raise NotImplementedError("This is a base class and should not be used directly.")
class FlexBertAbsoluteEmbeddings(FlexBertEmbeddingsBase):
"""Construct the embeddings with absolute positional embeddings."""
def __init__(self, config: FlexBertConfig):
super().__init__(config)
self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
self.norm = get_norm_layer(config) if config.embed_norm else nn.Identity()
self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()
self.register_buffer(
"position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
)
def _init_weights(self, reset_params: bool = False):
init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)
init_weights(self.config, self.position_embeddings, type_of_module=ModuleType.emb)
if reset_params:
if self.config.embed_norm:
self.norm.reset_parameters() # type: ignore
def forward(
self,
input_ids: torch.LongTensor,
position_ids: Optional[torch.LongTensor] = None,
) -> torch.Tensor:
if position_ids is None:
position_ids = self.position_ids[:, 0 : input_ids.shape[1]]
embeddings = self.tok_embeddings(input_ids)
position_embeddings = self.position_embeddings(position_ids)
embeddings = self.norm(embeddings + position_embeddings)
return self.drop(embeddings)
class FlexBertCompiledSansPositionEmbeddings(FlexBertEmbeddingsBase):
"""Construct the embeddings from token embeddings without any positional embeddings."""
def __init__(self, config: FlexBertConfig):
super().__init__(config)
self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
self.norm = get_norm_layer(config, compiled_norm=config.compile_model) if config.embed_norm else nn.Identity()
self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()
def _init_weights(self, reset_params: bool = False):
init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)
if reset_params:
if self.config.embed_norm:
self.norm.reset_parameters() # type: ignore
@torch.compile(dynamic=True)
def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
return self.drop(self.norm(self.tok_embeddings(input_ids)))
class FlexBertSansPositionEmbeddings(FlexBertEmbeddingsBase):
"""Construct the embeddings from token embeddings without any positional embeddings."""
def __init__(self, config: FlexBertConfig):
super().__init__(config)
self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
self.norm = get_norm_layer(config) if config.embed_norm else nn.Identity()
self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()
def _init_weights(self, reset_params: bool = False):
init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)
if reset_params:
if self.config.embed_norm:
self.norm.reset_parameters() # type: ignore
def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
return self.drop(self.norm(self.tok_embeddings(input_ids)))
EBB2CLS = {
"absolute_pos": FlexBertAbsoluteEmbeddings,
"sans_pos": FlexBertSansPositionEmbeddings,
}
def get_embedding_layer(config: FlexBertConfig) -> FlexBertEmbeddingsBase:
try:
if config.compile_model and config.embedding_layer == "sans_pos":
return FlexBertCompiledSansPositionEmbeddings(config)
elif config.compile_model:
raise ValueError(f"{config.compile_model=} only supports sans_pos embeddings.")
return EBB2CLS[config.embedding_layer](config)
except KeyError:
raise ValueError(f"Invalid embeddings layer type: {config.embedding_layer=}, must be one of {EBB2CLS.keys()}.")
|