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 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) -> torch.Tensor: seq_len = x.shape[-2] cos = rope_cos[:seq_len].to(x.dtype) sin = rope_sin[:seq_len].to(x.dtype) return x * cos + rotate_half(x) * sin 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) 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) def forward(self, x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor): batch_size, seq_len, _ = x.shape queries = self.query_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) keys = self.key_proj(x).unsqueeze(1) values = self.value_proj(x).unsqueeze(1) queries = apply_rope(queries, rope_cos, rope_sin) keys = apply_rope(keys, rope_cos, rope_sin) keys = keys.expand(batch_size, self.num_heads, seq_len, self.head_dim) values = values.expand(batch_size, self.num_heads, seq_len, 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=True) 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, rope_cos: torch.Tensor, rope_sin: torch.Tensor): shortcut = x x = self.norm1(x) x = self.attention(x, rope_cos, rope_sin) x = x + shortcut shortcut = x x = self.norm2(x) x = self.ff(x) x = x + shortcut return x 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"} 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, **kwargs): _, seq_len = input_ids.shape assert seq_len <= self.rope_cos.shape[0], ( f"seq_len {seq_len} exceeds cached 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) logits = self.out(self.norm(x)) return CausalLMOutputWithPast(logits=logits) AutoConfig.register("custom_llm", LLMConfig) AutoModelForCausalLM.register(LLMConfig, LLM) LLMConfig.register_for_auto_class() LLM.register_for_auto_class("AutoModelForCausalLM")