Delete veylon_model.py

veylon_model.py

@@ -1,515 +0,0 @@
-from __future__ import annotations
-
-import numpy as np
-import keras
-from keras import layers, ops
-import jax
-
-from veylon_attention import flash_splash_attention, decode_swa
-
-try:
-    from config import (
-        CONTEXT,
-        vocab_size as Vocab_size,
-        D_MODEL,
-        numberoflayers,
-        numberofheads,
-        d_Latent,
-        ffn_mult,
-        swa_window,
-        num_kv_heads,
-        use_moe,
-        moe_num_experts,
-        moe_top_k,
-        USE_REMAT,
-        MAX_GEN_TOKENS,
-    )
-except Exception:
-    CONTEXT = 2048
-    Vocab_size = 32000
-    D_MODEL = 512
-    numberoflayers = 8
-    numberofheads = 8
-    d_Latent = 128
-    ffn_mult = 3.5
-    swa_window = 1024
-    num_kv_heads = 2
-    use_moe = False
-    moe_num_experts = 8
-    moe_top_k = 2
-    USE_REMAT = True
-    MAX_GEN_TOKENS = 256
-
-# NOTE: mixed precision policy is set by train.py / finetune.py before model
-# creation.  Do NOT set it here — importing this module must not have side
-# effects.
-
-
-# ─────────────────────────────────────────────────────────────────────────────
-# Layers
-# ─────────────────────────────────────────────────────────────────────────────
-
-@keras.saving.register_keras_serializable()
-class RMSNorm(layers.Layer):
-    def __init__(self, epsilon=1e-5, **kwargs):
-        super().__init__(**kwargs)
-        self.epsilon = epsilon
-
-    def build(self, input_shape):
-        self.weight = self.add_weight(shape=(input_shape[-1],), initializer='ones', name='gamma')
-
-    def call(self, x):
-        x_fp32 = ops.cast(x, 'float32')
-        rms = ops.sqrt(ops.mean(ops.square(x_fp32), axis=-1, keepdims=True) + self.epsilon)
-        out = x_fp32 / rms
-        return ops.cast(out, x.dtype) * self.weight
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({'epsilon': self.epsilon})
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class RotaryEmbedding(layers.Layer):
-    """
-    RoPE with `max_seq_len + gen_headroom` so a full-context prompt can still
-    decode up to `gen_headroom` tokens without overflowing the cos/sin table.
-
-    `max_seq_len` is preserved in the config so checkpoints remain compatible —
-    only the (trainable=False) table shapes grow.
-    """
-
-    def __init__(self, max_seq_len, dim, theta=10000.0, gen_headroom=0, **kwargs):
-        super().__init__(**kwargs)
-        if dim % 2 != 0:
-            raise ValueError('RotaryEmbedding dim must be even.')
-        self.max_seq_len = max_seq_len
-        self.dim = dim
-        self.theta = theta
-        self.gen_headroom = int(gen_headroom)
-        self.table_size = max_seq_len + self.gen_headroom
-
-    def build(self, input_shape):
-        half = self.dim // 2
-        inv_freq = 1.0 / (self.theta ** (np.arange(0, self.dim, 2).astype(np.float32) / self.dim))
-        positions = np.arange(self.table_size, dtype=np.float32)
-        freqs = positions[:, None] * inv_freq[None, :]
-        self.cos = self.add_weight(
-            shape=(self.table_size, half),
-            initializer=keras.initializers.Constant(np.cos(freqs)),
-            trainable=False, dtype='float32', name='cos_table',
-        )
-        self.sin = self.add_weight(
-            shape=(self.table_size, half),
-            initializer=keras.initializers.Constant(np.sin(freqs)),
-            trainable=False, dtype='float32', name='sin_table',
-        )
-
-    def call(self, x, offset=0):
-        seq_len = x.shape[1]
-        if seq_len is None:
-            seq_len = ops.shape(x)[1]
-        half = self.dim // 2
-        if offset + seq_len > self.table_size:
-            raise ValueError(
-                f'RoPE table too small: offset={offset}, seq_len={seq_len}, '
-                f'max={self.table_size} (CONTEXT={self.max_seq_len} + gen_headroom={self.gen_headroom})'
-            )
-        cos = self.cos[offset:offset + seq_len, :half]
-        sin = self.sin[offset:offset + seq_len, :half]
-        cos = ops.cast(ops.reshape(cos, (1, seq_len, 1, half)), x.dtype)
-        sin = ops.cast(ops.reshape(sin, (1, seq_len, 1, half)), x.dtype)
-        x1 = x[..., :half]
-        x2 = x[..., half:]
-        return ops.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1)
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'max_seq_len': self.max_seq_len,
-            'dim': self.dim,
-            'theta': self.theta,
-            'gen_headroom': self.gen_headroom,
-        })
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class SwiGLUFFN(layers.Layer):
-    def __init__(self, d_model, hidden_mult=3.5, **kwargs):
-        super().__init__(**kwargs)
-        self.d_model_arg = d_model
-        self.hidden_mult = hidden_mult
-        self.hidden_dim = int(d_model * hidden_mult * 2 / 3)
-        self.hidden_dim = ((self.hidden_dim + 63) // 64) * 64
-
-    def build(self, input_shape):
-        d_model = input_shape[-1]
-        self.gate_up_proj = self.add_weight(shape=(d_model, 2 * self.hidden_dim), initializer='glorot_uniform', name='gate_up_proj')
-        self.down_proj = self.add_weight(shape=(self.hidden_dim, d_model), initializer='glorot_uniform', name='down_proj')
-
-    def call(self, x, training=False):
-        gate_up = ops.matmul(x, self.gate_up_proj)
-        gate, up = ops.split(gate_up, 2, axis=-1)
-        return ops.matmul(ops.silu(gate) * up, self.down_proj)
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({'d_model': self.d_model_arg, 'hidden_mult': self.hidden_mult})
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class MoE_FFN(layers.Layer):
-    """
-    Mixture-of-Experts FFN (top-k routing, masked dispatch).
-
-    For each routing slot k in [0, top_k) we run each expert once over the
-    full batch, zeroing inputs the expert was not selected for, then accumulate
-    the gated expert output.  The load-balancing loss uses top-1 argmax routing
-    for the `f` term (standard Switch-Transformer formulation); the previous
-    code used top-k membership which over-counted and produced the wrong loss.
-    """
-
-    def __init__(self, d_model, num_experts=8, top_k=2, hidden_mult=3.5, **kwargs):
-        super().__init__(**kwargs)
-        self.d_model_arg = d_model
-        self.num_experts = num_experts
-        self.top_k = top_k
-        self.hidden_mult = hidden_mult
-        self.experts = [SwiGLUFFN(d_model, hidden_mult) for _ in range(num_experts)]
-        self.router = layers.Dense(num_experts, use_bias=False)
-
-    def _load_balancing_loss(self, router_logits, top_idx):
-        # Standard Switch-Transformer aux loss.  `f` uses the argmax routing
-        # decision (top_idx[..., 0]); `p` uses the router softmax mean.
-        router_probs = ops.softmax(router_logits, axis=-1)
-        mask = ops.one_hot(top_idx[..., 0], self.num_experts)        # [..., num_experts]
-        f = ops.mean(ops.cast(mask, 'float32'), axis=tuple(range(mask.ndim - 1)))
-        p = ops.mean(router_probs, axis=tuple(range(router_probs.ndim - 1)))
-        return self.num_experts * ops.sum(f * p)
-
-    def call(self, x, training=False):
-        # x: [B, S, D]
-        B = ops.shape(x)[0]
-        S = ops.shape(x)[1]
-        D = x.shape[-1]
-
-        router_logits = self.router(x)                              # [B, S, E]
-        top_logits, top_idx = ops.top_k(router_logits, self.top_k)  # [B, S, k]
-        top_weights = ops.softmax(top_logits, axis=-1)              # [B, S, k]
-
-        if training:
-            self.add_loss(self._load_balancing_loss(router_logits, top_idx))
-
-        out = ops.zeros_like(x)
-        for kk in range(self.top_k):
-            idx_k = top_idx[..., kk]                                 # [B, S]
-            w_k = top_weights[..., kk][..., None]                    # [B, S, 1]
-            for e in range(self.num_experts):
-                mask_e = ops.cast(idx_k == e, x.dtype)[..., None]    # [B, S, 1]
-                # mask_e is always a tensor; skip the useless ops.is_tensor guard
-                # (it never triggered — removed FIX BUG11)
-                # Only run the expert where it is actually selected.
-                x_e = x * mask_e
-                expert_out = self.experts[e](x_e, training=training)
-                out = out + expert_out * w_k * mask_e
-        return out
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'd_model': self.d_model_arg, 'num_experts': self.num_experts,
-            'top_k': self.top_k, 'hidden_mult': self.hidden_mult,
-        })
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class MLAttention(layers.Layer):
-    def __init__(self, d_model, n_heads, d_latent, max_seq_len, num_kv_heads=2,
-                 swa_window=1024, attn_dropout=0.0, gen_headroom=0, **kwargs):
-        super().__init__(**kwargs)
-        if d_model % n_heads != 0:
-            raise ValueError('d_model must be divisible by n_heads')
-        if n_heads % num_kv_heads != 0:
-            raise ValueError('n_heads must be divisible by num_kv_heads')
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.num_kv_heads = num_kv_heads
-        self.group_size = n_heads // num_kv_heads
-        self.d_head = d_model // n_heads
-        self.d_latent = d_latent
-        self.max_seq_len = max_seq_len
-        self.swa_window = swa_window
-        self.gen_headroom = int(gen_headroom)
-        self.dropout = layers.Dropout(attn_dropout)
-        self.rope = RotaryEmbedding(max_seq_len, d_model // n_heads, gen_headroom=int(gen_headroom))  # FIX BUG10: must be in __init__ for Keras tracking
-
-    def build(self, input_shape):
-        self.W_qc = self.add_weight(shape=(self.d_model, self.d_model + self.d_latent), initializer='glorot_uniform', name='W_qc')
-        self.W_kv = self.add_weight(shape=(self.d_latent, self.num_kv_heads * 2 * self.d_head), initializer='glorot_uniform', name='W_kv')
-        self.W_o = self.add_weight(shape=(self.d_model, self.d_model), initializer='glorot_uniform', name='Wo')
-
-    def _project_kv(self, c):
-        kv = ops.matmul(c, self.W_kv)
-        return ops.split(kv, 2, axis=-1)
-
-    def call(self, x, training=False):
-        B = ops.shape(x)[0]
-        S = ops.shape(x)[1]
-        qc = ops.matmul(x, self.W_qc)
-        q_proj, c = ops.split(qc, [self.d_model], axis=-1)
-        q = ops.reshape(q_proj, (B, S, self.n_heads, self.d_head))
-        q = self.rope(q, offset=0)
-        q = ops.transpose(q, (0, 2, 1, 3))
-        k, v = self._project_kv(c)
-        k = ops.reshape(k, (B, S, self.num_kv_heads, self.d_head))
-        v = ops.reshape(v, (B, S, self.num_kv_heads, self.d_head))
-        k = self.rope(k, offset=0)
-        k = ops.transpose(k, (0, 2, 1, 3))
-        v = ops.transpose(v, (0, 2, 1, 3))
-
-        # Attention remat is REQUIRED during training: it is the per-block
-        # jax.checkpoint(process_block) that stops XLA from staging the hidden
-        # [S, B, Hkv, W, D] buffer (≈8.6 GB here) inside the fori_loop.  Block-
-        # level remat does NOT prevent that forward-time allocation, so this
-        # must follow `training` regardless of config.USE_REMAT.
-        out = flash_splash_attention(
-            q, k, v,
-            window_size=min(self.swa_window, self.max_seq_len),
-            backend=jax.default_backend(),
-            use_gqa=True,
-            use_remat=training,
-        )
-        out = ops.transpose(out, (0, 2, 1, 3))
-        out = ops.reshape(out, (B, S, self.d_model))
-        out = self.dropout(out, training=training)
-        return ops.matmul(out, self.W_o)
-
-    def generate_step(self, x, cache_k=None, cache_v=None, cache_pos=0):
-        B = ops.shape(x)[0]
-        S = ops.shape(x)[1]
-
-        qc = ops.matmul(x, self.W_qc)
-        q_proj, c = ops.split(qc, [self.d_model], axis=-1)
-
-        q = ops.reshape(q_proj, (B, S, self.n_heads, self.d_head))
-        q = self.rope(q, offset=cache_pos)
-        q = ops.transpose(q, (0, 2, 1, 3))
-
-        k, v = self._project_kv(c)
-        k = ops.reshape(k, (B, S, self.num_kv_heads, self.d_head))
-        v = ops.reshape(v, (B, S, self.num_kv_heads, self.d_head))
-        k = self.rope(k, offset=cache_pos)
-        k = ops.transpose(k, (0, 2, 1, 3))
-        v = ops.transpose(v, (0, 2, 1, 3))
-
-        if cache_k is None:
-            # Prefill: no remat (inference only), no window trim needed.
-            out = flash_splash_attention(
-                q, k, v,
-                window_size=min(self.swa_window, self.max_seq_len),
-                backend=jax.default_backend(),
-                use_gqa=True,
-                use_remat=False,
-            )
-            new_k = k[:, :, -self.swa_window:, :]
-            new_v = v[:, :, -self.swa_window:, :]
-        else:
-            if S != 1:
-                raise ValueError(f'generate_step with cache expects S=1, got S={S}')
-
-            k = ops.concatenate([cache_k, k], axis=2)
-            v = ops.concatenate([cache_v, v], axis=2)
-
-            k = k[:, :, -self.swa_window:, :]
-            v = v[:, :, -self.swa_window:, :]
-
-            out = decode_swa(q, k, v)
-            new_k = k
-            new_v = v
-
-        out = ops.transpose(out, (0, 2, 1, 3))
-        out = ops.reshape(out, (B, S, self.d_model))
-        out = ops.matmul(out, self.W_o)
-        return out, new_k, new_v
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'd_model': self.d_model, 'n_heads': self.n_heads,
-            'num_kv_heads': self.num_kv_heads, 'd_latent': self.d_latent,
-            'max_seq_len': self.max_seq_len, 'swa_window': self.swa_window,
-            'attn_dropout': self.dropout.rate, 'gen_headroom': self.gen_headroom,
-        })
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class TransformerBlock(layers.Layer):
-    def __init__(self, d_model, n_heads, d_latent, ffn_layer, max_seq_len,
-                 num_kv_heads=2, swa_window=1024, use_remat=True, gen_headroom=0, **kwargs):
-        super().__init__(**kwargs)
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.d_latent = d_latent
-        self.max_seq_len = max_seq_len
-        self.num_kv_heads = num_kv_heads
-        self.swa_window = swa_window
-        self.use_remat = use_remat
-        self.gen_headroom = int(gen_headroom)
-        self.ffn = keras.saving.deserialize_keras_object(ffn_layer) if isinstance(ffn_layer, dict) else ffn_layer
-        self.norm1 = RMSNorm()
-        self.norm2 = RMSNorm()
-        self.attn = MLAttention(
-            d_model, n_heads, d_latent, max_seq_len,
-            num_kv_heads=num_kv_heads, swa_window=swa_window, gen_headroom=self.gen_headroom,
-        )
-
-    def call(self, x, training=False):
-        # Block-level activation/gradient checkpointing (optional, via
-        # config.USE_REMAT).  NOTE: this is independent of attention's own
-        # remat, which MUST stay on whenever training=True — that's what stops
-        # XLA staging the hidden [S,B,Hkv,W,D] buffer (the OOM we hit).
-        if training and self.use_remat:
-            def _fwd(x_in):
-                attn_out = self.attn(self.norm1(x_in), training=True)
-                ffn_out = self.ffn(self.norm2(x_in), training=True)
-                return attn_out + ffn_out
-            return x + jax.checkpoint(_fwd)(x)
-        a = self.attn(self.norm1(x), training=training)
-        f = self.ffn(self.norm2(x), training=training)
-        return x + a + f
-
-    def generate_step(self, x, cache_k=None, cache_v=None, cache_pos=0):
-        # Must mirror call()'s PARALLEL residual structure exactly:
-        # attn and ffn both read from norm(x) computed on the SAME
-        # pre-block x, then both get added to that same original x.
-        # Previously this was written as a sequential block (x=x+attn;
-        # x=x+ffn(norm2(x))) — ffn ended up conditioned on norm2(x+attn_out)
-        # instead of norm2(x), an input distribution the FFN weights were
-        # never trained on. That silently produced a different model at
-        # generation time than at training time: teacher-forced loss (via
-        # call()) never exercises this path and looked fine, while actual
-        # autoregressive generation compounded the mismatch every step.
-        attn_out, nck, ncv = self.attn.generate_step(
-            self.norm1(x), cache_k=cache_k, cache_v=cache_v, cache_pos=cache_pos,
-        )
-        ffn_out = self.ffn(self.norm2(x), training=False)
-        x = x + attn_out + ffn_out
-        return x, nck, ncv
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'd_model': self.d_model, 'n_heads': self.n_heads, 'd_latent': self.d_latent,
-            'ffn_layer': keras.saving.serialize_keras_object(self.ffn),
-            'max_seq_len': self.max_seq_len, 'num_kv_heads': self.num_kv_heads,
-            'swa_window': self.swa_window, 'use_remat': self.use_remat,
-            'gen_headroom': self.gen_headroom,
-        })
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class VeylonModel(keras.Model):
-    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_latent, ffn_mult,
-                 max_seq_len, use_moe=False, moe_num_experts=8, moe_top_k=2,
-                 num_kv_heads=2, swa_window=1024, use_remat=True, gen_headroom=0, **kwargs):
-        super().__init__(**kwargs)
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.n_layers = n_layers
-        self.n_heads = n_heads
-        self.d_latent = d_latent
-        self.ffn_mult = ffn_mult
-        self.max_seq_len = max_seq_len
-        self.use_moe = use_moe
-        self.moe_num_experts = moe_num_experts
-        self.moe_top_k = moe_top_k
-        self.num_kv_heads = num_kv_heads
-        self.swa_window = swa_window
-        self.use_remat = use_remat
-        self.gen_headroom = int(gen_headroom)
-        self.embedding = layers.Embedding(vocab_size, d_model, name='token_embedding')
-        self.blocks = []
-        for i in range(n_layers):
-            ffn = (MoE_FFN(d_model, moe_num_experts, moe_top_k, ffn_mult)
-                   if use_moe else SwiGLUFFN(d_model, ffn_mult))
-            self.blocks.append(TransformerBlock(
-                d_model, n_heads, d_latent, ffn, max_seq_len,
-                num_kv_heads=num_kv_heads, swa_window=swa_window,
-                use_remat=use_remat, gen_headroom=self.gen_headroom, name=f'block_{i}',
-            ))
-        self.norm = RMSNorm()
-
-    def call(self, inputs, training=False):
-        x = self.embedding(inputs)
-        for block in self.blocks:
-            x = block(x, training=training)
-        x = self.norm(x)
-        embedding_weights = self.embedding.embeddings
-        logits = ops.matmul(x, ops.transpose(embedding_weights))
-        return ops.cast(logits, 'float32')
-
-    def generate_step(self, inputs, cache_k=None, cache_v=None, cache_pos=0):
-        x = self.embedding(inputs)
-        new_cache_k = []
-        new_cache_v = []
-
-        if cache_k is None:
-            cache_k = [None] * len(self.blocks)
-            cache_v = [None] * len(self.blocks)
-
-        for i, block in enumerate(self.blocks):
-            x, nck, ncv = block.generate_step(
-                x, cache_k=cache_k[i], cache_v=cache_v[i], cache_pos=cache_pos,
-            )
-            new_cache_k.append(nck)
-            new_cache_v.append(ncv)
-
-        x = self.norm(x)
-        embedding_weights = self.embedding.embeddings
-        logits = ops.matmul(x, ops.transpose(embedding_weights))
-        logits = ops.cast(logits, 'float32')
-        return logits, new_cache_k, new_cache_v
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'vocab_size': self.vocab_size, 'd_model': self.d_model,
-            'n_layers': self.n_layers, 'n_heads': self.n_heads,
-            'd_latent': self.d_latent, 'ffn_mult': self.ffn_mult,
-            'max_seq_len': self.max_seq_len, 'use_moe': self.use_moe,
-            'moe_num_experts': self.moe_num_experts, 'moe_top_k': self.moe_top_k,
-            'num_kv_heads': self.num_kv_heads, 'swa_window': self.swa_window,
-            'use_remat': self.use_remat, 'gen_headroom': self.gen_headroom,
-        })
-        return cfg
-
-
-# ─────────────────────────────────────────────────────────────────────────────
-# Factory
-# ─────────────────────────────────────────────────────────────────────────────
-
-def create_llm(
-    vocab_size=Vocab_size, d_model=D_MODEL, n_layers=numberoflayers, n_heads=numberofheads,
-    d_latent=d_Latent, ffn_mult=ffn_mult, max_seq_len=CONTEXT, use_moe=use_moe,
-    moe_num_experts=moe_num_experts, moe_top_k=moe_top_k, num_kv_heads=num_kv_heads,
-    swa_window=swa_window, use_remat=USE_REMAT, gen_headroom=MAX_GEN_TOKENS,
-):
-    """
-    Build a VeylonModel.
-
-    `gen_headroom` sizes the RoPE table to (max_seq_len + gen_headroom) so that
-    a full-length prompt can still generate `gen_headroom` tokens during
-    inference without "RoPE table too small" errors.  Defaults to MAX_GEN_TOKENS.
-    """
-    return VeylonModel(
-        vocab_size=vocab_size, d_model=d_model, n_layers=n_layers, n_heads=n_heads,
-        d_latent=d_latent, ffn_mult=ffn_mult, max_seq_len=max_seq_len, use_moe=use_moe,
-        moe_num_experts=moe_num_experts, moe_top_k=moe_top_k, num_kv_heads=num_kv_heads,
-        swa_window=swa_window, use_remat=use_remat, gen_headroom=gen_headroom,
-    )