import torch from torch import nn import torch.nn.functional as F from transformers import PretrainedConfig, PreTrainedModel, AutoConfig, AutoModelForCausalLM from transformers.modeling_outputs import CausalLMOutputWithPast from transformers import TokenizersBackend import time def build_rope_cache(head_dim: int, context_length: int, base: float = 10000.0, device=None): assert head_dim % 2 == 0, "RoPE rotates 2D planes, so head_dim must be even" plane_indices = torch.arange(head_dim // 2, dtype=torch.float32, device=device) inverse_frequencies = base ** (-2.0 * plane_indices / head_dim) positions = torch.arange(context_length, dtype=torch.float32, device=device) angles = positions[:, None] * inverse_frequencies[None, :] angles = torch.cat([angles, angles], dim=-1) return angles.cos(), angles.sin() def rotate_half(x: torch.Tensor) -> torch.Tensor: first_half, second_half = x.chunk(2, dim=-1) return torch.cat([-second_half, first_half], dim=-1) def apply_rope(x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor, start_pos: int) -> torch.Tensor: seq_len = x.shape[-2] cos = rope_cos[start_pos:start_pos + seq_len].to(x.dtype) sin = rope_sin[start_pos:start_pos + seq_len].to(x.dtype) return x * cos + rotate_half(x) * sin def apply_rope_to_queries_and_keys(queries, keys, rope_cos, rope_sin, start_pos): return apply_rope(queries, rope_cos, rope_sin, start_pos), apply_rope(keys, rope_cos, rope_sin, start_pos) class FeedForward(nn.Module): def __init__(self, cfg): super().__init__() self.gate = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False) self.up = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False) self.down = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False) def forward(self, x): return self.down(F.silu(self.gate(x)) * self.up(x)) class MoE(nn.Module): def __init__(self, cfg: dict[str, int | bool]): super().__init__() self.n_experts = cfg["n_experts"] self.top_k = cfg["top_k"] self.experts = nn.ModuleList( [FeedForward(cfg) for _ in range(self.n_experts)] ) self.router = nn.Linear(cfg["emb_dim"], self.n_experts, bias=False) self.aux_loss = 0.0 def forward(self, x: torch.Tensor): batch_size, seq_len, emb_dim = x.shape tokens = x.reshape(batch_size * seq_len, emb_dim) router_logits = self.router(tokens) router_probs = torch.softmax(router_logits, dim=-1) top_weights, top_experts = torch.topk(router_probs, self.top_k, dim=-1) top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True) expert_mask = F.one_hot(top_experts, self.n_experts) tokens_per_expert = torch.sum(expert_mask, dim=1).float().mean(dim=0) / self.top_k prob_per_expert = router_probs.mean(dim=0) aux_loss = self.n_experts * torch.sum(tokens_per_expert * prob_per_expert, dim=0) self.aux_loss = aux_loss output = torch.zeros_like(tokens) for expert_idx in range(self.n_experts): token_pos, slot_pos = torch.where(top_experts == expert_idx) selected_tokens = tokens[token_pos] expert_output = self.experts[expert_idx](selected_tokens) token_weights = top_weights[token_pos, slot_pos].unsqueeze(1) output.index_add_(dim=0, index=token_pos, source=expert_output * token_weights) return output.reshape(batch_size, seq_len, emb_dim) class MultiQueryAttention(nn.Module): def __init__(self, cfg: dict[str, int | bool]): super().__init__() assert cfg["emb_dim"] % cfg["n_heads"] == 0 self.emb_dim = cfg["emb_dim"] self.num_heads = cfg["n_heads"] self.head_dim = self.emb_dim // self.num_heads self.qkv_bias = cfg["qkv_bias"] self.drop_rate = cfg["drop_rate"] assert self.head_dim % 2 == 0, "RoPE requires an even head_dim" self.query_proj = nn.Linear(self.emb_dim, self.emb_dim, self.qkv_bias) self.key_proj = nn.Linear(self.emb_dim, self.head_dim, self.qkv_bias) self.value_proj = nn.Linear(self.emb_dim, self.head_dim, self.qkv_bias) self.out_proj = nn.Linear(self.emb_dim, self.emb_dim, bias=False) self.register_buffer("k_cache", None, persistent=False) self.register_buffer("v_cache", None, persistent=False) def reset_cache(self): self.k_cache, self.v_cache = None, None def project_x(self, x, batch_size, seq_len, num_heads, head_dim): 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) def forward(self, x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor, use_cached: bool = False): batch_size, seq_len, _ = x.shape start_pos = self.k_cache.shape[-2] if self.k_cache is not None and use_cached else 0 queries, new_keys, new_values = self.project_x(x, batch_size, seq_len, self.num_heads, self.head_dim) queries, new_keys = apply_rope_to_queries_and_keys(queries, new_keys, rope_cos, rope_sin, start_pos) if use_cached: if self.k_cache is None: self.k_cache, self.v_cache = new_keys, new_values else: self.k_cache = torch.cat([self.k_cache, new_keys], dim=2) self.v_cache = torch.cat([self.v_cache, new_values], dim=2) keys, values = self.k_cache, self.v_cache else: keys = new_keys values = new_values keys = keys.expand(batch_size, self.num_heads, keys.shape[2], self.head_dim) values = values.expand(batch_size, self.num_heads, keys.shape[2], self.head_dim) 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)) context_vecs = context_vecs.transpose(1, 2).reshape(batch_size, seq_len, self.emb_dim) return self.out_proj(context_vecs) class Transformer(nn.Module): def __init__(self, cfg): super().__init__() self.attention = MultiQueryAttention(cfg) self.ff = MoE(cfg) self.norm1 = nn.RMSNorm(cfg["emb_dim"]) self.norm2 = nn.RMSNorm(cfg["emb_dim"]) def forward(self, x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor, use_cached: bool): shortcut = x x = self.norm1(x) x = self.attention(x, rope_cos, rope_sin, use_cached) x = x + shortcut shortcut = x x = self.norm2(x) x = self.ff(x) x = x + shortcut return x def reset_cache(self): self.attention.reset_cache() class LLMConfig(PretrainedConfig): model_type = "custom_llm" def __init__(self, vocab_size: int = 32768, eos_token_id=0, bos_token_id=0, pad_token_id=0, context_length: int = 1024, emb_dim: int = 512, hidden_dim: int = 1024, n_heads: int = 8, n_layers: int = 14, qkv_bias: bool = False, drop_rate: float = 0.0, n_experts: int = 8, top_k: int = 2, rope_base: float = 10000.0, **kwargs): self.vocab_size = vocab_size self.context_length = context_length self.emb_dim = emb_dim self.n_heads = n_heads self.n_layers = n_layers self.qkv_bias = qkv_bias self.drop_rate = drop_rate self.hidden_dim = hidden_dim self.n_experts = n_experts self.top_k = top_k self.rope_base = rope_base kwargs.setdefault("tie_word_embeddings", True) super().__init__( eos_token_id=eos_token_id, bos_token_id=bos_token_id, pad_token_id=pad_token_id, **kwargs) def __getitem__(self, key): return getattr(self, key) class LLM(PreTrainedModel): config_class = LLMConfig _tied_weights_keys = {"out.weight": "tok_emb.weight"} _input_embed_layer = "tok_emb" def __init__(self, cfg: LLMConfig): super().__init__(cfg) self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"]) self.trans_blocks = nn.ModuleList( [Transformer(cfg) for _ in range(cfg["n_layers"])] ) rope_cos, rope_sin = build_rope_cache( head_dim=cfg["emb_dim"] // cfg["n_heads"], context_length=cfg["context_length"], base=cfg["rope_base"], ) self.register_buffer("rope_cos", rope_cos, persistent=True) self.register_buffer("rope_sin", rope_sin, persistent=True) self.norm = nn.RMSNorm(cfg["emb_dim"]) self.out = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False) self.out.weight = self.tok_emb.weight self.post_init() def forward(self, input_ids: torch.Tensor, use_cached: bool = False, **kwargs) -> CausalLMOutputWithPast: _, seq_len = input_ids.shape cache = self.trans_blocks[0].attention.k_cache cached_len = 0 if cache is None else cache.shape[-2] assert cached_len + seq_len <= self.rope_cos.shape[0], ( f"position {cached_len + seq_len} exceeds context_length {self.rope_cos.shape[0]}" ) x = self.tok_emb(input_ids) for block in self.trans_blocks: x = block(x, self.rope_cos, self.rope_sin, use_cached) logits = self.out(self.norm(x)) return CausalLMOutputWithPast(logits=logits) def reset_cache(self): for t in self.trans_blocks: t.reset_cache() def get_output_embeddings(self): return self.out def set_output_embeddings(self, new_embeddings): self.out = new_embeddings @property def aux_loss(self): """ Return the summed aux loss across all MoE modules in LLM. """ return sum(module.aux_loss for module in self.modules() if isinstance(module, MoE)) def format_prompt(self, prompt: str) -> str: return "<|user|>" + prompt + "<|end|>" + "<|assistant|>" @torch.inference_mode() 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): """ 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. 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). :param prompt: A string to run through the model. :param tokenizer: A tokenizer to encode and decode the prompt. :param device: Either "cpu" or "cuda". :param max_new_tokens: Caps how many iterations the model predicts tokens. :param top_k: Picks the *top_k* most likely tokens for addition. :param temp: Divided the probs from softmax. :param use_cached: Bool determining usage of KV cache. :param print_text: Bool (*True or False*) determining whether to print newly generated tokens. """ def get_next_id(ids: torch.Tensor, top_k: int, temp: float | int, use_cached: bool) -> torch.Tensor: """ Run ids through LLM forward, generating logits. Use top_k & and temp to determine next token. """ logits = self.forward(ids, use_cached) logits = logits.logits.squeeze(0) last_logits = logits[-1] / temp top_logits, top_ids = torch.topk(last_logits, top_k, dim=-1) probs = torch.softmax(top_logits, dim=-1) top_id_index = torch.multinomial(probs, 1).squeeze() return top_ids[top_id_index] def get_last_id(ids: torch.Tensor) -> torch.Tensor: """ Returns the last token id from ids, in shape *[B, 1]*. """ return ids[:, -1:] def join_next_id(ids: torch.Tensor, id: torch.Tensor) -> torch.Tensor: """ Concatenates id to ids. :param ids: Expected shape: *[B, seq_len]* :param id: Expected shape: *single int* """ return torch.cat(tensors=[ids, id.unsqueeze(0).unsqueeze(0)], dim=1) time_to_first_tok_start = time.time() sec_per_tok = [] self.reset_cache() # Pass all ids through model first to generate KV cache, then only the last token formatted_prompt = self.format_prompt(prompt) ids = tokenizer(formatted_prompt, return_tensors="pt").input_ids.to(device) next_id = get_next_id(ids=ids, top_k=top_k, temp=temp, use_cached=use_cached) if print_text: # Print decoded new token text = tokenizer.decode(next_id.unsqueeze(0)) print(text, end="") ids = join_next_id(ids=ids, id=next_id) # Calculate time to first token time_to_first_tok = time.time() - time_to_first_tok_start for _ in range(max_new_tokens - 1): start = time.time() # Generate next token model_input = get_last_id(ids=ids) if use_cached else ids next_id = get_next_id(ids=model_input, top_k=top_k, temp=temp, use_cached=use_cached) # Calculate seconds per token end = time.time() - start sec_per_tok.append(end) if next_id == tokenizer.eos_token_id: break if print_text: # Print decoded new token text = tokenizer.decode(next_id.unsqueeze(0)) print(text, end="") ids = join_next_id(ids=ids, id=next_id) if len(ids[-1]) >= self.config.context_length: break # Calculate tokens per second from seconds per token total_times = len(sec_per_tok) avg_sec_per_tok = sum(sec_per_tok) / total_times avg_tok_per_sec = 1 / avg_sec_per_tok return tokenizer.decode(ids.squeeze(0)), avg_tok_per_sec, time_to_first_tok AutoConfig.register("custom_llm", LLMConfig) AutoModelForCausalLM.register(LLMConfig, LLM) LLMConfig.register_for_auto_class() LLM.register_for_auto_class("AutoModelForCausalLM")