optimum-intel-internal-testing/tiny-random-phi-4-multimodal
014k
1# Copyright (c) Microsoft Corporation.2# Licensed under the MIT license.3 4#!/usr/bin/env python35 6# activation_checkpointing.py7"""helper function for activation checkpointing"""8 9from typing import Union, Dict, Callable10from functools import partial11from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (12 checkpoint_wrapper,13 offload_wrapper,14 CheckpointImpl,15)16 17 18# utils.py19"""cascade basic blocks"""20 21import math22#import backoff23import random24import numpy as np25from typing import Optional, Tuple, Union26import torch27from torch import nn28from torch import Tensor29import torch.nn.functional as F30 31 32# conformer_encoder.py33"""ConformerEncoder Module"""34 35from typing import Optional, Tuple, List, Literal36import abc37import math38import numpy as np39 40import torch41from torch import nn, Tensor42 43from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper44from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel45 46 47# activation_checkpointing.py48def validate_checkpointing_config(activation_checkpointing):49 """validate activation checkpointing configuration"""50 if isinstance(activation_checkpointing, str):51 assert activation_checkpointing in (52 "",53 "checkpoint",54 "offload",55 ), "activation_checkpointing has to be a dict or a str in ('', 'checkpoint', 'offload')."56 elif isinstance(activation_checkpointing, dict):57 assert activation_checkpointing.get("module", "transformer") in (58 "transformer",59 "attention",60 ), "module in activation_checkpointing has to be in ('transformer', 'attention')."61 else:62 raise ValueError("activation_checkpointing has to be a str or dict.")63 64 65def embedding_checkpoint_wrapper(66 activation_checkpointing: Union[str, Dict],67) -> Callable:68 """return encoder embedding activation checkpoint wrapper"""69 validate_checkpointing_config(activation_checkpointing)70 71 if isinstance(activation_checkpointing, str):72 if activation_checkpointing:73 if activation_checkpointing == "offload":74 return offload_wrapper75 return partial(checkpoint_wrapper)76 return lambda x: x77 78 if isinstance(activation_checkpointing, dict):79 enabled = activation_checkpointing.get("embed", False)80 if enabled:81 offloading = activation_checkpointing.get("offload", False)82 if offloading:83 return offload_wrapper84 impl = (85 CheckpointImpl.REENTRANT86 if activation_checkpointing.get("reentrant", False)87 else CheckpointImpl.NO_REENTRANT88 )89 return partial(checkpoint_wrapper, checkpoint_impl=impl)90 return lambda x: x91 raise ValueError("Invalid activation_checkpointing config")92 93 94def encoder_checkpoint_wrapper(95 activation_checkpointing: Union[str, Dict],96 layer_cls: type,97 idx: int = 0,98) -> Callable:99 """return encoder activation checkpoint wrapper"""100 validate_checkpointing_config(activation_checkpointing)101 102 if isinstance(activation_checkpointing, str):103 if activation_checkpointing:104 if activation_checkpointing == "offload":105 return offload_wrapper106 return partial(checkpoint_wrapper)107 return lambda x: x108 109 if isinstance(activation_checkpointing, dict):110 target_layer_cls = activation_checkpointing.get("module", "transformer")111 if target_layer_cls.lower() == "transformer":112 target_layer_cls = (113 "EncoderLayer",114 "ConformerEncoderLayer",115 )116 elif target_layer_cls.lower() == "attention":117 target_layer_cls = ("MultiHeadedAttention", "MultiHeadAttention")118 checkpointing_interval = activation_checkpointing.get("interval", 1)119 offloading = activation_checkpointing.get("offload", False)120 impl = (121 CheckpointImpl.REENTRANT122 if activation_checkpointing.get("reentrant", True)123 else CheckpointImpl.NO_REENTRANT124 )125 126 if idx % checkpointing_interval == 0 and layer_cls.__name__ in target_layer_cls:127 if offloading:128 return offload_wrapper129 return partial(checkpoint_wrapper, checkpoint_impl=impl)130 return lambda x: x131 132 raise ValueError("Invalid activation_checkpointing config")133 134 135def attn_checkpointing(activation_checkpointing: Union[str, Dict], i) -> Union[str, Dict]:136 """return activation checkpointing config for attention layer"""137 if isinstance(activation_checkpointing, str):138 return ""139 140 if isinstance(activation_checkpointing, dict):141 target_layer_cls = activation_checkpointing.get("module", "transformer")142 checkpointing_interval = activation_checkpointing.get("interval", 1)143 if target_layer_cls == "attention" and i % checkpointing_interval == 0:144 return activation_checkpointing145 return ""146 147 raise ValueError("Invalid activation_checkpointing config")148 149 150# utils.py151class Block(nn.Module):152 """Block abstract module"""153 154 def __init__(self, input_size, output_size):155 super().__init__()156 self.input_size = input_size157 self.output_size = output_size158 159def get_activation(name="relu"):160 """Select an activation function by name161 162 Args:163 name: str164 activation function name,165 one of ["relu", "gelu", "swish", "sigmoid"],166 default "relu".167 """168 name = name.lower()169 if name == "relu":170 return nn.ReLU(inplace=True)171 if name == "gelu":172 return nn.GELU()173 if name == "swish":174 return Swish()175 if name == "sigmoid":176 return torch.nn.Sigmoid()177 return nn.Identity()178 179def adaptive_enc_mask(x_len, chunk_start_idx, left_window=0, right_window=0):180 """181 The function is very important for Transformer Transducer Streaming mode182 Args:183 xs_len (int): sequence length184 chunk_start_idx (list): first idx of each chunk, such as [0,18,36,48]. It also supports adaptive chunk size [0,10,15,45]185 left_window (int): how many left chunks can be seen186 right_window (int): how many right chunks can be seen. It is used for chunk overlap model.187 Returns:188 mask (torch.Tensor): a mask tensor for streaming model189 Torch 1.0.1190 tensor([[1., 1., 0., 0.],191 [0., 1., 1., 0.],192 [0., 0., 1., 1.]])193 Torch 1.4.1194 tensor([[True., True., False., False.],195 [False., True., True., False.],196 [False., False., True., True.]])197 """198 chunk_start_idx = torch.Tensor(199 chunk_start_idx200 ).long() # first idx of each chunk, such as [0,18,36,48].201 start_pad = torch.nn.functional.pad(202 chunk_start_idx, (1, 0)203 ) # append 0 to the beginning, so it becomes [0, 0, 18, 36, 48]204 end_pad = torch.nn.functional.pad(205 chunk_start_idx, (0, 1), value=x_len206 ) # append x_len to the end, so it becomes [0,18,36,48, x_len]207 seq_range = torch.arange(0, x_len).unsqueeze(-1) # seq_range size: [x_len, 1]208 idx = ((seq_range < end_pad) & (seq_range >= start_pad)).nonzero()[:, 1] # idx size: [x_len]209 boundary = end_pad[idx] # boundary size: [x_len]210 seq_range_expand = (211 torch.arange(0, x_len).unsqueeze(0).expand(x_len, -1)212 ) # seq_range_expand size [x_len, x_len]213 idx_left = idx - left_window214 idx_left[idx_left < 0] = 0215 boundary_left = start_pad[idx_left]216 mask_left = seq_range_expand >= boundary_left.unsqueeze(-1)217 idx_right = idx + right_window218 idx_right[idx_right > len(chunk_start_idx)] = len(chunk_start_idx)219 boundary_right = end_pad[idx_right]220 mask_right = seq_range_expand < boundary_right.unsqueeze(-1)221 return mask_left & mask_right222 223class Swish(nn.Module):224 """Implement Swish activation module.225 From https://arxiv.org/pdf/2005.03191.pdf226 227 """228 229 def __init__(self) -> None:230 super().__init__()231 self.act_fn = nn.Sigmoid()232 233 def forward(self, x: Tensor) -> Tensor:234 """Apply Swish function235 236 Args:237 x: torch.Tensor238 Input.239 """240 return x * self.act_fn(x)241 242class GLU(nn.Module):243 """Implement Gated Linear Unit (GLU) module"""244 245 def __init__(self, dim: int = -1, act_name: str = "sigmoid") -> None:246 super().__init__()247 self.dim = dim248 self.act_name = act_name.lower()249 250 if self.act_name == "relu":251 self.act_fn = nn.ReLU(inplace=True)252 elif self.act_name == "gelu":253 self.act_fn = nn.GELU()254 elif self.act_name == "swish":255 self.act_fn = Swish()256 elif self.act_name == "sigmoid":257 self.act_fn = nn.Sigmoid()258 else:259 self.act_fn = nn.Identity()260 261 def forward(self, x: Tensor) -> Tensor:262 """GLU forward263 Apply Swish function on the first half of input matrices264 with sigmoid of the second half.265 266 Args:267 x: torch.Tensor268 Input.269 270 """271 half_x, gate = x.chunk(2, dim=self.dim)272 return half_x * self.act_fn(gate)273 274# TODO: Abdel, this can be improved using GLU module275class GLUPointWiseConv(nn.Module):276 """GLUPointWiseConv module277 used for conformer architecture,278 for more details see:279 https://arxiv.org/pdf/2005.08100v1.pdf280 281 Args:282 input_dim: int283 input channel size.284 output_dim: int285 output channel size.286 kernel_size: int287 kernel size288 glu_type: str, optional289 activation function one of290 ["sigmoid", "relu", "gelu"]291 default "sigmoid".292 bias_in_glu: bool, optional293 use addtive bias in glu294 causal: bool, optional295 if set to True, padding is set to the half of296 kernel size, ie, convolution can't see future frames.297 default False.298 299 """300 301 def __init__(302 self, input_dim, output_dim, kernel_size, glu_type="sigmoid", bias_in_glu=True, causal=False303 ):304 super().__init__()305 306 self.glu_type = glu_type307 self.output_dim = output_dim308 self.bias_in_glu = bias_in_glu309 if causal:310 self.ext_pw_conv_1d = nn.Conv1d(311 input_dim, output_dim * 2, kernel_size, 1, padding=(kernel_size - 1)312 )313 else:314 self.ext_pw_conv_1d = nn.Conv1d(315 input_dim, output_dim * 2, kernel_size, 1, padding=(kernel_size - 1) // 2316 )317 318 if glu_type == "sigmoid":319 self.glu_act = nn.Sigmoid()320 elif glu_type == "relu":321 self.glu_act = nn.ReLU()322 elif glu_type == "gelu":323 self.glu_act = nn.GELU()324 elif glu_type == "swish":325 self.glu_act = Swish()326 else:327 raise ValueError(f"Unsupported activation type {self.glu_act}")328 329 if bias_in_glu:330 self.b1 = nn.Parameter(torch.zeros(1, output_dim, 1))331 self.b2 = nn.Parameter(torch.zeros(1, output_dim, 1))332 333 def forward(self, x):334 """335 Args:336 x: torch.Tensor337 input tensor338 """339 # to be consistent with GLULinear, we assume the input always has the #channel (#dim) in the last dimension of the tensor, so need to switch the dimension first for 1D-Conv case340 x = x.permute([0, 2, 1])341 x = self.ext_pw_conv_1d(x)342 if self.glu_type == "bilinear":343 if self.bias_in_glu:344 x = (x[:, 0 : self.output_dim, :] + self.b1) * (345 x[:, self.output_dim : self.output_dim * 2, :] + self.b2346 )347 else:348 x = (x[:, 0 : self.output_dim, :]) * (349 x[:, self.output_dim : self.output_dim * 2, :]350 )351 else:352 if self.bias_in_glu:353 x = (x[:, 0 : self.output_dim, :] + self.b1) * self.glu_act(354 x[:, self.output_dim : self.output_dim * 2, :] + self.b2355 )356 else:357 x = (x[:, 0 : self.output_dim, :]) * self.glu_act(358 x[:, self.output_dim : self.output_dim * 2, :]359 )360 361 x = x.permute([0, 2, 1])362 return x363 364 365class DepthWiseSeperableConv1d(nn.Module):366 """DepthWiseSeperableConv1d module used in Convnet module367 for the conformer, for more details see:368 https://arxiv.org/pdf/2005.08100v1.pdf369 370 Args:371 input_dim: int372 input channel size.373 depthwise_seperable_out_channel: int374 if set different to 0, the number of depthwise_seperable_out_channel375 will be used as a channel_out of the second conv1d layer.376 otherwise, it equal to 0, the second conv1d layer is skipped.377 kernel_size: int378 kernel_size379 depthwise_multiplier: int380 number of input_dim channels duplication. this value381 will be used to compute the hidden channels of the Conv1D.382 padding: int, optional383 padding for the conv1d,384 default: 0.385 386 """387 388 def __init__(389 self,390 input_dim,391 depthwise_seperable_out_channel,392 kernel_size,393 depthwise_multiplier,394 padding=0,395 ):396 super().__init__()397 398 self.dw_conv = nn.Conv1d(399 input_dim,400 input_dim * depthwise_multiplier,401 kernel_size,402 1,403 padding=padding,404 groups=input_dim,405 )406 407 if depthwise_seperable_out_channel != 0:408 self.pw_conv = nn.Conv1d(409 input_dim * depthwise_multiplier, depthwise_seperable_out_channel, 1, 1, 0410 )411 else:412 self.pw_conv = nn.Identity()413 self.depthwise_seperable_out_channel = depthwise_seperable_out_channel414 415 def forward(self, x):416 """417 418 Args:419 x: torch.Tensor420 input tensor421 """422 x = self.dw_conv(x)423 if self.depthwise_seperable_out_channel != 0:424 x = self.pw_conv(x)425 return x426 427 428class ConvModule(nn.Module):429 """ConvModule Module for the conformer block.430 for more details see:431 https://arxiv.org/pdf/2005.08100v1.pdf432 433 Args:434 input_dim: int435 input channel size.436 ext_pw_out_channel: int437 if > 0, ext_pw_out_channel is a dim channel size438 for the last pointwise conv after swish activation.439 depthwise_seperable_out_channel: int440 if set different to 0, the number of depthwise_seperable_out_channel441 will be used as a channel_out of the second conv1d layer.442 otherwise, it equal to 0, the second conv1d layer is skipped.443 ext_pw_kernel_size: int444 kernel size of the conv pointwise of the conformer.445 kernel_size: int446 kernel size.447 depthwise_multiplier: int448 number of input_dim channels duplication. this value449 will be used to compute the hidden channels of the Conv1D.450 dropout_rate: float451 dropout rate.452 causal: bool, optional453 if set to True, convolution have no access454 to future frames. default False.455 batch_norm: bool, optional456 if set to True, apply batchnorm before activation.457 default False458 chunk_se: int, optional459 0 for offline SE.460 1 for streaming SE, where mean is computed461 by accumulated history until current chunk_se.462 2 for streaming SE, where mean is computed463 by only the current chunk.464 chunk_size: int, optional465 chunk size for cnn. default 18466 activation: str, optional467 activation function used in ConvModule,468 default: "relu".469 glu_type: str, optional470 activation function used for the glu,471 default: "sigmoid".472 bias_in_glu: bool, optional473 if set to True, use additive bias in the weight module474 before GLU.475 linear_glu_in_convm: bool, optional476 if set to True, use GLULinear module,477 otherwise, used GLUPointWiseConv module.478 default to False.479 export: bool, optional,480 if set to True, padding is equal to 0. This is for inference,481 or onnx export. Typically this is set by the export program or482 the decoder program, and it isn't present in your config file.483 default False484 """485 486 def __init__(487 self,488 input_dim,489 ext_pw_out_channel,490 depthwise_seperable_out_channel,491 ext_pw_kernel_size,492 kernel_size,493 depthwise_multiplier,494 dropout_rate,495 causal=False,496 batch_norm=False,497 chunk_se=0,498 chunk_size=18,499 activation="relu",500 glu_type="sigmoid",501 bias_in_glu=True,502 linear_glu_in_convm=False,503 export=False,504 ):505 super().__init__()506 self.layer_norm = nn.LayerNorm(input_dim)507 self.input_dim = input_dim508 self.ext_pw_out_channel = ext_pw_out_channel509 self.ext_pw_kernel_size = ext_pw_kernel_size510 self.depthwise_seperable_out_channel = depthwise_seperable_out_channel511 self.glu_type = glu_type512 self.bias_in_glu = bias_in_glu513 self.linear_glu_in_convm = linear_glu_in_convm514 self.causal = causal515 516 self._add_ext_pw_layer()517 518 self.batch_norm = batch_norm519 self.kernel_size = kernel_size520 521 if batch_norm:522 self.bn_layer = nn.BatchNorm1d(input_dim)523 524 self.act = get_activation(activation)525 self.dropout = nn.Dropout(dropout_rate)526 self.export = export527 528 if causal:529 if export: # Inference only.530 padding = 0 # A cache is concatenated to the left. No padding in the kernel.531 else:532 # Training only. Padding will be added symmetrically on both sides.533 # After convolution, clip off kernel_size-1 points on the right.534 padding = kernel_size - 1535 else:536 padding = (kernel_size - 1) // 2537 538 self.dw_sep_conv_1d = DepthWiseSeperableConv1d(539 input_dim,540 depthwise_seperable_out_channel,541 kernel_size,542 depthwise_multiplier,543 padding=padding,544 )545 546 if depthwise_seperable_out_channel != 0:547 if input_dim != depthwise_seperable_out_channel:548 self.ln2 = nn.Linear(depthwise_seperable_out_channel, input_dim)549 else:550 if depthwise_multiplier != 1:551 self.ln2 = nn.Linear(input_dim * depthwise_multiplier, input_dim)552 553 def _add_ext_pw_layer(self):554 """555 This function is an extension of __init__ function556 and dedicated to the convolution module creation557 of the conformer.558 """559 self.ln1 = self.glu = self.bn_layer = self.ext_pw_conv_1d = nn.Identity() # jit hacks.560 self.squeeze_excitation = nn.Identity() # jit.561 self.apply_ln1 = self.fix_len1 = False # jit.562 563 if self.ext_pw_out_channel != 0:564 if self.causal:565 self.ext_pw_conv_1d = nn.Conv1d(566 self.input_dim,567 self.ext_pw_out_channel,568 self.ext_pw_kernel_size,569 1,570 padding=(self.ext_pw_kernel_size - 1),571 )572 if self.ext_pw_kernel_size > 1:573 self.fix_len1 = True574 else:575 self.fix_len1 = False576 else:577 self.ext_pw_conv_1d = nn.Conv1d(578 self.input_dim,579 self.ext_pw_out_channel,580 self.ext_pw_kernel_size,581 1,582 padding=(self.ext_pw_kernel_size - 1) // 2,583 )584 self.fix_len1 = False585 586 if self.linear_glu_in_convm:587 self.glu = GLULinear(588 self.input_dim, self.ext_pw_out_channel, self.glu_type, self.bias_in_glu589 )590 else:591 self.glu = GLUPointWiseConv(592 self.input_dim,593 self.ext_pw_out_channel,594 self.ext_pw_kernel_size,595 self.glu_type,596 self.bias_in_glu,597 self.causal,598 )599 600 if self.input_dim != self.ext_pw_out_channel:601 self.apply_ln1 = True602 self.ln1 = nn.Linear(self.ext_pw_out_channel, self.input_dim)603 else:604 self.apply_ln1 = False605 else:606 self.pw_conv_simplify_w = torch.nn.Parameter(torch.ones(3))607 self.pw_conv_simplify_b = torch.nn.Parameter(torch.zeros(3))608 609 def forward(self, x):610 """ConvModule Forward.611 612 Args:613 x: torch.Tensor614 input tensor.615 """616 x = self.layer_norm(x)617 618 if self.ext_pw_out_channel != 0:619 x = self.glu(x)620 if self.causal and self.ext_pw_kernel_size > 1:621 x = x[:, : -(self.ext_pw_kernel_size - 1), :]622 if self.apply_ln1:623 x = self.ln1(x)624 else:625 x_0 = x * self.pw_conv_simplify_w[0] + self.pw_conv_simplify_b[0]626 x_1 = x * self.pw_conv_simplify_w[1] + self.pw_conv_simplify_b[1]627 x = x_0 + x_1628 629 x = x.permute([0, 2, 1])630 631 x = self.dw_sep_conv_1d(x)632 if self.causal and self.kernel_size > 1:633 x = x[:, :, : -(self.kernel_size - 1)]634 if hasattr(self, "ln2"):635 x = x.permute([0, 2, 1])636 x = self.ln2(x)637 x = x.permute([0, 2, 1])638 if self.batch_norm:639 x = self.bn_layer(x)640 x = self.act(x)641 642 if self.ext_pw_out_channel != 0:643 x = self.ext_pw_conv_1d(x)644 if self.fix_len1:645 x = x[:, :, : -(self.ext_pw_kernel_size - 1)]646 647 if self.apply_ln1:648 x = x.permute([0, 2, 1])649 x = self.ln1(x)650 x = x.permute([0, 2, 1])651 652 x = x.permute([0, 2, 1])653 else:654 x = x.unsqueeze(1).permute([0, 1, 3, 2])655 x = x * self.pw_conv_simplify_w[2] + self.pw_conv_simplify_b[2]656 x = x.squeeze(1)657 658 x = self.dropout(x)659 return x660 661class GLULinear(nn.Module):662 """Linear + GLU module663 664 Args:665 input_dim: int666 input size667 output_dim: int668 output size.669 glu_type:670 activation function name used in glu module.671 default "sigmoid" (swish function).672 bias_in_glu: bool, optional673 If True, the addtive bias is added. Default False.674 """675 676 def __init__(677 self,678 input_dim,679 output_dim,680 glu_type="sigmoid",681 bias_in_glu=True,682 ):683 super().__init__()684 self.linear = nn.Linear(input_dim, output_dim * 2, bias_in_glu)685 self.glu_act = GLU(-1, glu_type)686 687 def forward(self, x):688 """GLULinear forward689 690 Args:691 x: torch.Tensor692 inpute tensor.693 """694 x = self.linear(x)695 return self.glu_act(x)696 697class FeedForward(nn.Module):698 """FeedForward Module.699 For more details see Conformer paper:700 https://arxiv.org/pdf/2005.08100.pdf701 702 Args:703 d_model: int704 input size.705 d_inner: int706 output size.707 dropout_rate: float,708 dropout rate.709 activation: str,710 activation function name,711 one of ["relu", "swish", "sigmoid"],712 sigmoid activation is only used with "glu_in_fnn=True",713 default "sigmoid".714 bias_in_glu: bool, optional715 """716 717 def __init__(718 self,719 d_model,720 d_inner,721 dropout_rate,722 activation="sigmoid",723 bias_in_glu=True,724 ):725 super().__init__()726 self.d_model = d_model727 self.d_inner = d_inner728 729 self.layer_norm = nn.LayerNorm(d_model)730 module = GLULinear(d_model, d_inner, activation, bias_in_glu)731 self.net = nn.Sequential(732 module,733 nn.Dropout(dropout_rate),734 nn.Linear(d_inner, d_model),735 nn.Dropout(dropout_rate),736 )737 738 def forward(self, x):739 """FeedForward forward function.740 741 Args:742 x: torch.Tensor743 input tensor.744 """745 out = self.net(self.layer_norm(x))746 747 return out748 749#### positional encoding starts here750def _pre_hook(751 state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs752):753 """Perform pre-hook in load_state_dict for backward compatibility.754 755 Note:756 We saved self.pe until v.0.5.2 but we have omitted it later.757 Therefore, we remove the item "pe" from `state_dict` for backward compatibility.758 759 """760 k = prefix + "pe"761 if k in state_dict:762 state_dict.pop(k)763 764class T5RelativeAttentionLogitBias(nn.Module):765 """766 This module implements the relative position bias described in Section 2.1 of767 the T5 paper: https://arxiv.org/pdf/1910.10683.pdf768 769 The Huggingface implementation is used as a reference770 https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/t5/modeling_t5.py#L435771 772 Modifies attention as Q*K^T + B, where B is a learned scalar bias based on relative position773 of the query and key. It is HxNxN, where H is the number of heads, N is the sequence length.774 775 I've made these modifications to the original T5 bias:776 - Skipping of the bucketing step. Original T5 bias converted rel position distances into777 logarithmically increasing buckets. This is supposed to help with length generalization.778 - I just directly use rel position index as bias values, as we don't need length779 generalization (40s max is good enough for ASR encoder), and it keeps ONNX export simple.780 - I've also extended it so that biases can be asymmetric, the default implementation treats781 L->R and R->L the same. Asymmetric was found to yield better results in my experiments.782 783 Args:784 num_heads: int785 Number of attention heads786 num_buckets: int787 Number of buckets to use for relative attention bias. This is the size of the learnable788 bias parameter. Bucketing is not yet supported, so this defaults to -1 which means789 no bucketing is used (max_distance determines size of bias param).790 max_distance: int791 Maximum distance to use for relative attention bias. With num_buckets=-1, this directly792 controls the max size of the bias parameter. When num_buckets > 0 is supported, this793 will control the maximum distance for logarithmic bucketing after which all positions794 are in the same bucket.795 symmetric: bool796 Whether to use symmetric or asymmetric biases. symmetric=False uses 2x number of bias797 params to distinguish L->R from R->L. This was found to be better for the encoder.798 """799 800 def __init__(self, num_heads, num_buckets=-1, max_distance=1000, symmetric=False):801 super().__init__()802 self.num_heads = num_heads803 self.num_buckets = num_buckets804 self.max_distance = max_distance805 self.symmetric = symmetric806 self._skip_bucketing = self.num_buckets < 0807 if self._skip_bucketing:808 self.num_buckets = max_distance809 else:810 raise NotImplementedError("T5 attention bias with bucketed positions is not yet tested")811 if not self.symmetric:812 self.num_buckets *= 2813 self.bias_values = nn.Embedding(self.num_buckets, self.num_heads)814 815 def forward(self, x):816 # instantiate bias compatible with shape of x817 maxpos = x.size(1)818 context_position = torch.arange(maxpos, device=x.device, dtype=torch.long)[:, None]819 memory_position = torch.arange(maxpos, device=x.device, dtype=torch.long)[None, :]820 relative_position = memory_position - context_position821 # clipping to a maximum distance using ops that play well with ONNX export822 relative_position = relative_position.masked_fill(823 relative_position < -self.max_distance, -self.max_distance824 )825 relative_position = relative_position.masked_fill(826 relative_position > self.max_distance - 1, self.max_distance - 1827 )828 829 # mapping from relative position to index in the bias parameter830 if self._skip_bucketing:831 bias_idx = relative_position832 else:833 bias_idx = self._bucket_relative_position(relative_position)834 if self.symmetric:835 bias_idx = bias_idx.abs()836 else:837 bias_idx += self.num_buckets // 2838 839 t5_rel_att_bias = self.bias_values(bias_idx) # [L, L, H]840 t5_rel_att_bias = t5_rel_att_bias.permute(2, 0, 1).unsqueeze(0) # [1, H, L, L]841 842 return t5_rel_att_bias843 844 def _bucket_relative_position(self, relative_position):845 # this is a placeholder (isn't tested, likely buggy) using HuggingFace implem as a reference846 # this also needs to be extended to support asymmetric +/- ve positions847 relative_buckets = 0848 if not self.causal:849 num_buckets //= 2850 relative_buckets += (relative_position > 0).to(torch.long) * num_buckets851 relative_position = torch.abs(relative_position)852 else:853 relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))854 # now relative_position is in the range [0, inf)855 856 # half of the buckets are for exact increments in positions857 max_exact = num_buckets // 2858 is_small = relative_position < max_exact859 860 # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance861 relative_position_if_large = max_exact + (862 torch.log(relative_position.float() / max_exact)863 / math.log(self.max_distance / max_exact)864 * (num_buckets - max_exact)865 ).to(torch.long)866 relative_position_if_large = torch.min(867 relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)868 )869 870 relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)871 return relative_buckets872 873class AbsolutePositionalEncoding(nn.Module):874 """Absolute Positional encoding module.875 This module implement Absolute sinusoidal positional encoding876 from: https://arxiv.org/pdf/1706.03762.pdf877 878 Args:879 d_model: int880 Input embedding size.881 dropout_rate: float882 dropout rate883 max_len: int, optional884 Maximum input length sequence, Default 5000885 886 """887 888 def __init__(self, d_model, dropout_rate, max_len=5000):889 """Construct an PositionalEncoding object."""890 super().__init__()891 self.d_model = d_model892 self.xscale = math.sqrt(self.d_model)893 self.dropout = torch.nn.Dropout(p=dropout_rate)894 self.pe = None895 self.extend_pe(torch.tensor(0.0).expand(1, max_len))896 self._register_load_state_dict_pre_hook(_pre_hook)897 898 def extend_pe(self, x):899 """Reset the positional encodings.900 901 Args:902 x: torch.Tensor903 """904 if self.pe is not None:905 if self.pe.size(1) >= x.size(1):906 if self.pe.dtype != x.dtype or self.pe.device != x.device:907 self.pe = self.pe.to(dtype=x.dtype, device=x.device)908 return909 pe = torch.zeros(x.size(1), self.d_model)910 position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)911 div_term = torch.exp(912 torch.arange(0, self.d_model, 2, dtype=torch.float32)913 * -(math.log(10000.0) / self.d_model)914 )915 pe[:, 0::2] = torch.sin(position * div_term)916 pe[:, 1::2] = torch.cos(position * div_term)917 pe = pe.unsqueeze(0)918 self.pe = pe.to(device=x.device, dtype=x.dtype)919 920 def forward(self, x: torch.Tensor):921 """Add positional encoding.922 923 Args:924 x: torch.Tensor925 Input tensor. shape is (batch, time, ...)926 927 Returns:928 torch.Tensor: Encoded tensor. Its shape is (batch, time, ...)929 930 """931 self.extend_pe(x)932 x = x * self.xscale + self.pe[:, : x.size(1)]933 return self.dropout(x)934 935#### forward embedding layers starts here936 937#@backoff.on_exception(backoff.expo, Exception, max_tries=10)938def np_loadtxt_with_retry(filepath):939 """np.loadtxt with retry940 941 Args:942 filepath: str943 file path to the numpy array.944 """945 result = np.loadtxt(filepath, dtype="f")946 return result947 948class MeanVarianceNormLayer(nn.Module):949 """Mean/variance normalization layer.950 951 Will substract mean and multiply input by inverted standard deviation.952 Typically used as a very first layer in a model.953 954 Args:955 input_size: int956 layer input size.957 """958 959 def __init__(self, input_size):960 super().__init__()961 self.input_size = input_size962 self.register_buffer("global_mean", torch.zeros(input_size))963 self.register_buffer("global_invstd", torch.ones(input_size))964 self.global_mean: Optional[Tensor]965 self.global_invstd: Optional[Tensor]966 967 def forward(self, input_: Tensor) -> Tensor:968 """MeanVarianceNormLayer Forward969 970 Args:971 input_: torch.Tensor972 input tensor.973 """974 return (input_ - self.global_mean) * self.global_invstd975 976 def load_mean_invstd(self, mean_file, invstd_file, cuside_features=False):977 """Load feature mean and variance used for normalization.978 979 Args:980 mean_file: str981 path to the feature mean statistics file.982 invstd_file: str983 path to the features inverted standard deviation984 statistics file.985 cuside_features: bool986 Boolean that indicates CUSIDE is being used.987 The statistics of CUSIDE features are copied988 from the normal features989 """990 self.global_mean.data = torch.from_numpy(np_loadtxt_with_retry(mean_file))991 self.global_invstd.data = torch.from_numpy(np_loadtxt_with_retry(invstd_file))992 993 if cuside_features:994 self.global_mean.data = torch.cat((self.global_mean.data, self.global_mean.data), 0)995 self.global_invstd.data = torch.cat(996 (self.global_invstd.data, self.global_invstd.data), 0997 )998 999class CausalConv1D(nn.Conv1d):1000 """1001 A causal version of nn.Conv1d where each step would have limited access to locations on its right or left1002 All arguments are the same as nn.Conv1d except padding.1003 1004 If padding is set None, then paddings are set automatically to make it a causal convolution where each location would not see any steps on its right.1005 1006 If padding is set as a list (size of 2), then padding[0] would be used as left padding and padding[1] as right padding.1007 It would make it possible to control the number of steps to be accessible on the right and left.1008 This mode is not supported when stride > 1. padding[0]+padding[1] should be equal to (kernel_size - 1).1009 """1010 1011 def __init__(1012 self,1013 in_channels: int,1014 out_channels: int,1015 kernel_size: int,1016 stride: int = 1,1017 padding: Union[str, int] = 0,1018 dilation: int = 1,1019 groups: int = 1,1020 bias: bool = True,1021 padding_mode: str = "zeros",1022 device=None,1023 dtype=None,1024 ) -> None:1025 self.cache_drop_size = None1026 if padding is None:1027 self._left_padding = kernel_size - 11028 self._right_padding = stride - 11029 else:1030 if stride != 1 and padding != kernel_size - 1:1031 raise ValueError("No striding allowed for non-symmetric convolutions!")1032 if isinstance(padding, int):1033 self._left_padding = padding1034 self._right_padding = padding1035 elif (1036 isinstance(padding, list)1037 and len(padding) == 21038 and padding[0] + padding[1] == kernel_size - 11039 ):1040 self._left_padding = padding[0]1041 self._right_padding = padding[1]1042 else:1043 raise ValueError(f"Invalid padding param: {padding}!")1044 1045 self._max_cache_len = self._left_padding1046 1047 super().__init__(1048 in_channels=in_channels,1049 out_channels=out_channels,1050 kernel_size=kernel_size,1051 stride=stride,1052 padding=0,1053 dilation=dilation,1054 groups=groups,1055 bias=bias,1056 padding_mode=padding_mode,1057 device=device,1058 dtype=dtype,1059 )1060 1061 def update_cache(self, x, cache=None):1062 if cache is None:1063 new_x = F.pad(x, pad=(self._left_padding, self._right_padding))1064 next_cache = cache1065 else:1066 new_x = F.pad(x, pad=(0, self._right_padding))1067 new_x = torch.cat([cache, new_x], dim=-1)1068 if self.cache_drop_size > 0:1069 next_cache = new_x[:, :, : -self.cache_drop_size]1070 else:1071 next_cache = new_x1072 next_cache = next_cache[:, :, -cache.size(-1) :]1073 return new_x, next_cache1074 1075 def forward(self, x, cache=None):1076 x, cache = self.update_cache(x, cache=cache)1077 x = super().forward(x)1078 if cache is None:1079 return x1080 else:1081 return x, cache1082 1083 1084class CausalConv2D(nn.Conv2d):1085 """1086 A causal version of nn.Conv2d where each location in the 2D matrix would have no access to locations on its right or down1087 All arguments are the same as nn.Conv2d except padding which should be set as None1088 """1089 1090 def __init__(1091 self,1092 in_channels: int,1093 out_channels: int,1094 kernel_size: int,1095 stride: int = 1,1096 padding: Union[str, int] = 0,1097 dilation: int = 1,1098 groups: int = 1,1099 bias: bool = True,1100 padding_mode: str = "zeros",1101 device=None,1102 dtype=None,1103 ) -> None:1104 if padding is not None:1105 raise ValueError("Argument padding should be set to None for CausalConv2D.")1106 self._left_padding = kernel_size - 11107 self._right_padding = stride - 11108 1109 padding = 01110 super().__init__(1111 in_channels,1112 out_channels,1113 kernel_size,1114 stride,1115 padding,1116 dilation,1117 groups,1118 bias,1119 padding_mode,1120 device,1121 dtype,1122 )1123 1124 def forward(1125 self,1126 x,1127 ):1128 if self.training:1129 x = F.pad(1130 x,1131 pad=(1132 self._left_padding,1133 self._right_padding,1134 self._left_padding,1135 self._right_padding,1136 ),1137 )1138 else:1139 x = F.pad(1140 x,1141 pad=(self._left_padding, self._right_padding, 0, 0),1142 )1143 x = super().forward(x)1144 return x1145 1146 1147class NemoConvSubsampling(torch.nn.Module):1148 """Convlutional subsampling module, taken from NeMo ASR1149 (https://github.com/NVIDIA/NeMo/blob/b367413645d5c72db3c2c96e46e95a34501479cf/nemo/collections/asr/parts/submodules/subsampling.py)1150 1151 Striding Subsampling: "Speech-Transformer: A No-Recurrence Sequence-to-Sequence Model for1152 Speech Recognition" by Linhao Dong et al. (https://ieeexplore.ieee.org/document/8462506)1153 1154 1155 Compared with the EncoderConv2D (`input_layer: custom`), this is a much simplified approach,1156 and uses no LayerNorm and far fewer Conv2Ds. Moreover, depthwise convolutions are used to reduce1157 FLOPs, but the first layer is kept as a regular convolution so as not to degrade accuracy.1158 1159 `Striding` and `dw_striding` are the same except that the latter uses depthwise convolutions1160 after the first layer, whereas the former does not.1161 1162 Args:1163 subsampling_factor (int): Time reduction factor1164 feat_in (int): size of the input features1165 feat_out (int): size of the output features1166 subsampling (str): The subsampling technique, choose from1167 {"striding", "dw-striding", "striding_conv1d", "dw_striding_conv1d"}1168 conv_channels (int): Number of channels for the convolution layers, default is 256.1169 subsampling_conv_chunking_factor (int): Input chunking factor which can be -1 (no chunking)1170 1 (auto) or a power of 2. Default is 11171 activation (Module): activation function, default is nn.ReLU()1172 is_causal (bool): whether to use causal Conv1/2D, where each step will have limited access1173 to locations on its right or left1174 """1175 1176 def __init__(1177 self,1178 feat_in,1179 feat_out,1180 subsampling_factor=4,1181 subsampling="dw_striding",1182 conv_channels=256,1183 subsampling_conv_chunking_factor=1,1184 activation=nn.ReLU(),1185 is_causal=False,1186 ):1187 super().__init__()1188 self._subsampling = subsampling1189 self._conv_channels = conv_channels1190 self._feat_in = feat_in1191 self._feat_out = feat_out1192 1193 if subsampling_factor % 2 != 0:1194 raise ValueError("Sampling factor should be a multiply of 2!")1195 self._sampling_num = int(math.log(subsampling_factor, 2))1196 self.subsampling_factor = subsampling_factor1197 self.is_causal = is_causal1198 self.subsampling_causal_cond = subsampling in ("dw_striding", "striding", "striding_conv1d")1199 1200 if (