1# flake8: noqa: E266, C417, B950 2from dataclasses import dataclass 3from typing import Optional 4 5import torch 6import torch.nn as nn 7from torch import Tensor 8from torch.nn import functional as F 9 10 11def find_multiple(n: int, k: int) -> int: 12 if n % k == 0: 13 return n 14 return n + k - (n % k) 15 16 17@dataclass 18class ModelArgs: 19 block_size: int = 2048 20 vocab_size: int = 32000 21 n_layer: int = 32 22 n_head: int = 32 23 dim: int = 4096 24 intermediate_size: int = None 25 n_local_heads: int = -1 26 head_dim: int = 64 27 rope_base: float = 10000 28 norm_eps: float = 1e-5 29 30 def __post_init__(self): 31 if self.n_local_heads == -1: 32 self.n_local_heads = self.n_head 33 if self.intermediate_size is None: 34 hidden_dim = 4 * self.dim 35 n_hidden = int(2 * hidden_dim / 3) 36 self.intermediate_size = find_multiple(n_hidden, 256) 37 self.head_dim = self.dim // self.n_head 38 39 @classmethod 40 def from_name(cls, name: str): 41 if name in transformer_configs: 42 return cls(**transformer_configs[name]) 43 # fuzzy search 44 config = [ 45 config 46 for config in transformer_configs 47 if config in str(name).upper() or config in str(name) 48 ] 49 50 # We may have two or more configs matched (e.g. "7B" and "Mistral-7B"). Find the best config match, 51 # take longer name (as it have more symbols matched) 52 if len(config) > 1: 53 config.sort(key=len, reverse=True) 54 assert len(config[0]) != len( 55 config[1] 56 ), name # make sure only one 'best' match 57 58 return cls(**transformer_configs[config[0]]) 59 60 61transformer_configs = { 62 "CodeLlama-7b-Python-hf": dict( 63 block_size=16384, vocab_size=32000, n_layer=32, dim=4096, rope_base=1000000 64 ), 65 "7B": dict(n_layer=32, n_head=32, dim=4096), 66 "13B": dict(n_layer=40, n_head=40, dim=5120), 67 "30B": dict(n_layer=60, n_head=52, dim=6656), 68 "34B": dict( 69 n_layer=48, 70 n_head=64, 71 dim=8192, 72 vocab_size=32000, 73 n_local_heads=8, 74 intermediate_size=22016, 75 rope_base=1000000, 76 ), # CodeLlama-34B-Python-hf 77 "70B": dict( 78 n_layer=80, n_head=64, dim=8192, n_local_heads=8, intermediate_size=28672 79 ), 80 "Mistral-7B": dict( 81 n_layer=32, 82 n_head=32, 83 n_local_heads=8, 84 dim=4096, 85 intermediate_size=14336, 86 vocab_size=32000, 87 ), 88} 89 90 91class KVCache(nn.Module): 92 def __init__( 93 self, max_batch_size, max_seq_length, n_heads, head_dim, dtype=torch.bfloat16 94 ): 95 super().__init__() 96 cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim) 97 self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=dtype)) 98 self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=dtype)) 99 100 def update(self, input_pos, k_val, v_val): 101 # input_pos: [S], k_val: [B, H, S, D] 102 assert input_pos.shape[0] == k_val.shape[2] 103 104 k_out = self.k_cache 105 v_out = self.v_cache 106 k_out[:, :, input_pos] = k_val 107 v_out[:, :, input_pos] = v_val 108 109 return k_out, v_out 110 111 112class Transformer(nn.Module): 113 def __init__(self, config: ModelArgs) -> None: 114 super().__init__() 115 self.config = config 116 117 self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim) 118 self.layers = nn.ModuleList( 119 TransformerBlock(config) for _ in range(config.n_layer) 120 ) 121 self.norm = RMSNorm(config.dim, eps=config.norm_eps) 122 self.output = nn.Linear(config.dim, config.vocab_size, bias=False) 123 124 self.freqs_cis: Optional[Tensor] = None 125 self.mask_cache: Optional[Tensor] = None 126 self.max_batch_size = -1 127 self.max_seq_length = -1 128 129 def setup_caches(self, max_batch_size, max_seq_length): 130 if ( 131 self.max_seq_length >= max_seq_length 132 and self.max_batch_size >= max_batch_size 133 ): 134 return 135 head_dim = self.config.dim // self.config.n_head 136 max_seq_length = find_multiple(max_seq_length, 8) 137 self.max_seq_length = max_seq_length 138 self.max_batch_size = max_batch_size 139 for b in self.layers: 140 b.attention.kv_cache = KVCache( 141 max_batch_size, max_seq_length, self.config.n_local_heads, head_dim 142 ) 143 144 self.freqs_cis = precompute_freqs_cis( 145 self.config.block_size, 146 self.config.dim // self.config.n_head, 147 self.config.rope_base, 148 ) 149 self.causal_mask = torch.tril( 150 torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool) 151 ) 152 153 def forward(self, idx: Tensor, input_pos: Optional[Tensor] = None) -> Tensor: 154 assert self.freqs_cis is not None, "Caches must be initialized first" 155 mask = self.causal_mask[None, None, input_pos] 156 freqs_cis = self.freqs_cis[input_pos] 157 x = self.tok_embeddings(idx) 158 159 for i, layer in enumerate(self.layers): 160 x = layer(x, input_pos, freqs_cis, mask) 161 x = self.norm(x) 162 logits = self.output(x) 163 return logits 164 165 @classmethod 166 def from_name(cls, name: str): 167 return cls(ModelArgs.from_name(name)) 168 169 170class TransformerBlock(nn.Module): 171 def __init__(self, config: ModelArgs) -> None: 172 super().__init__() 173 self.attention = Attention(config) 174 self.feed_forward = FeedForward(config) 175 self.ffn_norm = RMSNorm(config.dim, config.norm_eps) 176 self.attention_norm = RMSNorm(config.dim, config.norm_eps) 177 178 def forward( 179 self, x: Tensor, input_pos: Tensor, freqs_cis: Tensor, mask: Tensor 180 ) -> Tensor: 181 h = x + self.attention(self.attention_norm(x), freqs_cis, mask, input_pos) 182 out = h + self.feed_forward(self.ffn_norm(h)) 183 return out 184 185 186class Attention(nn.Module): 187 def __init__(self, config: ModelArgs): 188 super().__init__() 189 assert config.dim % config.n_head == 0 190 191 total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim 192 # key, query, value projections for all heads, but in a batch 193 self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False) 194 self.wo = nn.Linear(config.dim, config.dim, bias=False) 195 self.kv_cache = None 196 197 self.n_head = config.n_head 198 self.head_dim = config.head_dim 199 self.n_local_heads = config.n_local_heads 200 self.dim = config.dim 201 self._register_load_state_dict_pre_hook(self.load_hook) 202 203 def load_hook(self, state_dict, prefix, *args): 204 if prefix + "wq.weight" in state_dict: 205 wq = state_dict.pop(prefix + "wq.weight") 206 wk = state_dict.pop(prefix + "wk.weight") 207 wv = state_dict.pop(prefix + "wv.weight") 208 state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv]) 209 210 def forward( 211 self, 212 x: Tensor, 213 freqs_cis: Tensor, 214 mask: Tensor, 215 input_pos: Optional[Tensor] = None, 216 ) -> Tensor: 217 bsz, seqlen, _ = x.shape 218 219 kv_size = self.n_local_heads * self.head_dim 220 q, k, v = self.wqkv(x).split([self.dim, kv_size, kv_size], dim=-1) 221 222 q = q.view(bsz, seqlen, self.n_head, self.head_dim) 223 k = k.view(bsz, seqlen, self.n_local_heads, self.head_dim) 224 v = v.view(bsz, seqlen, self.n_local_heads, self.head_dim) 225 226 q = apply_rotary_emb(q, freqs_cis) 227 k = apply_rotary_emb(k, freqs_cis) 228 229 q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) 230 231 if self.kv_cache is not None: 232 k, v = self.kv_cache.update(input_pos, k, v) 233 234 k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1) 235 v = v.repeat_interleave(self.n_head // self.n_local_heads, dim=1) 236 y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0) 237 238 y = y.transpose(1, 2).contiguous().view(bsz, seqlen, self.dim) 239 240 y = self.wo(y) 241 return y 242 243 244class FeedForward(nn.Module): 245 def __init__(self, config: ModelArgs) -> None: 246 super().__init__() 247 self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False) 248 self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False) 249 self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False) 250 251 def forward(self, x: Tensor) -> Tensor: 252 return self.w2(F.silu(self.w1(x)) * self.w3(x)) 253 254 255class RMSNorm(nn.Module): 256 def __init__(self, dim: int, eps: float = 1e-5): 257 super().__init__() 258 self.eps = eps 259 self.weight = nn.Parameter(torch.ones(dim)) 260 261 def _norm(self, x): 262 return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps) 263 264 def forward(self, x: Tensor) -> Tensor: 265 output = self._norm(x.float()).type_as(x) 266 return output * self.weight 267 268 269def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000) -> Tensor: 270 freqs = 1.0 / ( 271 base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem) 272 ) 273 t = torch.arange(seq_len, device=freqs.device) 274 freqs = torch.outer(t, freqs) 275 freqs_cis = torch.polar(torch.ones_like(freqs), freqs) 276 cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1) 277 return cache.to(dtype=torch.bfloat16) 278 279 280def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor: 281 xshaped = x.float().reshape(*x.shape[:-1], -1, 2) 282 freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2) 283 x_out2 = torch.stack( 284 [ 285 xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1], 286 xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1], 287 ], 288 -1, 289 ) 290 291 x_out2 = x_out2.flatten(3) 292 return x_out2.type_as(x) 293