| """CORTEX model definition — the trainable Frankenstein-Labs architecture. |
| |
| This module defines the model the CORTEX training pipeline builds and trains. It is a |
| standard decoder-only Transformer with: |
| |
| * RMSNorm (pre-norm) |
| * rotary position embeddings (RoPE) |
| * grouped-query attention (GQA) |
| * a SwiGLU feed-forward network |
| |
| Nothing here loads, mutates or depends on the distributed 1.65T checkpoint that this |
| repository also hosts. Weights produced from this module are initialised from scratch |
| and are owned by Frankenstein-Labs. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import math |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| __all__ = ["CortexConfig", "CortexForCausalLM", "count_parameters"] |
|
|
|
|
| @dataclass |
| class CortexConfig: |
| """Hyper-parameters of a trainable CORTEX model.""" |
|
|
| vocab_size: int = 129280 |
| hidden_size: int = 768 |
| num_hidden_layers: int = 12 |
| num_attention_heads: int = 12 |
| num_key_value_heads: int = 4 |
| intermediate_size: int = 2048 |
| hidden_act: str = "silu" |
| max_position_embeddings: int = 2048 |
| rope_theta: float = 10000.0 |
| rms_norm_eps: float = 1e-5 |
| attention_bias: bool = False |
| attention_dropout: float = 0.0 |
| tie_word_embeddings: bool = True |
| initializer_range: float = 0.02 |
| bos_token_id: int = 0 |
| eos_token_id: int = 1 |
| pad_token_id: int = 2 |
| model_name: str = "cortex-dev-1" |
|
|
| @property |
| def head_dim(self) -> int: |
| if self.hidden_size % self.num_attention_heads: |
| raise ValueError( |
| f"hidden_size ({self.hidden_size}) must be divisible by " |
| f"num_attention_heads ({self.num_attention_heads})" |
| ) |
| return self.hidden_size // self.num_attention_heads |
|
|
| def validate(self) -> None: |
| if self.num_attention_heads % self.num_key_value_heads: |
| raise ValueError( |
| f"num_attention_heads ({self.num_attention_heads}) must be a multiple of " |
| f"num_key_value_heads ({self.num_key_value_heads})" |
| ) |
| if self.hidden_size % 2: |
| raise ValueError("hidden_size must be even so RoPE can split the head dimension") |
| if self.head_dim % 2: |
| raise ValueError("head_dim must be even for RoPE") |
| if self.vocab_size <= 0: |
| raise ValueError("vocab_size must be positive") |
| if self.num_hidden_layers <= 0: |
| raise ValueError("num_hidden_layers must be positive") |
|
|
| @classmethod |
| def from_json(cls, path: str | Path) -> "CortexConfig": |
| raw = json.loads(Path(path).read_text(encoding="utf-8")) |
| fields = set(cls.__dataclass_fields__) |
| return cls(**{k: v for k, v in raw.items() if k in fields}) |
|
|
| def to_json(self, path: str | Path) -> None: |
| Path(path).write_text( |
| json.dumps(asdict(self), indent=2, ensure_ascii=False) + "\n", encoding="utf-8" |
| ) |
|
|
| def parameter_breakdown(self) -> dict: |
| """Analytical parameter count, including norms and (tied) embeddings.""" |
| h = self.hidden_size |
| kv = self.num_key_value_heads * self.head_dim |
| q = self.num_attention_heads * self.head_dim |
| attn = h * q + h * kv * 2 + q * h |
| if self.attention_bias: |
| attn += q + kv * 2 |
| mlp = 3 * h * self.intermediate_size |
| norms = 2 * h |
| per_layer = attn + mlp + norms |
| embeds = self.vocab_size * h if self.tie_word_embeddings else self.vocab_size * h * 2 |
| total = embeds + per_layer * self.num_hidden_layers + h |
| return { |
| "embedding_and_head_shared": self.vocab_size * h, |
| "per_layer_attention": attn, |
| "per_layer_mlp": mlp, |
| "per_layer_norms": norms, |
| "total_per_layer": per_layer, |
| "total_estimated": total, |
| } |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-5): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| dtype = x.dtype |
| x = x.float() |
| x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| return x.to(dtype) * self.weight |
|
|
|
|
| def build_rope_cache(head_dim: int, max_position_embeddings: int, theta: float, device, dtype): |
| """Precompute cos/sin tables of shape [max_pos, head_dim].""" |
| inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) |
| positions = torch.arange(max_position_embeddings, device=device).float() |
| freqs = torch.outer(positions, inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| return emb.cos().to(dtype), emb.sin().to(dtype) |
|
|
|
|
| def rotate_half(x: torch.Tensor) -> torch.Tensor: |
| half = x.shape[-1] // 2 |
| x1, x2 = x[..., :half], x[..., half:] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| def apply_rope(q, k, cos, sin): |
| cos = cos.unsqueeze(0).unsqueeze(0) |
| sin = sin.unsqueeze(0).unsqueeze(0) |
| return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin |
|
|
|
|
| class CortexAttention(nn.Module): |
| def __init__(self, config: CortexConfig): |
| super().__init__() |
| self.config = config |
| self.num_heads = config.num_attention_heads |
| self.num_kv_heads = config.num_key_value_heads |
| self.head_dim = config.head_dim |
| self.num_kv_groups = self.num_heads // self.num_kv_heads |
|
|
| self.q_proj = nn.Linear( |
| config.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias |
| ) |
| self.k_proj = nn.Linear( |
| config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias |
| ) |
| self.v_proj = nn.Linear( |
| config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias |
| ) |
| self.o_proj = nn.Linear( |
| self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias |
| ) |
| self.attention_dropout = config.attention_dropout |
|
|
| def forward(self, hidden_states, cos, sin, attention_mask=None): |
| batch, seq, _ = hidden_states.shape |
| q = self.q_proj(hidden_states).view(batch, seq, self.num_heads, self.head_dim) |
| k = self.k_proj(hidden_states).view(batch, seq, self.num_kv_heads, self.head_dim) |
| v = self.v_proj(hidden_states).view(batch, seq, self.num_kv_heads, self.head_dim) |
|
|
| q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) |
| q, k = apply_rope(q, k, cos, sin) |
|
|
| if self.num_kv_groups > 1: |
| k = k.repeat_interleave(self.num_kv_groups, dim=1) |
| v = v.repeat_interleave(self.num_kv_groups, dim=1) |
|
|
| dropout_p = self.attention_dropout if self.training else 0.0 |
| attn = F.scaled_dot_product_attention( |
| q, k, v, attn_mask=attention_mask, dropout_p=dropout_p, |
| is_causal=attention_mask is None, |
| ) |
| attn = attn.transpose(1, 2).reshape(batch, seq, self.num_heads * self.head_dim) |
| return self.o_proj(attn) |
|
|
|
|
| class CortexMLP(nn.Module): |
| def __init__(self, config: CortexConfig): |
| super().__init__() |
| self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) |
| self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) |
| self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) |
| self.act = nn.SiLU() if config.hidden_act == "silu" else nn.GELU() |
|
|
| def forward(self, x): |
| return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class CortexDecoderLayer(nn.Module): |
| def __init__(self, config: CortexConfig): |
| super().__init__() |
| self.self_attn = CortexAttention(config) |
| self.mlp = CortexMLP(config) |
| self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
|
|
| def forward(self, hidden_states, cos, sin, attention_mask=None): |
| residual = hidden_states |
| hidden_states = self.self_attn( |
| self.input_layernorm(hidden_states), cos, sin, attention_mask |
| ) |
| hidden_states = residual + hidden_states |
|
|
| residual = hidden_states |
| hidden_states = self.mlp(self.post_attention_layernorm(hidden_states)) |
| return residual + hidden_states |
|
|
|
|
| class CortexForCausalLM(nn.Module): |
| """Decoder-only causal language model, the CORTEX training target.""" |
|
|
| def __init__(self, config: CortexConfig): |
| super().__init__() |
| config.validate() |
| self.config = config |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList( |
| [CortexDecoderLayer(config) for _ in range(config.num_hidden_layers)] |
| ) |
| self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| if config.tie_word_embeddings: |
| self.lm_head.weight = self.embed_tokens.weight |
| self.apply(self._init_weights) |
| for name, param in self.named_parameters(): |
| if name.endswith(("o_proj.weight", "down_proj.weight")): |
| nn.init.normal_( |
| param, mean=0.0, |
| std=config.initializer_range / math.sqrt(2 * config.num_hidden_layers), |
| ) |
| self._rope_cache = None |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) |
|
|
| def _rope(self, seq_len, device, dtype): |
| if ( |
| self._rope_cache is None |
| or self._rope_cache[0].shape[0] < seq_len |
| or self._rope_cache[0].device != device |
| ): |
| self._rope_cache = build_rope_cache( |
| self.config.head_dim, self.config.max_position_embeddings, |
| self.config.rope_theta, device, dtype, |
| ) |
| cos, sin = self._rope_cache |
| return cos[:seq_len], sin[:seq_len] |
|
|
| def forward(self, input_ids, attention_mask=None, labels=None): |
| batch, seq = input_ids.shape |
| hidden_states = self.embed_tokens(input_ids) |
| cos, sin = self._rope(seq, input_ids.device, hidden_states.dtype) |
|
|
| causal = None |
| if attention_mask is not None: |
| causal = torch.tril( |
| torch.ones(seq, seq, dtype=torch.bool, device=input_ids.device) |
| )[None, None, :, :] |
| causal = causal & attention_mask[:, None, None, :].bool() |
|
|
| for layer in self.layers: |
| hidden_states = layer(hidden_states, cos, sin, causal) |
|
|
| logits = self.lm_head(self.norm(hidden_states)) |
| loss = None |
| if labels is not None: |
| loss = self._loss(logits, labels) |
| return {"logits": logits, "loss": loss} |
|
|
| @staticmethod |
| def _loss(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: |
| """Cross-entropy computed in chunks over the vocabulary. |
| |
| A full fp32 logits tensor of [batch*seq, vocab] is several hundred megabytes |
| for the CORTEX vocabulary, so the softmax is done in slices. |
| """ |
| shift_logits = logits[:, :-1] |
| shift_labels = labels[:, 1:] |
| total, count = 0.0, 0 |
| chunk = 8192 |
| flat_logits = shift_logits.reshape(-1, shift_logits.size(-1)) |
| flat_labels = shift_labels.reshape(-1) |
| for start in range(0, flat_labels.size(0), chunk): |
| sl = flat_logits[start:start + chunk].float() |
| lb = flat_labels[start:start + chunk] |
| valid = lb.ne(-100) |
| if not valid.any(): |
| continue |
| total = total + F.cross_entropy(sl, lb, ignore_index=-100, reduction="sum") |
| count += int(valid.sum()) |
| if count == 0: |
| return torch.zeros((), device=logits.device, requires_grad=True) |
| return total / count |
|
|
| @torch.no_grad() |
| def generate(self, input_ids, max_new_tokens=32, temperature=1.0, eos_token_id=None): |
| self.eval() |
| for _ in range(max_new_tokens): |
| ctx = input_ids[:, -self.config.max_position_embeddings:] |
| logits = self.forward(ctx)["logits"][:, -1, :] / max(temperature, 1e-5) |
| probs = torch.softmax(logits.float(), dim=-1) |
| next_id = torch.multinomial(probs, num_samples=1) |
| input_ids = torch.cat([input_ids, next_id], dim=1) |
| if eos_token_id is not None and bool((next_id == eos_token_id).all()): |
| break |
| return input_ids |
|
|
|
|
| def count_parameters(model: nn.Module) -> int: |
| """Count unique parameters, so tied weights are not double counted.""" |
| seen, total = set(), 0 |
| for param in model.parameters(): |
| if id(param) in seen: |
| continue |
| seen.add(id(param)) |
| total += param.numel() |
| return total |
|
|
