| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Conformer encoder components for FireRedASR2. |
| 4 | |
| 5 | The audio encoder architecture is |
| 6 | (Conv2dSubsampling → RelPositionalEncoding → N × RelPosEmbConformerBlock). |
| 7 | """ |
| 8 | |
| 9 | import torch |
| 10 | import torch.nn.functional as F |
| 11 | from torch import nn |
| 12 | |
| 13 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 14 | |
| 15 | |
| 16 | class Conv2dSubsampling(nn.Module): |
| 17 | def __init__(self, idim: int, d_model: int, out_channels: int = 32): |
| 18 | super().__init__() |
| 19 | self.conv = nn.Sequential( |
| 20 | nn.Conv2d(1, out_channels, 3, 2), |
| 21 | nn.ReLU(), |
| 22 | nn.Conv2d(out_channels, out_channels, 3, 2), |
| 23 | nn.ReLU(), |
| 24 | ) |
| 25 | subsample_idim = ((idim - 1) // 2 - 1) // 2 |
| 26 | self.out = ReplicatedLinear( |
| 27 | input_size=out_channels * subsample_idim, |
| 28 | output_size=d_model, |
| 29 | bias=True, |
| 30 | ) |
| 31 | |
| 32 | self.subsampling = 4 |
| 33 | left_context = right_context = 3 # both exclude current frame |
| 34 | self.context = left_context + 1 + right_context # 7 |
| 35 | |
| 36 | def forward( |
| 37 | self, x: torch.Tensor, x_mask: torch.Tensor |
| 38 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 39 | x = x.unsqueeze(1) |
| 40 | x = self.conv(x) |
| 41 | N, C, T, D = x.size() |
| 42 | x, _ = self.out(x.transpose(1, 2).contiguous().view(N, T, C * D)) |
| 43 | mask = x_mask[:, :, :-2:2][:, :, :-2:2] |
| 44 | input_lengths = mask[:, -1, :].sum(dim=-1) |
| 45 | return x, input_lengths, mask |
| 46 | |
| 47 | |
| 48 | class Swish(nn.Module): |
| 49 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 50 | return x * torch.sigmoid(x) |
| 51 | |
| 52 | |
| 53 | class RelPositionalEncoding(nn.Module): |
| 54 | def __init__(self, d_model: int, max_len: int = 5000): |
| 55 | super().__init__() |
| 56 | pe_positive = torch.zeros(max_len, d_model, requires_grad=False) |
| 57 | pe_negative = torch.zeros(max_len, d_model, requires_grad=False) |
| 58 | position = torch.arange(0, max_len).unsqueeze(1).float() |
| 59 | div_term = torch.exp( |
| 60 | torch.arange(0, d_model, 2).float() |
| 61 | * -(torch.log(torch.tensor(10000.0)).item() / d_model) |
| 62 | ) |
| 63 | pe_positive[:, 0::2] = torch.sin(position * div_term) |
| 64 | pe_positive[:, 1::2] = torch.cos(position * div_term) |
| 65 | pe_negative[:, 0::2] = torch.sin(-1 * position * div_term) |
| 66 | pe_negative[:, 1::2] = torch.cos(-1 * position * div_term) |
| 67 | |
| 68 | pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0) |
| 69 | pe_negative = pe_negative[1:].unsqueeze(0) |
| 70 | self.register_buffer( |
| 71 | "pe", torch.cat([pe_positive, pe_negative], dim=1), persistent=False |
| 72 | ) |
| 73 | |
| 74 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 75 | # Tmax = 2 * max_len - 1 |
| 76 | Tmax, T = self.pe.size(1), x.size(1) |
| 77 | pos_emb = self.pe[:, Tmax // 2 - T + 1 : Tmax // 2 + T].clone().detach() |
| 78 | return pos_emb |
| 79 | |
| 80 | |
| 81 | class ConformerFeedForward(nn.Module): |
| 82 | def __init__(self, d_model: int): |
| 83 | super().__init__() |
| 84 | self.pre_layer_norm = nn.LayerNorm(d_model) |
| 85 | self.linear_expand = ReplicatedLinear( |
| 86 | input_size=d_model, |
| 87 | output_size=d_model * 4, |
| 88 | bias=True, |
| 89 | ) |
| 90 | self.nonlinear = Swish() |
| 91 | self.linear_project = ReplicatedLinear( |
| 92 | input_size=d_model * 4, |
| 93 | output_size=d_model, |
| 94 | bias=True, |
| 95 | ) |
| 96 | |
| 97 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 98 | residual = x |
| 99 | x = self.pre_layer_norm(x) |
| 100 | x, _ = self.linear_expand(x) |
| 101 | x = self.nonlinear(x) |
| 102 | x, _ = self.linear_project(x) |
| 103 | return x + residual |
| 104 | |
| 105 | |
| 106 | class EncoderMultiHeadAttention(nn.Module): |
| 107 | def __init__(self, n_head: int, d_model: int): |
| 108 | super().__init__() |
| 109 | assert d_model % n_head == 0 |
| 110 | self.n_head = n_head |
| 111 | self.d_k = d_model // n_head |
| 112 | self.d_v = self.d_k |
| 113 | |
| 114 | self.w_qs = ReplicatedLinear(d_model, n_head * self.d_k, bias=False) |
| 115 | self.w_ks = ReplicatedLinear(d_model, n_head * self.d_k, bias=False) |
| 116 | self.w_vs = ReplicatedLinear(d_model, n_head * self.d_v, bias=False) |
| 117 | |
| 118 | self.layer_norm_q = nn.LayerNorm(d_model) |
| 119 | self.layer_norm_k = nn.LayerNorm(d_model) |
| 120 | self.layer_norm_v = nn.LayerNorm(d_model) |
| 121 | |
| 122 | self.fc = ReplicatedLinear(n_head * self.d_v, d_model, bias=False) |
| 123 | |
| 124 | def forward_qkv( |
| 125 | self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor |
| 126 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 127 | d_k, d_v, n_head = self.d_k, self.d_v, self.n_head |
| 128 | sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1) |
| 129 | |
| 130 | q = self.layer_norm_q(q) |
| 131 | k = self.layer_norm_k(k) |
| 132 | v = self.layer_norm_v(v) |
| 133 | |
| 134 | q = self.w_qs(q)[0].view(sz_b, len_q, n_head, d_k) |
| 135 | k = self.w_ks(k)[0].view(sz_b, len_k, n_head, d_k) |
| 136 | v = self.w_vs(v)[0].view(sz_b, len_v, n_head, d_v) |
| 137 | q = q.transpose(1, 2) |
| 138 | k = k.transpose(1, 2) |
| 139 | v = v.transpose(1, 2) |
| 140 | return q, k, v |
| 141 | |
| 142 | def forward_output( |
| 143 | self, |
| 144 | output: torch.Tensor, |
| 145 | residual: torch.Tensor, |
| 146 | sz_b: int, |
| 147 | len_q: int, |
| 148 | ) -> torch.Tensor: |
| 149 | output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1) |
| 150 | fc_out, _ = self.fc(output) |
| 151 | return fc_out + residual |
| 152 | |
| 153 | def forward_attention( |
| 154 | self, |
| 155 | attn: torch.Tensor, |
| 156 | v: torch.Tensor, |
| 157 | mask: torch.Tensor | None = None, |
| 158 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 159 | if mask is not None: |
| 160 | mask = mask.unsqueeze(1) |
| 161 | mask = mask.eq(0) |
| 162 | attn = attn.masked_fill(mask, -float("inf")) |
| 163 | attn = torch.softmax(attn, dim=-1).masked_fill(mask, 0.0) |
| 164 | else: |
| 165 | attn = torch.softmax(attn, dim=-1) |
| 166 | output = torch.matmul(attn, v) |
| 167 | return output, attn |
| 168 | |
| 169 | |
| 170 | class RelPosMultiHeadAttention(EncoderMultiHeadAttention): |
| 171 | def __init__(self, n_head: int, d_model: int): |
| 172 | super().__init__(n_head, d_model) |
| 173 | d_k = d_model // n_head |
| 174 | self.scale = 1.0 / (d_k**0.5) |
| 175 | self.linear_pos = ReplicatedLinear(d_model, n_head * d_k, bias=False) |
| 176 | self.pos_bias_u = nn.Parameter(torch.empty([n_head, d_k])) |
| 177 | self.pos_bias_v = nn.Parameter(torch.empty([n_head, d_k])) |
| 178 | |
| 179 | def _rel_shift(self, x): |
| 180 | N, H, T1, T2 = x.size() |
| 181 | zero_pad = torch.zeros((N, H, T1, 1), device=x.device, dtype=x.dtype) |
| 182 | x_padded = torch.cat([zero_pad, x], dim=-1) |
| 183 | x_padded = x_padded.view(N, H, T2 + 1, T1) |
| 184 | x = x_padded[:, :, 1:].view_as(x) |
| 185 | x = x[:, :, :, : x.size(-1) // 2 + 1] |
| 186 | return x |
| 187 | |
| 188 | def forward( |
| 189 | self, |
| 190 | q: torch.Tensor, |
| 191 | k: torch.Tensor, |
| 192 | v: torch.Tensor, |
| 193 | pos_emb: torch.Tensor, |
| 194 | mask: torch.Tensor | None = None, |
| 195 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 196 | sz_b, len_q = q.size(0), q.size(1) |
| 197 | residual = q |
| 198 | q, k, v = self.forward_qkv(q, k, v) |
| 199 | |
| 200 | q = q.transpose(1, 2) |
| 201 | n_batch_pos = pos_emb.size(0) |
| 202 | p = self.linear_pos(pos_emb)[0].view(n_batch_pos, -1, self.n_head, self.d_k) |
| 203 | p = p.transpose(1, 2) |
| 204 | |
| 205 | q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) |
| 206 | q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) |
| 207 | |
| 208 | matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) |
| 209 | matrix_bd = self._rel_shift(matrix_bd) |
| 210 | |
| 211 | batch_size, num_heads, query_len, key_len = matrix_bd.shape |
| 212 | attn_scores = torch.baddbmm( |
| 213 | matrix_bd.reshape(-1, query_len, key_len), |
| 214 | q_with_bias_u.reshape(-1, query_len, self.d_k), |
| 215 | k.transpose(-2, -1).reshape(-1, self.d_k, key_len), |
| 216 | beta=self.scale, |
| 217 | alpha=self.scale, |
| 218 | ).view( |
| 219 | batch_size, |
| 220 | num_heads, |
| 221 | query_len, |
| 222 | key_len, |
| 223 | ) |
| 224 | |
| 225 | output, attn = self.forward_attention(attn_scores, v, mask=mask) |
| 226 | output = self.forward_output(output, residual, sz_b, len_q) |
| 227 | return output, attn |
| 228 | |
| 229 | |
| 230 | class ConformerConvolution(nn.Module): |
| 231 | def __init__(self, d_model: int, kernel_size: int = 33): |
| 232 | super().__init__() |
| 233 | assert kernel_size % 2 == 1 |
| 234 | self.pre_layer_norm = nn.LayerNorm(d_model) |
| 235 | self.pointwise_conv1 = nn.Conv1d( |
| 236 | d_model, d_model * 4, kernel_size=1, bias=False |
| 237 | ) |
| 238 | self.padding = (kernel_size - 1) // 2 |
| 239 | self.depthwise_conv = nn.Conv1d( |
| 240 | d_model * 2, |
| 241 | d_model * 2, |
| 242 | kernel_size, |
| 243 | stride=1, |
| 244 | padding=self.padding, |
| 245 | groups=d_model * 2, |
| 246 | bias=False, |
| 247 | ) |
| 248 | self.batch_norm = nn.LayerNorm(d_model * 2) |
| 249 | self.swish = Swish() |
| 250 | self.pointwise_conv2 = nn.Conv1d( |
| 251 | d_model * 2, d_model, kernel_size=1, bias=False |
| 252 | ) |
| 253 | |
| 254 | def forward( |
| 255 | self, x: torch.Tensor, mask: torch.Tensor | None = None |
| 256 | ) -> torch.Tensor: |
| 257 | residual = x |
| 258 | out = self.pre_layer_norm(x) |
| 259 | out = out.transpose(1, 2) |
| 260 | if mask is not None: |
| 261 | out.masked_fill_(mask.ne(1), 0.0) |
| 262 | out = self.pointwise_conv1(out) |
| 263 | out = F.glu(out, dim=1) |
| 264 | out = self.depthwise_conv(out) |
| 265 | out = out.transpose(1, 2) |
| 266 | out = self.swish(self.batch_norm(out)) |
| 267 | out = out.transpose(1, 2) |
| 268 | out = self.pointwise_conv2(out) |
| 269 | if mask is not None: |
| 270 | out.masked_fill_(mask.ne(1), 0.0) |
| 271 | out = out.transpose(1, 2) |
| 272 | return out + residual |
| 273 | |
| 274 | |
| 275 | class RelPosEmbConformerBlock(nn.Module): |
| 276 | def __init__(self, d_model: int, n_head: int, kernel_size: int = 33): |
| 277 | super().__init__() |
| 278 | self.ffn1 = ConformerFeedForward(d_model) |
| 279 | self.mhsa = RelPosMultiHeadAttention(n_head, d_model) |
| 280 | self.conv = ConformerConvolution(d_model, kernel_size) |
| 281 | self.ffn2 = ConformerFeedForward(d_model) |
| 282 | self.layer_norm = nn.LayerNorm(d_model) |
| 283 | |
| 284 | def forward( |
| 285 | self, |
| 286 | x: torch.Tensor, |
| 287 | pos_emb: torch.Tensor, |
| 288 | slf_attn_mask: torch.Tensor | None = None, |
| 289 | pad_mask: torch.Tensor | None = None, |
| 290 | ) -> torch.Tensor: |
| 291 | out = 0.5 * x + 0.5 * self.ffn1(x) |
| 292 | out = self.mhsa(out, out, out, pos_emb, mask=slf_attn_mask)[0] |
| 293 | out = self.conv(out, pad_mask) |
| 294 | out = 0.5 * out + 0.5 * self.ffn2(out) |
| 295 | out = self.layer_norm(out) |
| 296 | return out |
| 297 | |
| 298 | |
| 299 | class ConformerEncoder(nn.Module): |
| 300 | """Conformer encoder used by FireRedASR2.""" |
| 301 | |
| 302 | def __init__( |
| 303 | self, |
| 304 | idim: int, |
| 305 | n_layers_enc: int, |
| 306 | n_head: int, |
| 307 | d_model: int, |
| 308 | kernel_size: int = 33, |
| 309 | pe_maxlen: int = 5000, |
| 310 | ): |
| 311 | super().__init__() |
| 312 | self.odim = d_model |
| 313 | |
| 314 | self.input_preprocessor = Conv2dSubsampling(idim, d_model) |
| 315 | self.positional_encoding = RelPositionalEncoding(d_model, max_len=pe_maxlen) |
| 316 | |
| 317 | self.layer_stack = nn.ModuleList() |
| 318 | for _ in range(n_layers_enc): |
| 319 | block = RelPosEmbConformerBlock(d_model, n_head, kernel_size) |
| 320 | self.layer_stack.append(block) |
| 321 | |
| 322 | def forward( |
| 323 | self, |
| 324 | padded_input: torch.Tensor, |
| 325 | input_lengths: torch.Tensor, |
| 326 | pad: bool = True, |
| 327 | ): |
| 328 | if pad: |
| 329 | padded_input = F.pad( |
| 330 | padded_input, |
| 331 | (0, 0, 0, self.input_preprocessor.context - 1), |
| 332 | "constant", |
| 333 | 0.0, |
| 334 | ) |
| 335 | src_mask = self.padding_position_is_0(padded_input, input_lengths) |
| 336 | |
| 337 | embed_output, input_lengths, src_mask = self.input_preprocessor( |
| 338 | padded_input, src_mask |
| 339 | ) |
| 340 | enc_output = embed_output |
| 341 | |
| 342 | pos_emb = self.positional_encoding(embed_output) |
| 343 | |
| 344 | for enc_layer in self.layer_stack: |
| 345 | enc_output = enc_layer( |
| 346 | enc_output, pos_emb, slf_attn_mask=src_mask, pad_mask=src_mask |
| 347 | ) |
| 348 | |
| 349 | return enc_output, input_lengths, src_mask |
| 350 | |
| 351 | def padding_position_is_0( |
| 352 | self, padded_input: torch.Tensor, input_lengths: torch.Tensor |
| 353 | ) -> torch.Tensor: |
| 354 | N, T = padded_input.size()[:2] |
| 355 | # Use broadcasting instead of a Python loop for efficiency. |
| 356 | positions = torch.arange(T, device=padded_input.device).unsqueeze(0) |
| 357 | mask = (positions < input_lengths.unsqueeze(1)).to(torch.uint8) |
| 358 | return mask.unsqueeze(1) |