Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2023 The Mega Authors and The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch MEGA model."""16 17import math18from typing import Optional, Union19 20import torch21import torch.nn.functional as F22from torch import nn23from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss24 25from ....activations import ACT2FN26from ....cache_utils import Cache27from ....modeling_outputs import (28 BaseModelOutputWithPoolingAndCrossAttentions,29 CausalLMOutputWithCrossAttentions,30 MaskedLMOutput,31 MultipleChoiceModelOutput,32 QuestionAnsweringModelOutput,33 SequenceClassifierOutput,34 TokenClassifierOutput,35)36from ....modeling_utils import PreTrainedModel37from ....utils import (38 add_code_sample_docstrings,39 add_start_docstrings,40 add_start_docstrings_to_model_forward,41 logging,42 replace_return_docstrings,43)44from ....utils.deprecation import deprecate_kwarg45from .configuration_mega import MegaConfig46 47 48logger = logging.get_logger(__name__)49 50_CHECKPOINT_FOR_DOC = "mnaylor/mega-base-wikitext"51_CONFIG_FOR_DOC = "MegaConfig"52 53 54class MegaEmbeddings(nn.Module):55 """56 Mega's basic implementation does not incorporate token type embeddings, so this is a stripped-down version of57 RoBERTa's embeddings which optionally includes token types58 """59 60 def __init__(self, config: MegaConfig):61 super().__init__()62 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)63 self.use_token_types = config.add_token_type_embeddings64 if self.use_token_types:65 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)66 # registering a buffer here allows model tracing when not passing optional token type IDs67 # more info at transformers issue #566468 self.register_buffer(69 "token_type_ids", torch.zeros(config.max_positions, dtype=torch.long).expand((1, -1)), persistent=False70 )71 72 self.padding_idx = config.pad_token_id73 74 def forward(self, input_ids=None, token_type_ids=None, inputs_embeds=None):75 if (input_ids is None) and (inputs_embeds is None):76 raise ValueError("Must provide one of input_ids or inputs_embeds")77 elif input_ids is not None:78 input_shape = input_ids.size()79 device = input_ids.device80 81 # get the word embeddings if only IDs are provided82 inputs_embeds = self.word_embeddings(input_ids)83 else:84 input_shape = inputs_embeds.size()[:-1]85 device = inputs_embeds.device86 87 # the original Mega implementation did not include token type embeddings, so we add88 # an option to use them if desired; if embeddings are present and token type IDs are89 # not provided, we will use a registered buffer (which helps with tracing)90 if self.use_token_types:91 if token_type_ids is None:92 if hasattr(self, "token_type_ids"):93 buffered_token_type_ids = self.token_type_ids[:, : input_shape[1]]94 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], input_shape[1])95 token_type_ids = buffered_token_type_ids_expanded96 else:97 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)98 99 # access token type embeddings100 token_type_embeddings = self.token_type_embeddings(token_type_ids)101 # add the token type embeddings to the word embeddings102 embeddings = inputs_embeds + token_type_embeddings103 else:104 embeddings = inputs_embeds105 return embeddings106 107 108class MegaSimpleRelativePositionalBias(nn.Module):109 """110 Simple relative positional embeddings copied from the Mega repo; renamed variables for better readability111 """112 113 def __init__(self, config: MegaConfig):114 super().__init__()115 self.config = config116 self.max_positions = self.config.max_positions if self.config.chunk_size < 0 else self.config.chunk_size117 self.rel_pos_bias = nn.Parameter(torch.Tensor(2 * config.max_positions - 1))118 119 def forward(self, seq_len):120 if seq_len > self.max_positions:121 raise ValueError(f"Sequence length {seq_len} going beyond max length {self.max_positions}")122 123 # seq_len * 2 - 1124 bias = self.rel_pos_bias[(self.max_positions - seq_len) : (self.max_positions + seq_len - 1)]125 # seq_len * 3 - 1126 tile = F.pad(bias, (0, seq_len))127 # (seq_len * 3 - 1) * seq_len128 tile = torch.tile(tile, (seq_len,))129 tile = tile[:-seq_len]130 # seq_len x (3 * seq_len - 2)131 tile = tile.view(seq_len, 3 * seq_len - 2)132 start = (2 * seq_len - 1) // 2133 end = tile.size(1) - start134 tile = tile[:, start:end]135 return tile136 137 138class MegaRotaryRelativePositionalBias(nn.Module):139 """140 Rotary relative bias for positional information; similar in concept to RoPE (i.e. RoFormer) but taken from the Mega141 repo due to differences in implementation.142 143 When initialized, produces a positional bias which ranges from position 0 to config.max_positions, but can144 extrapolate to longer sequences. Can be indexed according to input position IDs145 """146 147 def __init__(self, config: MegaConfig):148 super().__init__()149 if config.hidden_size % 2 != 0:150 raise RuntimeError("Rotary positional bias requires `hidden_size` to be a multiple of 2")151 self.config = config152 self.embed_dim = config.shared_representation_size153 self.max_positions = self.config.max_positions if self.config.chunk_size < 0 else self.config.chunk_size154 self.sine, self.cosine = MegaRotaryRelativePositionalBias.get_sinusoid_embeddings(155 config.max_positions, self.embed_dim156 )157 # alpha and beta parameters for the rotary bias; beta renamed to b_param to avoid clashes with tf/flax weight handling158 # in loading pretrained weights159 self.alpha = nn.Parameter(torch.Tensor(1, self.embed_dim))160 self.b_param = nn.Parameter(torch.Tensor(1, self.embed_dim))161 self.register_buffer("_float_tensor", torch.FloatTensor([0.0]))162 163 @staticmethod164 def get_sinusoid_embeddings(max_positions: int, embedding_dim: int):165 half_dim = embedding_dim // 2166 emb = math.log(10000) / half_dim167 emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb)168 emb = torch.arange(max_positions, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)169 return torch.sin(emb), torch.cos(emb)170 171 def rotary(self, input):172 seq_len, embed_dim = input.size()173 chunk_1, chunk_2 = torch.chunk(input, 2, dim=-1)174 if self.sine is None or seq_len > self.sine.size(0):175 self.sine, self.cosine = MegaRotaryRelativePositionalBias.get_sinusoid_embeddings(seq_len, embed_dim)176 self.max_positions = seq_len177 self.sine = self.sine.to(self._float_tensor)178 self.cosine = self.cosine.to(self._float_tensor)179 180 sin = self.sine[:seq_len]181 cos = self.cosine[:seq_len]182 return torch.cat([chunk_1 * cos - chunk_2 * sin, chunk_2 * cos + chunk_1 * sin], dim=1)183 184 def forward(self, seq_len):185 rotary_alpha = self.rotary(self.alpha.expand(seq_len, self.embed_dim))186 rotary_beta = self.rotary(self.b_param.expand(seq_len, self.embed_dim))187 bias = torch.einsum("mk,nk->mn", rotary_alpha, rotary_beta)188 return bias189 190 191class MegaDropout(nn.Module):192 """193 A unified class for standard dropout functionality and featurewise dropout.194 195 The original fairseq Mega repo used 2 classes for these, which included some unnecessary handling of training logic196 and an unused `inplace` option. The original implementation used torch.nn.functional instead of submodules, which197 is retained here as well.198 """199 200 def __init__(self, dropout_probability, is_featurewise=False):201 super().__init__()202 self.dropout_probability = dropout_probability203 self.is_featurewise = is_featurewise204 205 def forward(self, input, batch_first: bool = False):206 if self.is_featurewise:207 if batch_first:208 # (batch_size X sequence_length X feature_dimension)209 # -> (batch_size X feature_dimension X sequence_length)210 # -> (batch_size X sequence_length X feature_dimension)211 return F.dropout2d(212 input.transpose(-1, -2), p=self.dropout_probability, training=self.training213 ).transpose(-1, -2)214 else:215 if input.dim() != 3:216 raise ValueError(217 "Feature dropout inputs must be exactly 3-dimensional if inputs are ordered [sequence length, batch size, hidden dimension]"218 )219 # (sequence_length X batch_size X feature_dimension)220 # -> (batch_size X feature_dimension X sequence_length)221 # -> (sequence_length X batch_size X feature_dimension)222 return F.dropout2d(input.permute(1, 2, 0), p=self.dropout_probability, training=self.training).permute(223 2, 0, 1224 )225 else:226 return F.dropout(input, p=self.dropout_probability, training=self.training)227 228 229class MegaRMSNorm(nn.Module):230 """231 RMSNorm used in Mega implementation. Differs from T5's RMSNorm by applying the weight prior to taking the square232 root (as opposed to after in T5)233 """234 235 def __init__(self, number_features, eps=1e-6, affine=True):236 super().__init__()237 self.num_features = number_features238 self.eps = eps239 self.affine = affine240 if affine:241 self.weight = nn.Parameter(torch.Tensor(self.num_features))242 else:243 self.register_parameter("weight", None)244 245 def forward(self, input):246 mean_square = torch.mean(torch.square(input), dim=-1, keepdim=True)247 if self.weight is not None:248 input = input * self.weight249 250 input * torch.rsqrt(mean_square + self.eps)251 return input252 253 def extra_repr(self):254 return f"{self.num_features}, eps={self.eps}, affine={self.affine}"255 256 257class MegaScaleNorm(nn.Module):258 """259 Scale normalization introduced in MEGA which is similar to RMSNorm, but uses a single parameter for scalar260 multiplication instead of a vector, and applies over a specified dimension261 """262 263 def __init__(self, dim, eps=1e-6, affine=True):264 super().__init__()265 self.dim = dim266 self.eps = eps267 self.affine = affine268 if affine:269 self.scalar = nn.Parameter(torch.Tensor(1))270 else:271 self.register_parameter("scalar", None)272 273 def forward(self, input):274 mean_square = torch.mean(torch.square(input), dim=self.dim, keepdim=True)275 if self.scalar is not None:276 input = self.scalar * input277 278 output = input * torch.rsqrt(mean_square + self.eps)279 return output280 281 282class MegaSequenceNorm(nn.Module):283 """284 A wrapper class for various layer normalization options used in Mega. Used to handle differences in expectations on285 input axis locations for different normalization methods.286 """287 288 def __init__(self, norm_type, embedding_dim, eps=1e-5, affine=True, export=False):289 super().__init__()290 if norm_type == "layernorm":291 self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine=affine)292 elif norm_type == "scalenorm":293 self.norm = MegaScaleNorm(dim=-1, eps=eps, affine=affine)294 elif norm_type == "rmsnorm":295 self.norm = MegaRMSNorm(embedding_dim, eps=eps, affine=affine)296 elif norm_type == "batchnorm":297 self.norm = nn.BatchNorm1d(embedding_dim, eps=eps, affine=affine)298 elif norm_type == "syncbatchnorm":299 self.norm = nn.SyncBatchNorm(embedding_dim, eps=eps, affine=affine)300 else:301 raise ValueError(f"Unknown norm type: {norm_type}")302 303 def forward(self, input):304 if isinstance(self.norm, nn.modules.batchnorm._BatchNorm):305 if input.dim() != 3:306 raise ValueError("BatchNorm inputs must be exactly 3-dimensional")307 input = input.permute(1, 2, 0)308 input = self.norm(input)309 return input.permute(2, 0, 1)310 else:311 return self.norm(input)312 313 314class MegaMultiDimensionDampedEma(nn.Module):315 """316 Mega's Exponential Moving Average layer, largely left unmodified from the original repo with the exception of317 variable names and moving away from the stateful representation of incremental decoding state. See318 "https://huggingface.co/papers/2209.10655" for more details.319 """320 321 def __init__(self, config: MegaConfig):322 super().__init__()323 324 self.config = config325 326 self.embed_dim = config.hidden_size327 self.ndim = config.ema_projection_size328 self.bidirectional = config.bidirectional329 self.truncation = config.truncation330 self.scale = math.sqrt(1.0 / self.ndim)331 332 kernel_dim = 2 * config.hidden_size if self.bidirectional else config.hidden_size333 # renamed delta (damping_factor) and alpha (decay_factor) to be more descriptive of what the parameters are doing334 self.damping_factor = nn.Parameter(torch.Tensor(kernel_dim, self.ndim, 1))335 self.decay_factor = nn.Parameter(torch.Tensor(kernel_dim, self.ndim, 1))336 # renamed gamma (kernel_projection_matrix) and beta (ema_expansion_matrix) respectively to avoid HF renaming337 # things and align with the paper's description of these params' behavior338 self.ema_expansion_matrix = nn.Parameter(torch.Tensor(kernel_dim, self.ndim, 1))339 self.kernel_projection_matrix = nn.Parameter(torch.Tensor(kernel_dim, self.ndim))340 # renamed omega to residual_weight to describe what it's doing341 self.residual_weight = nn.Parameter(torch.Tensor(config.hidden_size))342 self._kernel = None343 self._coeffs = None344 345 def _compute_ema_coefficients(self):346 self._coeffs = None347 # convert the alpha and delta parameters (kernel_dim x EMA projection size x 1) to [0, 1] with sigmoid348 damping_factor = torch.sigmoid(self.damping_factor)349 decay_factor = torch.sigmoid(self.decay_factor)350 previous_timestep_weight = 1.0 - damping_factor * decay_factor351 return damping_factor, previous_timestep_weight352 353 def _compute_efficient_ema_kernel(self, length: int):354 # computes the kernel used for efficient damped EMA applied via FFT convolution355 self._kernel = None356 # p and q have shape (kernel_dim x ema_projection_size x 1)357 damping_factor, previous_timestep_weight = self._compute_ema_coefficients()358 # extend the kernel to (kernel_dim X ema_projection_size X sequence_length) and359 # multiply q by sequential ints up to the sequence length360 vander = torch.arange(length).to(damping_factor).view(1, 1, length) * torch.log(previous_timestep_weight)361 kernel = (damping_factor * self.ema_expansion_matrix) * torch.exp(vander)362 # (kernel_dim X ema_projection_size X sequence_length) -> (kernel_dim, sequence_length)363 return torch.einsum("dnl,dn->dl", kernel, self.kernel_projection_matrix * self.scale)364 365 def get_ema_coefficients(self):366 if self.training:367 return self._compute_ema_coefficients()368 else:369 if self._coeffs is None:370 self._coeffs = self._compute_ema_coefficients()371 return self._coeffs372 373 def get_ema_kernel(self, length: int):374 kernel_size = length if self.truncation is None else min(self.truncation, length)375 if self.training:376 return self._compute_efficient_ema_kernel(kernel_size)377 else:378 if self._kernel is None or self._kernel.size(-1) < kernel_size:379 self._kernel = self._compute_efficient_ema_kernel(kernel_size)380 return self._kernel[..., :kernel_size]381 382 def fft_convolution(self, inputs, kernel, length):383 # this is a wrapper for repeated use of EMA calculation via FFT (fast Fourier transform) convolution384 inputs_fft = torch.fft.rfft(inputs.float(), n=2 * length)385 kernel_fft = torch.fft.rfft(kernel.float(), n=2 * length)386 convolved_sequence = torch.fft.irfft(inputs_fft * kernel_fft, n=2 * length)387 return convolved_sequence388 389 def ema_step(self, inputs, length, past_state=None):390 if length == 1:391 return self.one_ema_step(inputs, past_state=past_state)392 393 # (kernel_dim X ema_projection_size X 1)394 damping_factor, previous_timestep_weight = self.get_ema_coefficients()395 # (kernel_dim X ema_projection_size X 1+sequence_length)396 vander = torch.arange(length + 1).to(damping_factor).view(1, 1, length + 1) * torch.log(397 previous_timestep_weight398 )399 vander = torch.exp(vander)400 if past_state is not None:401 # (kernel_dim X ema_projection_size X sequence_length) * (kernel_dim X ema_projection_size X 1)402 # -> (kernel_dim X ema_projection_size X sequence_length)403 past_ema_proj = vander[:, :, 1:] * (self.kernel_projection_matrix * self.scale).unsqueeze(-1)404 # past_state will be (batch_size, kernel_dim, ema_projection_size)405 past_ema_state = torch.einsum("bdn,dnl->bdl", past_state, past_ema_proj)406 # (kernel_dim X ema_projection_size) * (batch_size X kernel_dim X ema_projection_size)407 # -> (batch_size X kernel_dim X ema_projection_size)408 past_vandermonde = vander[:, :, -1] * past_state409 else:410 past_ema_state = None411 past_vandermonde = None412 413 # (kernel_dim X ema_projection_size X sequence_length)414 vander = vander[:, :, :-1]415 kernel = (damping_factor * self.ema_expansion_matrix) * vander416 kernel_proj = torch.einsum("dnl,dn->dl", kernel, self.kernel_projection_matrix * self.scale)417 418 ema_output = self.fft_convolution(inputs, kernel_proj, length=length)[..., 0:length]419 ema_output = ema_output.type_as(inputs)420 if past_ema_state is not None:421 ema_output = ema_output + past_ema_state422 423 updated_hidden_state = torch.einsum("bdl,dnl->bdn", inputs, torch.flip(kernel, dims=[2]))424 if past_vandermonde is not None:425 updated_hidden_state = updated_hidden_state + past_vandermonde426 # return a tuple:427 # (sequence_length, batch_size, kernel_dim)428 # (batch_size, kernel_dim, ema_projection_size)429 return ema_output.permute(2, 0, 1), updated_hidden_state430 431 def one_ema_step(self, inputs, past_state=None):432 damping_factor, previous_timestep_weight = self.get_ema_coefficients()433 # (kernel_dim X ema_projection_size) x (batch_size X kernel_dim X 1)434 # -> (batch_size X kernel_dim X ema_projection_size)435 updated_state = (damping_factor * self.ema_expansion_matrix).squeeze(-1) * inputs436 if past_state is not None:437 updated_state = updated_state + previous_timestep_weight.squeeze(-1) * past_state438 # (batch_size X kernel_dim)439 out = torch.einsum("bdn,dn->bd", updated_state, self.kernel_projection_matrix * self.scale)440 # (1 X batch_size X kernel_dim), (batch_size X kernel_dim X ema_projection_size)441 return out.unsqueeze(0), updated_state442 443 def forward(444 self,445 inputs,446 attention_mask: Optional[torch.Tensor] = None,447 prev_state: Optional[torch.Tensor] = None,448 use_cache: bool = False,449 ) -> torch.Tensor:450 """451 Mega's exponential moving average (EMA) sub-layer applied prior to single-headed (traditional) self-attention452 453 Args:454 inputs (`torch.Tensor` of shape `(sequence_length, batch_size, hidden_size)`):455 Hidden state / embedding input to update via EMA based on FFT convolution456 attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):457 Indicates which inputs are to be ignored (mostly due to padding), where elements are either 1 for *not458 masked* or 0 for *masked*459 prev_state (`torch.Tensor` of shape `(batch_size, config.ndim)`, *optional*):460 The hidden state returned from the previous timestep during incremental decoding.461 use_cache (`bool`, default `False`):462 Whether to perform incremental decoding; uses `prev_state` as the prior timestep, and returns the463 updated EMA hidden state for use in the next step464 465 Returns:466 `tuple(torch.FloatTensor)` containing various elements depending on configuration ([`MegaConfig`]) and467 inputs:468 - **hidden_states** (`torch.FloatTensor` of shape `(sequence_length, batch_size, hidden_size)`) -- Hidden469 states updated by EMA, with same shapes as inputs470 - **updated_state** (*optional*, returned when `use_cache=True`) `torch.FloatTensor of shape `(batch_size,471 config.ndim)` -- The incremental EMA state for use in the next step of incremental decoding472 """473 474 seq_len, bsz, embed_dim = inputs.size()475 if embed_dim != self.embed_dim:476 raise ValueError(477 f"Unexpected embedding dimension received: input is {embed_dim}, model expects {self.embed_dim}"478 )479 480 # sequence_length X batch_size X hidden_size481 residual = inputs * self.residual_weight482 483 # (sequence_length x batch_size x hidden_size) -> (batch_size x hidden_size x sequence_length)484 inputs = inputs.permute(1, 2, 0)485 # mask the input: output is a tensor with 0 in the masked positions486 if attention_mask is not None:487 inputs = inputs * (attention_mask.unsqueeze(1).type_as(inputs))488 489 if self.bidirectional and use_cache:490 raise RuntimeError("Bidirectional EMA does not support incremental state")491 492 if use_cache:493 out, updated_state = self.ema_step(inputs, seq_len, past_state=prev_state)494 495 # (batch_size X hidden_size) -> (1 x batch_size x hidden_size)496 out = F.silu(out + residual)497 498 # if incremental decoding, return the new state along with the output499 return out, updated_state500 else:501 # (hidden_size x sequence_length)502 kernel = self.get_ema_kernel(seq_len)503 fft_len = seq_len504 s_index = 0505 kernel_size = kernel.size(1)506 if self.bidirectional:507 # split the kernel for each direction of EMA508 k1, k2 = torch.split(kernel, [self.embed_dim, self.embed_dim], dim=0)509 # (hidden_size X 2*sequence_length - 1)510 kernel = F.pad(k1, (kernel_size - 1, 0)) + F.pad(k2.flip(-1), (0, kernel_size - 1))511 inputs = F.pad(inputs, (kernel_size - 1, 0))512 fft_len = fft_len + kernel_size - 1513 s_index = 2 * kernel_size - 2514 515 ema_output = self.fft_convolution(inputs, kernel, length=fft_len)[..., s_index : s_index + seq_len]516 ema_output = ema_output.type_as(inputs)517 # (batch_size X hidden_size X sequence_length) -> (sequence_length X batch_size X hidden_size)518 gated_ema_output = F.silu(ema_output.permute(2, 0, 1) + residual)519 520 return gated_ema_output, None521 522 523class MegaGatedCrossAttention(nn.Module):524 """525 Gated Structured State Attention for use in encoder-decoder model. See Mega paper for more details. Only526 modifications from original implementation are variable names, removing the unnecessary `before_attn_fn` and527 `static_kv` arguments, and the stateful representation of incremental decoder state.528 """529 530 def __init__(self, config: MegaConfig):531 super().__init__()532 533 self.config = config534 self.activation = ACT2FN[self.config.activation]535 self.attention_activation = self.config.attention_activation536 self.scaling = self.config.shared_representation_size**-0.5 if self.attention_activation == "softmax" else None537 538 self.dropout = MegaDropout(self.config.dropout_prob, is_featurewise=self.config.use_feature_dropout)539 self.hidden_dropout = MegaDropout(540 self.config.hidden_dropout_prob, is_featurewise=self.config.use_feature_dropout541 )542 # Attention dropout is standard dropout543 self.attention_dropout = MegaDropout(self.config.attention_probs_dropout_prob, is_featurewise=False)544 545 self.prenorm = self.config.normalize_before_mega546 self.norm = MegaSequenceNorm(547 self.config.normalization_type, self.config.hidden_size, affine=self.config.norm_affine548 )549 550 self.k_proj = nn.Linear(self.config.hidden_size, self.config.shared_representation_size)551 self.v_proj = nn.Linear(self.config.hidden_size, self.config.hidden_size)552 self.q_proj = nn.Linear(553 self.config.hidden_size, 2 * self.config.hidden_size + self.config.shared_representation_size554 )555 self.h_proj = nn.Linear(self.config.hidden_size, self.config.hidden_size)556 557 if self.config.relative_positional_bias == "simple":558 self.rel_pos_bias = MegaSimpleRelativePositionalBias(config)559 elif self.config.relative_positional_bias == "rotary":560 self.rel_pos_bias = MegaRotaryRelativePositionalBias(config)561 else:562 raise ValueError(f"unknown relative position bias: {self.config.relative_positional_bias}")563 564 self.softmax = nn.Softmax(dim=-1)565 566 def element_attention(self, query, key, key_padding_mask, pidx):567 bsz, src_len, _ = key.size()568 tgt_len = query.size(1) if pidx is None else pidx + 1569 if key_padding_mask is not None:570 # (batch_size X source_sequence_length) --> (batch_size X 1 X 1)571 lengths = key_padding_mask.sum(dim=-1).view(bsz, 1, 1)572 else:573 lengths = src_len574 575 # (target_sequence_length X source_sequence_length)576 bias = self.rel_pos_bias(max(tgt_len, src_len))[:, :src_len]577 if pidx is not None:578 if query.size(1) != 1:579 raise ValueError("Position offset provided with queries longer than 1 token")580 # source_sequence_length581 bias = bias[pidx]582 else:583 # (target_sequence_length X source_sequence_length)584 bias = bias[:tgt_len]585 586 # (batch_size X target_sequence_length X source_sequence_length)587 qk = torch.bmm(query, key.transpose(1, 2)) / lengths + bias588 589 attn_weights = ACT2FN[self.attention_activation](qk).type_as(qk)590 591 if key_padding_mask is not None:592 attn_weights = attn_weights * key_padding_mask.unsqueeze(1)593 594 return attn_weights595 596 def softmax_attention(self, query, key, key_padding_mask, pidx):597 bsz, src_len, _ = key.size()598 tgt_len = query.size(1) if pidx is None else pidx + 1599 600 # (target_sequence_length X source_sequence_length)601 bias = self.rel_pos_bias(max(tgt_len, src_len))[:, :src_len]602 if pidx is not None:603 if query.size(1) != 1:604 raise ValueError("Position offset provided with queries longer than 1 token")605 # source_sequence_length606 bias = bias[pidx]607 else:608 # (target_sequence_length X source_sequence_length)609 bias = bias[:tgt_len]610 611 # scaled attention612 query = query * self.scaling613 # (batch_size X target_sequence_length X source_sequence_length)614 qk = torch.bmm(query, key.transpose(1, 2)) + bias615 616 if key_padding_mask is not None:617 qk = qk.masked_fill((1 - key_padding_mask).unsqueeze(1).to(torch.bool), float("-inf"))618 619 attn_weights = self.softmax(qk).type_as(qk)620 return attn_weights621 622 def forward(623 self,624 query,625 key: Optional[torch.Tensor],626 value: Optional[torch.Tensor],627 key_padding_mask: Optional[torch.Tensor] = None,628 past_key_values: Optional[Cache] = None,629 output_attentions: bool = False,630 use_cache: bool = False,631 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:632 """633 Gated cross-attention used in Mega634 635 Args:636 query (`torch.Tensor` of shape `(target_sequence_length, batch_size, hidden_size)`):637 The self (or target) sequence input used as query inputs for cross-attention638 key (`torch.Tensor` of shape `(source_sequence_length, batch_size, hidden_size)`):639 The cross (or source) sequence input with shape used as keys in cross-attention640 value (`torch.Tensor` of shape `(source_sequence_length, batch_size, hidden_size)`):641 The cross (or source) sequence input with shape used as values in cross-attention642 key_padding_mask (`torch.LongTensor` of shape `(batch_size, source_sequence_length)`, *optional*):643 Padding mask corresponding to the source sequence, where entries are 1 for *not masked* and 0 for644 *masked* tokens645 past_key_values (`tuple(torch.FloatTensor)`, *optional*):646 If provided, the hidden state returned from the previous timestep during incremental decoding; expects647 that prior cross-attention keys and values will be the last two items in the tuple648 output_attentions (`bool`, defaults to `False`):649 Whether or not to return the cross-attention weights.650 use_cache (`bool`, defaults to `False`):651 Whether to perform incremental decoding; uses `prev_state` as the prior timestep, and returns the652 updated EMA hidden state for use in the next step653 654 Returns:655 `tuple(torch.FloatTensor)` containing various elements depending on configuration ([`MegaConfig`]) and656 inputs:657 - **hidden_states** (`torch.FloatTensor` of shape `(target_sequence_length, batch_size, hidden_size)`) --658 Hidden states from target sequence updated by gated cross-attention659 - **attn_weights** (*optional*, returned when `output_attentions=True`) `torch.FloatTensor` of shape660 `(batch_size, source_sequence_length, target_sequence_length)` -- The pairwise cross-attention weights661 corresponding to each token in the source and target sequences662 - **cross_key** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,663 source_sequence_length, config.shared_representation_size)` -- The cross-attention key state for use in664 the next step of incremental decoding665 - **cross_value** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,666 source_sequence_length, config.hidden_size)` -- The cross-attention value state for use in the next step667 of incremental decoding668 """669 670 seq_len, bsz, embed_dim = query.size()671 if embed_dim != self.config.hidden_size:672 raise ValueError(673 f"Unexpected embedding dimension received: input is {embed_dim} but expected {self.config.hidden_size}"674 )675 676 if past_key_values is not None:677 # make sure the inputs only have a sequence length of 1 if we're doing incremental decoding678 if seq_len != 1:679 raise ValueError(f"Incremental decoding requested with self-sequence length > 1: {seq_len}")680 # expect past_key_values to have (self_key, self_value, self_ema, cross_key, cross_value)681 prev_cross_key, prev_cross_value = past_key_values[-2:]682 key = value = None683 684 # use the self-attention cache to get the position id of the current step685 prev_self_key = past_key_values[0]686 num_incremental_steps = prev_self_key.size(1) + 1687 else:688 prev_cross_key = prev_cross_value = None689 # we still need the position id if we're doing incremental decoding (past_key_values will be None for the first step)690 num_incremental_steps = 0 if use_cache and (seq_len == 1) else None691 692 full_query = query693 if self.prenorm:694 full_query = self.norm(full_query)695 696 # (target_sequence_length X batch_size X 2*hidden_size + shared_representation_size)697 query_projected = self.q_proj(full_query)698 # split the query projections into separate components699 # - residual_weight is passed through sigmoid and sent through elementwise multiplication to the gated/weighted targets prior to being added to the query directly700 # - target_gate is a silu-gated tensor that is multiplied by the attention-weighted target below prior to residual connection701 # - attention_query is the part that is passed to the attention function702 residual_weight, target_gate, attention_query = torch.split(703 query_projected,704 [self.config.hidden_size, self.config.hidden_size, self.config.shared_representation_size],705 dim=-1,706 )707 708 # (target_sequence_length X batch_size X hidden_size)709 residual_weight = torch.sigmoid(residual_weight)710 target_gate = F.silu(target_gate)711 712 if key is None:713 if value is not None:714 raise ValueError("Key and value must be `None` simultaneously")715 projected_key = projected_value = None716 else:717 # (source_sequence_length X batch_size X shared_representation_size)718 projected_key = self.k_proj(key)719 # (source_sequence_length X batch_size X hidden_size)720 projected_value = self.activation(self.v_proj(key))721 722 # (target_sequence_length X batch_size X shared_representation_size)723 # -> (batch_size X target_sequence_length X shared_representation_size)724 attention_query = attention_query.transpose(0, 1)725 if projected_key is not None:726 projected_key = projected_key.transpose(0, 1)727 if projected_value is not None:728 projected_value = projected_value.transpose(0, 1)729 730 # if we're doing incremental decoding, k and v are None and need to be overwritten with past values731 if past_key_values is not None:732 projected_key = prev_cross_key733 projected_value = prev_cross_value734 735 # if we're returning the cache for later use, store these now for later return (can be done without having past_key_values provided)736 if use_cache:737 updated_cross_key = projected_key738 updated_cross_value = projected_value739 740 ctx_len = projected_key.size(1)741 # This is part of a workaround to get around fork/join parallelism742 # not supporting Optional types.743 if key_padding_mask is not None and key_padding_mask.dim() == 0:744 key_padding_mask = None745 746 if key_padding_mask is not None:747 if key_padding_mask.size(0) != bsz:748 raise ValueError("Key padding mask does not align on the batch dimension")749 if key_padding_mask.size(1) != ctx_len:750 raise ValueError("Key padding mask does not align on the sequence length dimension")751 752 if self.attention_activation == "softmax":753 attn_weights = self.softmax_attention(754 attention_query, projected_key, key_padding_mask, num_incremental_steps755 )756 else:757 attn_weights = self.element_attention(758 attention_query, projected_key, key_padding_mask, num_incremental_steps759 )760 761 projected_value = self.hidden_dropout(projected_value, batch_first=True)762 kernel = self.attention_dropout(attn_weights)763 # (batch_size X target_sequence_length X hidden_size)764 # -> (target_sequence_length X batch_size X hidden_size)765 weighted_targets = torch.bmm(kernel, projected_value).transpose(0, 1)766 # (target_sequence_length X batch_size X hidden_size)767 weighted_targets = self.activation(self.h_proj(weighted_targets * target_gate))768 weighted_targets = self.dropout(weighted_targets)769 out = torch.addcmul(query, residual_weight, weighted_targets - query)770 771 if not self.prenorm:772 out = self.norm(out)773 774 outputs = (out, attn_weights) if output_attentions else (out,)775 if use_cache:776 outputs = outputs + (updated_cross_key, updated_cross_value)777 778 return outputs779 780 781class MegaMovingAverageGatedAttention(nn.Module):782 """783 Pure PyTorch implementation of Mega block; see https://huggingface.co/papers/2209.10655 and original fairseq implementation784 at https://github.com/facebookresearch/mega (copyright Meta Research, licensed under MIT License)785 786 Differences from original implementation include hidden state refactor and fixed inconsistency with additive /787 multiplicative attention masks788 """789 790 def __init__(self, config: MegaConfig):791 super().__init__()792 self.config = config793 self.activation = ACT2FN[self.config.activation]794 self.scaling = (795 self.config.shared_representation_size**-0.5 if self.config.attention_activation == "softmax" else None796 )797 self.dropout = MegaDropout(self.config.dropout_prob, is_featurewise=self.config.use_feature_dropout)798 self.hidden_dropout = MegaDropout(799 self.config.hidden_dropout_prob, is_featurewise=self.config.use_feature_dropout800 )801 # attention dropout is standard dropout802 self.attention_dropout = MegaDropout(self.config.attention_probs_dropout_prob, is_featurewise=False)803 804 self.norm = MegaSequenceNorm(805 self.config.normalization_type, self.config.hidden_size, affine=self.config.norm_affine806 )807 self.ema_gate = MegaMultiDimensionDampedEma(config)808 809 self.v_proj = nn.Linear(self.config.hidden_size, self.config.intermediate_size)810 self.mx_proj = nn.Linear(811 self.config.hidden_size,812 self.config.shared_representation_size + self.config.intermediate_size + 2 * self.config.hidden_size,813 )814 self.h_proj = nn.Linear(self.config.intermediate_size, self.config.hidden_size)815 816 self.qk_weight = nn.Parameter(torch.Tensor(2, self.config.shared_representation_size))817 self.qk_bias = nn.Parameter(torch.Tensor(2, self.config.shared_representation_size))818 819 if self.config.relative_positional_bias == "simple":820 self.rel_pos_bias = MegaSimpleRelativePositionalBias(config)821 elif self.config.relative_positional_bias == "rotary":822 self.rel_pos_bias = MegaRotaryRelativePositionalBias(config)823 else:824 raise ValueError(f"Unknown relative positional bias: {self.config.relative_positional_bias}")825 826 self.softmax = nn.Softmax(dim=-1)827 self.attention_function = (828 self.softmax_attention if self.config.attention_activation == "softmax" else self.element_attention829 )830 831 def element_attention(self, query, key, padding_mask, causal_mask):832 """833 Apply element-wise attention via relu^2 or laplace. Same as original implementation but with standardized834 causal attention mask. Expects the Hugging Face standard attention mask paradigm: 1 for not masked, and 0 for835 masked.836 """837 seq_len = key.size(2)838 if padding_mask is not None:839 # (batch_size X number of chunks X 1)840 lengths = padding_mask.sum(-1, keepdim=True)841 # (batch_size X number of chunks X 1 X 1)842 lengths = lengths.clamp(min=1.0).unsqueeze(-1)843 else:844 lengths = seq_len845 846 if causal_mask is not None:847 lengths = causal_mask.sum(dim=-1, keepdim=True)848 849 # (sequence_length X sequence_length)850 bias = self.rel_pos_bias(seq_len)851 if seq_len != query.size(2):852 if query.size(2) != 1:853 raise ValueError("Size mismatch between Q and K in element attention")854 # (1 X sequence_length)855 bias = bias[-1:]856 857 # (batch_size X number of chunks X sequence_length X sequence_length)858 qk = torch.matmul(query, key.transpose(2, 3)) / lengths + bias859 860 attn_weights = ACT2FN[self.config.attention_activation](qk).type_as(qk)861 862 if padding_mask is not None:863 attn_weights = attn_weights * padding_mask.unsqueeze(2)864 865 if causal_mask is not None:866 attn_weights = attn_weights * causal_mask867 868 return attn_weights869 870 def softmax_attention(self, query, key, padding_mask, causal_mask):871 "Standard softmax self-attention, as in the original Transformer paper"872 seq_len = key.size(2)873 # (sequence_length X sequence_length)874 bias = self.rel_pos_bias(seq_len)875 if seq_len != query.size(2):876 if query.size(2) != 1:877 raise ValueError("Size mismatch between Q and K in softmax attention")878 # (1 X sequence_length)879 bias = bias[-1:]880 881 # scaled attention882 query = query * self.scaling883 884 # (batch_size x number of chunks x chunk_size x chunk_size) if chunking885 # (batch_size x 1 x sequence_length x sequence_length) otherwise886 qk = torch.matmul(query, key.transpose(2, 3)) + bias887 888 # apply causal mask (presumed to be 1/0 for not masked / masked)889 # additive, but convert to 0/-inf (which is not explicitly in the Mega source code)890 if causal_mask is not None:891 additive_causal_mask = torch.zeros_like(causal_mask, dtype=qk.dtype)892 additive_causal_mask = additive_causal_mask.masked_fill((1 - causal_mask).bool(), float("-inf"))893 qk = qk + additive_causal_mask894 895 if padding_mask is not None:896 # 1 for tokens which are *not masked*897 # 0 for tokens which are *masked*898 # replace masked tokens with -inf to make softmax ignore them899 # need to invert the padding mask to match what mega original did900 padding_mask = 1 - padding_mask901 padding_mask_all = padding_mask.all(dim=-1, keepdim=True)902 padding_mask = torch.logical_and(padding_mask, ~padding_mask_all)903 qk = qk.masked_fill(padding_mask.unsqueeze(2).to(torch.bool), float("-inf"))904 905 attn_weights = self.softmax(qk).type_as(qk)906 return attn_weights907 908 def forward(909 self,910 input,911 padding_mask: Optional[torch.Tensor] = None,912 causal_mask: Optional[torch.Tensor] = None,913 past_key_values: Optional[Cache] = None,914 output_attentions=False,915 use_cache=False,916 ):917 """918 Mega's self-attention block, which combines multi-headed EMA with traditional self-attention919 920 Args:921 input (`torch.Tensor` of shape `(sequence_length, batch_size, hidden_size)`):922 Hidden states to be updated by Mega's self-attention923 padding_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):924 Indicates which inputs are to be ignored due to padding, where elements are either 1 for *not masked*925 or 0 for *masked*926 causal_mask (`torch.LongTensor` of shape `(sequence_length, sequence_length)`, *optional*):927 Indicates which inputs are to be ignored due to causal attention, where elements are either 1 for *not928 masked* or 0 for *masked*929 past_key_values (`tuple(torch.Tensor)`, *optional*):930 The hidden states returned from the previous timestep during incremental decoding; expects that931 self-attention key, value, and EMA states are the first 3 entries in the tuple932 output_attentions (`bool`, default `False`):933 Whether to return self-attention weights934 use_cache (`bool`, default `False`):935 Whether to perform incremental decoding; uses `past_key_values` as prior state, and returns the updated936 states for use in the next step937 938 Returns:939 `tuple(torch.FloatTensor)` containing various elements depending on configuration ([`MegaConfig`]) and940 inputs:941 - **hidden_states** (`torch.FloatTensor` of shape `(sequence_length, batch_size, hidden_size)`) -- Hidden942 states from target sequence updated by Mega's self-attention943 - **attn_weights** (*optional*, returned when `output_attentions=True`) `torch.FloatTensor` of shape944 `(batch_size, 1, sequence_length, sequence_length)` -- The self-attention weights corresponding to how945 each token in the input sequence attends to every other token946 - **self_key** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,947 sequence_length, config.shared_representation_size)` -- The self-attention key state for use in the next948 step of incremental decoding949 - **self_value** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape `(batch_size,950 sequence_length, config.hidden_size)` -- The self-attention value state for use in the next step of951 incremental decoding952 - **self_ema_state** (*optional*, returned when `use_cache=True`) `torch.FloatTensor` of shape953 `(batch_size, config.ndim)` The incremental EMA state for use in the next step of incremental decoding.954 """955 956 seq_len, bsz, embed_dim = input.size()957 if embed_dim != self.config.hidden_size:958 raise ValueError(f"Input embedding dimension should be {self.config.hidden_size}; received {embed_dim}")959 960 # store inputs for residual connection and handle pre-norm if requested961 residual = input962 if self.config.normalize_before_mega:963 input = self.norm(input)964 965 # (sequence_length X batch_size X hidden_size) -> (sequence_length X batch_size X intermediate_size)966 value = self.activation(self.v_proj(input))967 968 # unpack the incremental state if provided969 # assumed to be (self K, self V, self EMA state, cross K, cross V)970 # also assumes that incremental decoding is working one token at a time, so input sequence length must be 1971 if self.config.is_decoder and (past_key_values is not None):972 if seq_len > 1:973 raise ValueError(f"Incremental decoding only supports self sequence length of 1; received {seq_len}")974 # the first 3 items in the saved states will be these regardless of whether cross-attention is present975 prev_self_key, prev_self_value, prev_ema_state = past_key_values[0:3]976 else:977 prev_self_key = prev_self_value = prev_ema_state = None978 979 # ema output is (sequence_length x batch_size x hidden_size)980 # updated_ema_state will be None if use_cache=False; otherwise (batch_size, config.ndim)981 ema_out, updated_ema_state = self.ema_gate(982 input, attention_mask=padding_mask, prev_state=prev_ema_state, use_cache=use_cache983 )984 ema_out = self.dropout(ema_out)985 986 # (sequence_length X batch_size X hidden_size)987 # -> (sequence_length X batch_size X 2*hidden_size + config.shared_representation_size + config.intermediate_size)988 # - residual_weight -> sigmoid -> applied to residual connection in torch.addcmul989 # - query_key_gates -> split into two components: query_key becomes query and key for attention input, gates becomes gating for self-attention output990 # - intermediate_state -> added to weighted attention output, sent through activation, and has inputs subtracted during991 # torch.addcmul to create the final layer output992 base = self.mx_proj(ema_out)993 residual_weight, query_key_gates, intermediate_state = torch.split(994 base,995 [996 self.config.hidden_size,997 self.config.shared_representation_size + self.config.intermediate_size,998 self.config.hidden_size,999 ],1000 dim=-1,1001 )1002 1003 # (sequence_length X batch_size X hidden_size)1004 residual_weight = torch.sigmoid(residual_weight)1005 1006 # (sequence_length X batch_size X shared_representation_size + intermediate_size)1007 query_key_gates = F.silu(query_key_gates)1008 1009 # split into two different tensors: one for Q/K usage and the other for gating self-attention1010 query_key, attention_gate = torch.split(1011 query_key_gates, [self.config.shared_representation_size, self.config.intermediate_size], dim=-11012 )1013 1014 # (sequence_length X batch_size X shared_representation_size)1015 # -> (sequence_length X batch_size X 1 X shared_representation_size)1016 # -> (sequence_length X batch_size X 2 X shared_representation_size)1017 query_key = query_key.unsqueeze(2) * self.qk_weight + self.qk_bias1018 1019 # (sequence_length X batch_size X 2 X shared_representation_size)1020 # -> 2 tensors of (sequence_length X batch_size X shared_representation_size)1021 query, key = torch.unbind(query_key, dim=2)1022 1023 # (sequence_length X batch_size X dimension)1024 # -> (batch_size X sequence_length X dimension)1025 # where `dimension` is either shared_representation_size (queries and keys) or intermediate_size (values)1026 query = query.transpose(0, 1)1027 key = key.transpose(0, 1)1028 value = value.transpose(0, 1)1029 1030 if self.config.is_decoder:1031 # combine history and current to save updated state (if history is provided)1032 # when chunking is applied, the past states will be None at the end of the chunk, in1033 # which case, proceed as if no K/V history had been provided1034 # saved states are stored with shape (batch_size X sequence_length X dimension)1035 if prev_self_key is not None:1036 key = torch.cat([prev_self_key, key], dim=1)1037 if prev_self_value is not None:1038 value = torch.cat([prev_self_value, value], dim=1)1039 1040 # if not chunking, store as-is1041 if not self.config.use_chunking:1042 updated_self_key = key1043 updated_self_value = value1044 else:1045 curr_len = key.size(1) % self.config.chunk_size1046 if curr_len == 0:1047 # if we're chunking and have reached the end of a chunk, wipe out the saved state1048 updated_self_key = None1049 updated_self_value = None1050 else:1051 updated_self_key = key1052 updated_self_value = value1053 1054 ctx_len = key.size(1) # potentially differs from seq_len because of incremental decoding1055 if not self.config.use_chunking:1056 # if we're not chunking, treat the entire sequence as one long chunk1057 # (batch_size X sequence_length X dimension) -> (batch_size X 1 X sequence_length X dimension)1058 query = query.unsqueeze(1)1059 key = key.unsqueeze(1)1060 value = value.unsqueeze(1)1061 if padding_mask is not None:1062 # (batch_size X sequence_length) -> (batch_size X 1 X sequence_length)1063 padding_mask = padding_mask.unsqueeze(1)1064 else:1065 # otherwise, split the sequences in the batch into `n_chunks` chunks of size `chunk_size`1066 if seq_len < self.config.chunk_size:1067 query = query.unsqueeze(1)1068 else:1069 # (batch_size X sequence_length X dimension) -> (batch_size X n_chunks X chunk_size X dimension)1070 n_chunks = seq_len // self.config.chunk_size1071 query = query.reshape(bsz, n_chunks, self.config.chunk_size, self.config.shared_representation_size)1072 1073 if ctx_len < self.config.chunk_size:1074 key = key.unsqueeze(1)1075 value = value.unsqueeze(1)1076 if padding_mask is not None:1077 padding_mask = padding_mask.unsqueeze(1)1078 else:1079 # (batch_size X sequence_length X dimension) -> (batch_size X n_chunks X chunk_size X dimension)1080 n_chunks = ctx_len // self.config.chunk_size1081 key = key.reshape(bsz, n_chunks, self.config.chunk_size, self.config.shared_representation_size)1082 value = value.reshape(bsz, n_chunks, self.config.chunk_size, self.config.intermediate_size)1083 if padding_mask is not None:1084 padding_mask = padding_mask.view(bsz, n_chunks, self.config.chunk_size)1085 1086 # this is in the original Mega implementation to work around fork/join parallelism not supporting optional types1087 if padding_mask is not None and padding_mask.dim() == 0:1088 padding_mask = None1089 1090 attn_weights = self.attention_function(query, key, padding_mask=padding_mask, causal_mask=causal_mask)1091 1092 value = self.hidden_dropout(value, batch_first=True)1093 kernel = self.attention_dropout(attn_weights)1094 1095 # (batch_size x n_chunks x chunk_size x intermediate_size) -> (sequence_length X batch_size X intermediate_size)1096 weighted_self_output = (1097 torch.matmul(kernel, value).view(bsz, seq_len, self.config.intermediate_size).transpose(0, 1)1098 )1099 1100 # (sequence_length X batch_size X intermediate_size) -> (sequence_length X batch_size X hidden_size)1101 weighted_self_output = self.activation(intermediate_state + self.h_proj(weighted_self_output * attention_gate))1102 weighted_self_output = self.dropout(weighted_self_output)1103 # (sequence_length X batch_size X hidden_size)1104 out = torch.addcmul(residual, residual_weight, weighted_self_output - residual)1105 1106 if not self.config.normalize_before_mega:1107 out = self.norm(out)1108 1109 return_values = (out, attn_weights) if output_attentions else (out,)1110 1111 if self.config.is_decoder:1112 return_values = return_values + (updated_self_key, updated_self_value, updated_ema_state)1113 1114 return return_values1115 1116 1117class MegaNormalizedFeedForwardNetwork(nn.Module):1118 """1119 Normalized feed-forward network used in Mega blocks. Left as-is from original Mega repo aside from retrieving args1120 from Hugging Face config1121 """1122 1123 def __init__(self, config: MegaConfig):1124 super().__init__()1125 1126 self.config = config1127 self.hidden_dim = config.nffn_hidden_size1128 self.act_fn = config.activation1129 self.activation = ACT2FN[config.activation]1130 1131 self.dropout = MegaDropout(self.config.dropout_prob, is_featurewise=self.config.use_feature_dropout)1132 self.hidden_dropout = MegaDropout(1133 self.config.nffn_activation_dropout_prob, is_featurewise=self.config.use_feature_dropout1134 )1135 1136 self.prenorm = self.config.normalize_before_ffn1137 self.norm = MegaSequenceNorm(1138 self.config.normalization_type, self.config.hidden_size, affine=self.config.norm_affine1139 )1140 1141 self.fc1 = nn.Linear(self.config.hidden_size, self.config.nffn_hidden_size)1142 self.fc2 = nn.Linear(self.config.nffn_hidden_size, self.config.hidden_size)1143 1144 def forward(self, inputs):1145 residual = inputs1146 1147 if self.prenorm:1148 inputs = self.norm(inputs)1149 1150 hidden = self.activation(self.fc1(inputs))1151 hidden = self.hidden_dropout(hidden)1152 output = self.fc2(hidden)1153 output = self.dropout(output)1154 output = output + residual1155 1156 if not self.prenorm:1157 output = self.norm(output)1158 1159 return output1160 1161 1162class MegaBlock(nn.Module):1163 def __init__(self, config: MegaConfig):1164 super().__init__()1165 self.seq_len_dim = 11166 self.mega_layer = MegaMovingAverageGatedAttention(config)1167 self.nffn = MegaNormalizedFeedForwardNetwork(config) if config.use_normalized_ffn else None1168 self.is_decoder = config.is_decoder1169 self.add_cross_attention = config.add_cross_attention1170 if self.add_cross_attention:1171 if not self.is_decoder:1172 raise ValueError(f"{self} should be used as a decoder model if cross attention is added")1173 self.cross_attn = MegaGatedCrossAttention(config)1174 else:1175 self.cross_attn = None1176 1177 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")1178 def forward(1179 self,1180 hidden_states: torch.Tensor,1181 attention_mask: Optional[torch.LongTensor] = None,1182 causal_mask: Optional[torch.LongTensor] = None,1183 encoder_hidden_states: Optional[torch.FloatTensor] = None,1184 encoder_attention_mask: Optional[torch.FloatTensor] = None,1185 past_key_values: Optional[Cache] = None,1186 output_attentions: Optional[bool] = False,1187 use_cache: bool = False,1188 ) -> tuple[torch.Tensor]:1189 """1190 A single Mega layer: either encoder or decoder, with optional cross-attention and optional normalized1191 feed-forward layer1192 1193 Args:1194 hidden_states (`torch.Tensor` of shape `(target_sequence_length, batch_size, hidden_size)`):1195 Hidden states to be updated by the Mega block1196 attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):1197 Indicates which entries in the self/target sequence are to be ignored (mostly due to padding), where1198 elements are either 1 for *not masked* or 0 for *masked*. Causal attention is enforced internally.1199 causal_mask (`torch.LongTensor` of shape `(sequence_length, sequence_length)`, *optional*):1200 Indicates which inputs are to be ignored due to causal attention, where elements are either 1 for *not