OliverSundaram/MoE-Study-Applied
152
1import torch
2from torch import nn
3import torch.nn.functional as F
4
5from transformers import PretrainedConfig, PreTrainedModel, AutoConfig, AutoModelForCausalLM
6from transformers.modeling_outputs import CausalLMOutputWithPast
7
8from transformers import TokenizersBackend
9
10import time
11
12
13
14def build_rope_cache(head_dim: int, context_length: int, base: float = 10000.0, device=None):
15 assert head_dim % 2 == 0, "RoPE rotates 2D planes, so head_dim must be even"
16
17 plane_indices = torch.arange(head_dim // 2, dtype=torch.float32, device=device)
18 inverse_frequencies = base ** (-2.0 * plane_indices / head_dim)
19
20 positions = torch.arange(context_length, dtype=torch.float32, device=device)
21 angles = positions[:, None] * inverse_frequencies[None, :]
22
23 angles = torch.cat([angles, angles], dim=-1)
24
25 return angles.cos(), angles.sin()
26
27def rotate_half(x: torch.Tensor) -> torch.Tensor:
28 first_half, second_half = x.chunk(2, dim=-1)
29 return torch.cat([-second_half, first_half], dim=-1)
30
31def apply_rope(x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor, start_pos: int) -> torch.Tensor:
32 seq_len = x.shape[-2]
33 cos = rope_cos[start_pos:start_pos + seq_len].to(x.dtype)
34 sin = rope_sin[start_pos:start_pos + seq_len].to(x.dtype)
35 return x * cos + rotate_half(x) * sin
36
37def apply_rope_to_queries_and_keys(queries, keys, rope_cos, rope_sin, start_pos):
38 return apply_rope(queries, rope_cos, rope_sin, start_pos), apply_rope(keys, rope_cos, rope_sin, start_pos)
39
40
41
42class FeedForward(nn.Module):
43
44
45 def __init__(self, cfg):
46 super().__init__()
47 self.gate = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
48 self.up = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
49 self.down = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False)
50
51
52 def forward(self, x):
53 return self.down(F.silu(self.gate(x)) * self.up(x))
54
55
56
57class MoE(nn.Module):
58
59
60 def __init__(self, cfg: dict[str, int | bool]):
61 super().__init__()
62 self.n_experts = cfg["n_experts"]
63 self.top_k = cfg["top_k"]
64 self.experts = nn.ModuleList(
65 [FeedForward(cfg) for _ in range(self.n_experts)]
66 )
67 self.router = nn.Linear(cfg["emb_dim"], self.n_experts, bias=False)
68 self.aux_loss = 0.0
69
70
71 def forward(self, x: torch.Tensor):
72 batch_size, seq_len, emb_dim = x.shape
73 tokens = x.reshape(batch_size * seq_len, emb_dim)
74
75 router_logits = self.router(tokens)
76 router_probs = torch.softmax(router_logits, dim=-1)
77
78 top_weights, top_experts = torch.topk(router_probs, self.top_k, dim=-1)
79 top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)
80
81 expert_mask = F.one_hot(top_experts, self.n_experts)
82
83 tokens_per_expert = torch.sum(expert_mask, dim=1).float().mean(dim=0) / self.top_k
84 prob_per_expert = router_probs.mean(dim=0)
85
86 aux_loss = self.n_experts * torch.sum(tokens_per_expert * prob_per_expert, dim=0)
87
88 self.aux_loss = aux_loss
89
90 output = torch.zeros_like(tokens)
91
92 for expert_idx in range(self.n_experts):
93 token_pos, slot_pos = torch.where(top_experts == expert_idx)
94
95 selected_tokens = tokens[token_pos]
96 expert_output = self.experts[expert_idx](selected_tokens)
97
98 token_weights = top_weights[token_pos, slot_pos].unsqueeze(1)
99
100 output.index_add_(dim=0, index=token_pos, source=expert_output * token_weights)
101
102 return output.reshape(batch_size, seq_len, emb_dim)
103
104
105
106class MultiQueryAttention(nn.Module):
107
108
109 def __init__(self, cfg: dict[str, int | bool]):
110 super().__init__()
111
112 assert cfg["emb_dim"] % cfg["n_heads"] == 0
113
114 self.emb_dim = cfg["emb_dim"]
115 self.num_heads = cfg["n_heads"]
116 self.head_dim = self.emb_dim // self.num_heads
117 self.qkv_bias = cfg["qkv_bias"]
118 self.drop_rate = cfg["drop_rate"]
119
120 assert self.head_dim % 2 == 0, "RoPE requires an even head_dim"
121
122 self.query_proj = nn.Linear(self.emb_dim, self.emb_dim, self.qkv_bias)
123 self.key_proj = nn.Linear(self.emb_dim, self.head_dim, self.qkv_bias)
124 self.value_proj = nn.Linear(self.emb_dim, self.head_dim, self.qkv_bias)
125 self.out_proj = nn.Linear(self.emb_dim, self.emb_dim, bias=False)
126
127 self.register_buffer("k_cache", None, persistent=False)
128 self.register_buffer("v_cache", None, persistent=False)
129
130
131 def reset_cache(self):
132 self.k_cache, self.v_cache = None, None
133
134
135 def project_x(self, x, batch_size, seq_len, num_heads, head_dim):
136 return self.query_proj(x).reshape(batch_size, seq_len, num_heads, head_dim).transpose(1, 2), self.key_proj(x).unsqueeze(1), self.value_proj(x).unsqueeze(1)
137
138
139 def forward(self, x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor, use_cached: bool = False):
140 batch_size, seq_len, _ = x.shape
141 start_pos = self.k_cache.shape[-2] if self.k_cache is not None and use_cached else 0
142
143 queries, new_keys, new_values = self.project_x(x, batch_size, seq_len, self.num_heads, self.head_dim)
144 queries, new_keys = apply_rope_to_queries_and_keys(queries, new_keys, rope_cos, rope_sin, start_pos)
145
146 if use_cached:
147 if self.k_cache is None:
148 self.k_cache, self.v_cache = new_keys, new_values
149
150 else:
151 self.k_cache = torch.cat([self.k_cache, new_keys], dim=2)
152 self.v_cache = torch.cat([self.v_cache, new_values], dim=2)
153 keys, values = self.k_cache, self.v_cache
154 else:
155 keys = new_keys
156 values = new_values
157
158 keys = keys.expand(batch_size, self.num_heads, keys.shape[2], self.head_dim)
159 values = values.expand(batch_size, self.num_heads, keys.shape[2], self.head_dim)
160
161 context_vecs = F.scaled_dot_product_attention(queries, keys, values, dropout_p=self.drop_rate if self.training else 0.0, is_causal=(seq_len > 1))
162 context_vecs = context_vecs.transpose(1, 2).reshape(batch_size, seq_len, self.emb_dim)
163
164 return self.out_proj(context_vecs)
165
166
167
168class Transformer(nn.Module):
169
170
171 def __init__(self, cfg):
172 super().__init__()
173
174 self.attention = MultiQueryAttention(cfg)
175 self.ff = MoE(cfg)
176 self.norm1 = nn.RMSNorm(cfg["emb_dim"])
177 self.norm2 = nn.RMSNorm(cfg["emb_dim"])
178
179
180 def forward(self, x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor, use_cached: bool):
181
182 shortcut = x
183 x = self.norm1(x)
184 x = self.attention(x, rope_cos, rope_sin, use_cached)
185 x = x + shortcut
186
187 shortcut = x
188 x = self.norm2(x)
189 x = self.ff(x)
190 x = x + shortcut
191
192 return x
193
194
195 def reset_cache(self):
196 self.attention.reset_cache()
197
198
199
200class LLMConfig(PretrainedConfig):
201
202 model_type = "custom_llm"
203
204 def __init__(self,
205 vocab_size: int = 32768,
206 eos_token_id=0,
207 bos_token_id=0,
208 pad_token_id=0,
209 context_length: int = 1024,
210 emb_dim: int = 512,
211 hidden_dim: int = 1024,
212 n_heads: int = 8,
213 n_layers: int = 14,
214 qkv_bias: bool = False,
215 drop_rate: float = 0.0,
216 n_experts: int = 8,
217 top_k: int = 2,
218 rope_base: float = 10000.0,
219 **kwargs):
220
221 self.vocab_size = vocab_size
222 self.context_length = context_length
223 self.emb_dim = emb_dim
224 self.n_heads = n_heads
225 self.n_layers = n_layers
226 self.qkv_bias = qkv_bias
227 self.drop_rate = drop_rate
228 self.hidden_dim = hidden_dim
229 self.n_experts = n_experts
230 self.top_k = top_k
231 self.rope_base = rope_base
232
233 kwargs.setdefault("tie_word_embeddings", True)
234 super().__init__(
235 eos_token_id=eos_token_id,
236 bos_token_id=bos_token_id,
237 pad_token_id=pad_token_id,
238 **kwargs)
239
240
241 def __getitem__(self, key):
242 return getattr(self, key)
243
244
245
246class LLM(PreTrainedModel):
247
248 config_class = LLMConfig
249 _tied_weights_keys = {"out.weight": "tok_emb.weight"}
250 _input_embed_layer = "tok_emb"
251
252 def __init__(self, cfg: LLMConfig):
253 super().__init__(cfg)
254
255 self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
256 self.trans_blocks = nn.ModuleList(
257 [Transformer(cfg) for _ in range(cfg["n_layers"])]
258 )
259 rope_cos, rope_sin = build_rope_cache(
260 head_dim=cfg["emb_dim"] // cfg["n_heads"],
261 context_length=cfg["context_length"],
262 base=cfg["rope_base"],
263 )
264 self.register_buffer("rope_cos", rope_cos, persistent=True)
265 self.register_buffer("rope_sin", rope_sin, persistent=True)
266
267 self.norm = nn.RMSNorm(cfg["emb_dim"])
268 self.out = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
269 self.out.weight = self.tok_emb.weight
270
271 self.post_init()
272
273
274 def forward(self, input_ids: torch.Tensor, use_cached: bool = False, **kwargs) -> CausalLMOutputWithPast:
275 _, seq_len = input_ids.shape
276
277 cache = self.trans_blocks[0].attention.k_cache
278 cached_len = 0 if cache is None else cache.shape[-2]
279 assert cached_len + seq_len <= self.rope_cos.shape[0], (
280 f"position {cached_len + seq_len} exceeds context_length {self.rope_cos.shape[0]}"
281 )
282
283 x = self.tok_emb(input_ids)
284 for block in self.trans_blocks:
285 x = block(x, self.rope_cos, self.rope_sin, use_cached)
286
287 logits = self.out(self.norm(x))
288 return CausalLMOutputWithPast(logits=logits)
289
290
291 def reset_cache(self):
292 for t in self.trans_blocks:
293 t.reset_cache()
294
295
296 def get_output_embeddings(self):
297 return self.out
298
299
300 def set_output_embeddings(self, new_embeddings):
301 self.out = new_embeddings
302
303
304 @property
305 def aux_loss(self):
306 """
307 Return the summed aux loss across all MoE modules in LLM.
308 """
309 return sum(module.aux_loss for module in self.modules() if isinstance(module, MoE))
310
311
312 def format_prompt(self, prompt: str) -> str:
313 return "<|user|>" + prompt + "<|end|>" + "<|assistant|>"
314
315 @torch.inference_mode()
316 def generate(self, prompt: str, tokenizer: TokenizersBackend, device: torch.device, max_new_tokens: int, top_k: int = 40, temp: float | int = 1.3, use_cached: bool = True, print_text: bool = False):
317 """
318 Runs an inference loop with LLM model. Returns decoded text as a *string*. Starts by first passing the full encoded ids through the model, then only the final token.
319 This is to maintain shortcut usage, as well as building the KV cache. As well, returns TTFT (Time to First Token) and TPS (Tokens Per Second).
320 :param prompt: A string to run through the model.
321 :param tokenizer: A tokenizer to encode and decode the prompt.
322 :param device: Either "cpu" or "cuda".
323 :param max_new_tokens: Caps how many iterations the model predicts tokens.
324 :param top_k: Picks the *top_k* most likely tokens for addition.
325 :param temp: Divided the probs from softmax.
326 :param use_cached: Bool determining usage of KV cache.
327 :param print_text: Bool (*True or False*) determining whether to print newly generated tokens.
328 """
329
330
331 def get_next_id(ids: torch.Tensor, top_k: int, temp: float | int, use_cached: bool) -> torch.Tensor:
332 """
333 Run ids through LLM forward, generating logits. Use top_k & and temp to determine next token.
334 """
335 logits = self.forward(ids, use_cached)
336 logits = logits.logits.squeeze(0)
337
338 last_logits = logits[-1] / temp
339
340 top_logits, top_ids = torch.topk(last_logits, top_k, dim=-1)
341 probs = torch.softmax(top_logits, dim=-1)
342 top_id_index = torch.multinomial(probs, 1).squeeze()
343
344 return top_ids[top_id_index]
345
346
347 def get_last_id(ids: torch.Tensor) -> torch.Tensor:
348 """
349 Returns the last token id from ids, in shape *[B, 1]*.
350 """
351 return ids[:, -1:]
352
353
354 def join_next_id(ids: torch.Tensor, id: torch.Tensor) -> torch.Tensor:
355 """
356 Concatenates id to ids.
357 :param ids: Expected shape: *[B, seq_len]*
358 :param id: Expected shape: *single int*
359 """
360 return torch.cat(tensors=[ids, id.unsqueeze(0).unsqueeze(0)], dim=1)
361
362
363 time_to_first_tok_start = time.time()
364 sec_per_tok = []
365
366 self.reset_cache()
367
368 # Pass all ids through model first to generate KV cache, then only the last token
369 formatted_prompt = self.format_prompt(prompt)
370 ids = tokenizer(formatted_prompt, return_tensors="pt").input_ids.to(device)
371 next_id = get_next_id(ids=ids, top_k=top_k, temp=temp, use_cached=use_cached)
372
373 if print_text:
374 # Print decoded new token
375 text = tokenizer.decode(next_id.unsqueeze(0))
376 print(text, end="")
377
378 ids = join_next_id(ids=ids, id=next_id)
379
380 # Calculate time to first token
381 time_to_first_tok = time.time() - time_to_first_tok_start
382
383 for _ in range(max_new_tokens - 1):
384 start = time.time()
385
386 # Generate next token
387 model_input = get_last_id(ids=ids) if use_cached else ids
388 next_id = get_next_id(ids=model_input, top_k=top_k, temp=temp, use_cached=use_cached)
389
390 # Calculate seconds per token
391 end = time.time() - start
392 sec_per_tok.append(end)
393
394 if next_id == tokenizer.eos_token_id:
395 break
396
397 if print_text:
398 # Print decoded new token
399 text = tokenizer.decode(next_id.unsqueeze(0))
400 print(text, end="")
401
402 ids = join_next_id(ids=ids, id=next_id)
403
404 if len(ids[-1]) >= self.config.context_length:
405 break
406
407 # Calculate tokens per second from seconds per token
408 total_times = len(sec_per_tok)
409 avg_sec_per_tok = sum(sec_per_tok) / total_times
410 avg_tok_per_sec = 1 / avg_sec_per_tok
411
412 return tokenizer.decode(ids.squeeze(0)), avg_tok_per_sec, time_to_first_tok
413
414
415
416AutoConfig.register("custom_llm", LLMConfig)
417AutoModelForCausalLM.register(LLMConfig, LLM)
418LLMConfig.register_for_auto_class()
419LLM.register_for_auto_class("AutoModelForCausalLM")