modules.py

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")