CAMB-AI/MARS5-TTS
48076
1import torch2import torch.nn as nn3import torch.nn.functional as F4from torch import Tensor5import math6from dataclasses import dataclass7from typing import Optional8 9 10# --------------------------11# activation functions12 13class FNNSwiGLU(nn.Module):14 15 def __init__(self, dim, dim_ff) -> None:16 super().__init__()17 18 # we will receive in xW19 self.V = nn.Linear(dim, dim_ff, bias=False)20 self.W = nn.Linear(dim, dim_ff, bias=False)21 22 23 def forward(self, x: Tensor) -> Tensor:24 """ Compute SwiGLU output of x, the output of the first linear layer. i.e.25 FFNSwiGLU(x, W, V, W2) = (Swish1(xW) ⊗ xV )W2.26 NOTE: the transformer linear1 layer must be overwritten to identity. This layer only applies27 the Swish(xW) * xV. The W2 multiplication is done in the main transformer layer28 """29 return F.silu(self.W(x)) * self.V(x)30 31 32# ---------------------------------33# padding and position layers34 35class SinePositionalEmbedding(nn.Module):36 def __init__(37 self,38 dim_model: int,39 dropout: float = 0.0,40 scale: bool = False,41 alpha: bool = False,42 ):43 super().__init__()44 self.dim_model = dim_model45 self.x_scale = math.sqrt(dim_model) if scale else 1.046 self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)47 self.dropout = torch.nn.Dropout(p=dropout)48 49 self.reverse = False50 self.pe = None51 self.extend_pe(torch.tensor(0.0).expand(1, 4000))52 53 def extend_pe(self, x):54 """Reset the positional encodings."""55 if self.pe is not None:56 if self.pe.size(1) >= x.size(1):57 if self.pe.dtype != x.dtype or self.pe.device != x.device:58 self.pe = self.pe.to(dtype=x.dtype, device=x.device)59 return60 pe = torch.zeros(x.size(1), self.dim_model)61 if self.reverse:62 position = torch.arange(63 x.size(1) - 1, -1, -1.0, dtype=torch.float3264 ).unsqueeze(1)65 else:66 position = torch.arange(67 0, x.size(1), dtype=torch.float3268 ).unsqueeze(1)69 div_term = torch.exp(70 torch.arange(0, self.dim_model, 2, dtype=torch.float32)71 * -(math.log(10000.0) / self.dim_model)72 )73 pe[:, 0::2] = torch.sin(position * div_term)74 pe[:, 1::2] = torch.cos(position * div_term)75 pe = pe.unsqueeze(0)76 self.pe = pe.to(device=x.device, dtype=x.dtype).detach()77 78 def forward(self, x: torch.Tensor) -> torch.Tensor:79 """ Assumes x of shape (bs, seq_len, dim) """80 self.extend_pe(x)81 output = x.unsqueeze(-1) if x.ndim == 2 else x82 output = output * self.x_scale + self.alpha * self.pe[:, : x.size(1)]83 return self.dropout(output)84 85 86# --------------------------------87# kv cache blocks88 89class CacheView:90 def __init__(self, cache_k: torch.Tensor, cache_v: torch.Tensor):91 self.cache_k = cache_k92 self.cache_v = cache_v93 94 @property95 def sliding_window(self):96 return self.cache_k.shape[1]97 98class RotatingBufferCache:99 """100 This is an example that implements a less naive rotating buffer cache, allowing for variable length sequences.101 Allocated cache is rectangular which is wasteful (see PagedAttention for better mechanisms)102 """103 def __init__(self, n_layers: int, max_batch_size: int, sliding_window: int, n_kv_heads: int, head_dim: int):104 105 self.sliding_window = sliding_window106 self.n_kv_heads = n_kv_heads107 self.head_dim = head_dim108 109 self.cache_k = torch.empty((110 n_layers,111 max_batch_size,112 sliding_window,113 n_kv_heads,114 head_dim115 ))116 self.cache_v = torch.empty((117 n_layers,118 max_batch_size,119 sliding_window,120 n_kv_heads,121 head_dim122 ))123 124 def get_view(self, layer_id: int) -> CacheView:125 return CacheView(self.cache_k[layer_id], self.cache_v[layer_id])126 127 @property128 def device(self):129 return self.cache_k.device130 131 def to(self, device: torch.device, dtype: torch.dtype):132 self.cache_k = self.cache_k.to(device=device, dtype=dtype)133 self.cache_v = self.cache_v.to(device=device, dtype=dtype)134 return self135 136 137# --------------------------------138# Mistral transformer blocks139# Code for the follow blocks are adapted from 140# https://github.com/mistralai/mistral-src141# Thank you Mistral team!142 143@dataclass144class ModelArgs:145 vocab_size: int146 147 dim: int = 1152 # default for mars3 and before: 1024148 n_layers: int = 24149 head_dim: int = 64 # = dim/n_heads150 hidden_dim: int = 3584151 n_heads: int = 16152 n_kv_heads: int = 16 # default: 8153 sliding_window: int = 1792154 norm_eps: float = 1e-5155 156 max_batch_size: int = 256157 158 159def repeat_kv(keys: torch.Tensor, values: torch.Tensor, repeats: int):160 if repeats == 1: return keys, values161 keys = torch.repeat_interleave(keys, repeats=repeats, dim=2)162 values = torch.repeat_interleave(values, repeats=repeats, dim=2)163 return keys, values164 165 166def _reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:167 """168 freqs_cis: complex - (seq_len, head_dim / 2)169 x: complex - (bsz, seq_len, head_dim / 2)170 """171 ndim = x.ndim172 assert 1 < ndim173 assert freqs_cis.shape == (x.shape[1], x.shape[-1]), (174 freqs_cis.shape,175 (x.shape[1], x.shape[-1]),176 )177 shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]178 return freqs_cis.view(*shape)179 180 181def apply_rotary_emb(182 xq: torch.Tensor,183 xk: torch.Tensor,184 freqs_cis: torch.Tensor,185) -> tuple[torch.Tensor, torch.Tensor]:186 xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))187 xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))188 freqs_cis = _reshape_for_broadcast(freqs_cis, xq_)189 xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)190 xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)191 return xq_out.type_as(xq), xk_out.type_as(xk)192 193 194def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor:195 freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))196 t = torch.arange(end, device=freqs.device) # type: ignore197 freqs = torch.outer(t, freqs).float() # type: ignore198 return torch.polar(torch.ones_like(freqs), freqs) # complex64199 200 201class Attention(nn.Module):202 def __init__(self, args: ModelArgs):203 super().__init__()204 self.args = args205 206 self.n_heads: int = args.n_heads207 self.n_kv_heads: int = args.n_kv_heads208 209 self.repeats = self.n_heads // self.n_kv_heads210 self.sliding_window = self.args.sliding_window211 212 self.scale = self.args.head_dim**-0.5213 214 self.wq = nn.Linear(215 args.dim,216 args.n_heads * args.head_dim,217 bias=False218 )219 self.wk = nn.Linear(220 args.dim,221 args.n_kv_heads * args.head_dim,222 bias=False223 )224 self.wv = nn.Linear(225 args.dim,226 args.n_kv_heads * args.head_dim,227 bias=False228 )229 self.wo = nn.Linear(230 args.n_heads * args.head_dim,231 args.dim,232 bias=False233 )234 235 def forward(236 self, x: torch.Tensor, freqs_cis: torch.Tensor, positions: torch.Tensor, mask: Optional[torch.Tensor], cache: Optional[CacheView]237 ) -> torch.Tensor:238 239 bsz, seqlen, _ = x.shape240 241 xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)242 xq = xq.view(bsz, seqlen, self.n_heads, self.args.head_dim)243 xk = xk.view(bsz, seqlen, self.n_kv_heads, self.args.head_dim)244 xv = xv.view(bsz, seqlen, self.n_kv_heads, self.args.head_dim)245 xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)246 247 # The cache is a rotating buffer248 if cache is not None:249 scatter_pos = (positions[-self.sliding_window:] % self.sliding_window)[None, :, None, None]250 scatter_pos = scatter_pos.repeat(bsz, 1, self.n_kv_heads, self.args.head_dim)251 cache.cache_k[:bsz].scatter_(dim=1, index=scatter_pos, src=xk[:, -self.sliding_window:])252 cache.cache_v[:bsz].scatter_(dim=1, index=scatter_pos, src=xv[:, -self.sliding_window:])253 254 if positions.shape[0] > 1:255 # prefill256 key, value = repeat_kv(xk, xv, self.repeats)257 else:258 cur_pos = positions[-1].item() + 1259 key, value = repeat_kv(cache.cache_k[:bsz, :cur_pos, ...], cache.cache_v[:bsz, :cur_pos, ...], self.repeats)260 261 # print(f"Internal: {xq.shape}, key: {key.shape}, mask: {mask.shape} | {mask.dtype} | xq: {xq.dtype} | mask: {mask} ")262 # if mask is not None: 263 # mask = mask[None, None, ...].expand(bsz, self.n_heads, -1, -1)264 # mask = mask.to(key.dtype)265 266 query = xq.transpose(1, 2)267 key = key.transpose(1, 2)268 value = value.transpose(1, 2)269 # # scores : [bsz, n_heads, seqlen | 1, seqlen]270 # scores = torch.matmul(query, key.transpose(2, 3)) * self.scale271 272 output = F.scaled_dot_product_attention(query, key, value, mask) # (bs, n_local_heads, slen, head_dim)273 output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)274 return self.wo(output)275 276 277class FeedForward(nn.Module):278 def __init__(self, args: ModelArgs):279 super().__init__()280 281 self.w1 = nn.Linear(282 args.dim,283 args.hidden_dim,284 bias=False285 )286 self.w2 = nn.Linear(287 args.hidden_dim,288 args.dim,289 bias=False290 )291 self.w3 = nn.Linear(292 args.dim,293 args.hidden_dim,294 bias=False295 )296 297 def forward(self, x) -> torch.Tensor:298 return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))299 300 301class RMSNorm(torch.nn.Module):302 def __init__(self, dim: int, eps: float = 1e-6):303 super().__init__()304 self.eps = eps305 self.weight = nn.Parameter(torch.ones(dim))306 307 def _norm(self, x):308 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)309 310 def forward(self, x):311 output = self._norm(x.float()).type_as(x)312 return output * self.weight313 314 315class TransformerBlock(nn.Module):316 def __init__(self, args: ModelArgs):317 super().__init__()318 self.n_heads = args.n_heads319 self.dim = args.dim320 self.attention = Attention(args)321 self.feed_forward = FeedForward(args=args)322 self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)323 self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)324 self.args = args325 326 def forward(327 self, x: torch.Tensor, freqs_cis: torch.Tensor, positions: torch.Tensor, mask: Optional[torch.Tensor], cache: Optional[CacheView]328 ) -> torch.Tensor:329 r = self.attention.forward(self.attention_norm(x), freqs_cis, positions, mask, cache)330 h = x + r331 r = self.feed_forward.forward(self.ffn_norm(h))332 out = h + r333 return out334 335 336class MistralTransformer(nn.Module):337 def __init__(self, args: ModelArgs):338 super().__init__()339 self.args = args340 self.vocab_size = args.vocab_size341 self.n_layers = args.n_layers342 assert self.vocab_size > 0343 344 # self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim)345 346 self.layers = torch.nn.ModuleList(347 [TransformerBlock(args=args) for _ in range(args.n_layers)]348 )349 350 self.norm = RMSNorm(args.dim, eps=args.norm_eps)351 352 self.output = nn.Linear(353 args.dim,354 args.vocab_size,355 bias=False356 )357 358 # self.freqs_cis 359 self.freqs_cis = precompute_freqs_cis(self.args.head_dim, 128_000)360 361 @property362 def dtype(self) -> torch.dtype:363 return self.tok_embeddings.weight.dtype364 365 @property366 def device(self) -> torch.device:367 return self.tok_embeddings.weight.device368 369 def forward(370 self,371 input_ids: torch.Tensor,372 positions: torch.Tensor,373 cache: Optional[RotatingBufferCache]374 ):375 h = input_ids376 if self.freqs_cis.device != h.device:377 self.freqs_cis = self.freqs_cis.to(h.device)378 freqs_cis = self.freqs_cis[positions]379 380 mask: Optional[torch.Tensor] = None381 if input_ids.shape[1] > 1:382 seqlen = input_ids.shape[1]383 tensor = torch.full(384 (seqlen, seqlen),385 dtype=h.dtype,386 fill_value=1,387 device=h.device,388 )389 mask = torch.tril(tensor, diagonal=0).to(h.dtype)390 # make the mask banded to account for sliding window391 mask = torch.triu(mask, diagonal=-self.args.sliding_window)392 mask = torch.log(mask)393 394 for layer_id, layer in enumerate(self.layers):395 cache_view = None if cache is None else cache.get_view(layer_id)396 h = layer(h, freqs_cis, positions, mask, cache_view)397 398 return self.output(self.norm(h))399 400 401