Create modeling_momo.py

modeling_momo.py

@@ -0,0 +1,284 @@
+# modeling_momo.py
+# 🌸 Momo-336M — HuggingFace compatible model definition
+# Upload this file to your HF repo alongside config.json and configuration_momo.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_momo import MomoConfig
+
+
+# ════════════════════════════════════════════════════════════════
+#  COMPONENTS
+# ════════════════════════════════════════════════════════════════
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x.float() * rms).to(x.dtype) * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_seq=512, theta=10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+        self._cache(max_seq)
+
+    def _cache(self, n):
+        t = torch.arange(n, device=self.inv_freq.device).float()
+        freq = torch.outer(t, self.inv_freq)
+        emb = torch.cat([freq, freq], dim=-1)
+        self.register_buffer('cos_c', emb.cos()[None, None])
+        self.register_buffer('sin_c', emb.sin()[None, None])
+
+    def forward(self, x, seq_len):
+        if seq_len > self.cos_c.shape[2]:
+            self._cache(seq_len)
+        return (
+            self.cos_c[:, :, :seq_len].to(x.dtype),
+            self.sin_c[:, :, :seq_len].to(x.dtype),
+        )
+
+
+def rot_half(x):
+    a, b = x.chunk(2, dim=-1)
+    return torch.cat([-b, a], dim=-1)
+
+
+def apply_rope(q, k, cos, sin):
+    return (q * cos) + (rot_half(q) * sin), (k * cos) + (rot_half(k) * sin)
+
+
+# ════════════════════════════════════════════════════════════════
+#  ATTENTION — Grouped Query Attention (GQA)
+# ════════════════════════════════════════════════════════════════
+
+class MomoAttention(nn.Module):
+    def __init__(self, cfg: MomoConfig):
+        super().__init__()
+        self.nh  = cfg.num_attention_heads
+        self.nkv = cfg.num_key_value_heads
+        self.hd  = cfg.hidden_size // cfg.num_attention_heads
+        self.grp = self.nh // self.nkv
+        self.sc  = self.hd ** -0.5
+        H = cfg.hidden_size
+        self.q    = nn.Linear(H, self.nh  * self.hd, bias=False)
+        self.k    = nn.Linear(H, self.nkv * self.hd, bias=False)
+        self.v    = nn.Linear(H, self.nkv * self.hd, bias=False)
+        self.o    = nn.Linear(self.nh * self.hd, H,  bias=False)
+        self.rope = RotaryEmbedding(self.hd, cfg.max_position_embeddings, cfg.rope_theta)
+
+    def forward(self, x, mask=None, past=None, use_cache=False):
+        B, T, _ = x.shape
+        q = self.q(x).view(B, T, self.nh,  self.hd).transpose(1, 2)
+        k = self.k(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
+        v = self.v(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
+
+        past_len = past[0].shape[2] if past is not None else 0
+        cos, sin = self.rope(q, past_len + T)
+        cos = cos[:, :, past_len:past_len + T]
+        sin = sin[:, :, past_len:past_len + T]
+        q, k = apply_rope(q, k, cos, sin)
+
+        if self.grp > 1:
+            k = k[:, None].expand(-1, self.grp, -1, -1, -1).reshape(B, self.nh, T, self.hd)
+            v = v[:, None].expand(-1, self.grp, -1, -1, -1).reshape(B, self.nh, T, self.hd)
+
+        if past is not None:
+            pk, pv = past
+            k = torch.cat([pk, k], 2)
+            v = torch.cat([pv, v], 2)
+
+        pres = (k, v) if use_cache else None
+        S = k.shape[2]
+        a = torch.matmul(q, k.transpose(-2, -1)) * self.sc
+        causal = torch.triu(
+            torch.full((T, S), float('-inf'), device=x.device),
+            diagonal=S - T + 1
+        )
+        a = a + causal
+        if mask is not None:
+            a = a + mask
+        a = F.softmax(a, dim=-1)
+        out = torch.matmul(a, v).transpose(1, 2).reshape(B, T, -1)
+        return self.o(out), pres
+
+
+# ════════════════════════════════════════════════════════════════
+#  FEED-FORWARD — SwiGLU
+# ════════════════════════════════════════════════════════════════
+
+class MomoFFN(nn.Module):
+    def __init__(self, cfg: MomoConfig):
+        super().__init__()
+        self.gate = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
+        self.up   = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
+        self.down = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+# ════════════════════════════════════════════════════════════════
+#  TRANSFORMER BLOCK
+# ════════════════════════════════════════════════════════════════
+
+class MomoBlock(nn.Module):
+    def __init__(self, cfg: MomoConfig):
+        super().__init__()
+        self.attn  = MomoAttention(cfg)
+        self.ffn   = MomoFFN(cfg)
+        self.norm1 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.norm2 = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+
+    def forward(self, x, mask=None, past=None, use_cache=False):
+        a, p = self.attn(self.norm1(x), mask, past, use_cache)
+        x = x + a
+        x = x + self.ffn(self.norm2(x))
+        return x, p
+
+
+# ════════════════════════════════════════════════════════════════
+#  🌸 MOMO FOR CAUSAL LM — Main Model
+# ════════════════════════════════════════════════════════════════
+
+class MomoForCausalLM(PreTrainedModel):
+    config_class = MomoConfig
+    _no_split_modules = ['MomoBlock']
+
+    def __init__(self, cfg: MomoConfig):
+        super().__init__(cfg)
+        self.embed    = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.layers   = nn.ModuleList([MomoBlock(cfg) for _ in range(cfg.num_hidden_layers)])
+        self.norm     = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.lm_head  = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+        self.lm_head.weight = self.embed.weight   # weight tying
+        self.grad_ckpt = cfg.use_gradient_checkpointing
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, std=0.02)
+
+    def get_input_embeddings(self):
+        return self.embed
+
+    def set_input_embeddings(self, value):
+        self.embed = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        past_key_values=None,
+        use_cache=False,
+        **kwargs,
+    ):
+        x = self.embed(input_ids)
+        pkvs = past_key_values or [None] * len(self.layers)
+        cache = []
+
+        for layer, past in zip(self.layers, pkvs):
+            if self.grad_ckpt and self.training:
+                def _fn(layer):
+                    def fn(x):
+                        out, _ = layer(x, mask=attention_mask, use_cache=False)
+                        return out
+                    return fn
+                x = torch.utils.checkpoint.checkpoint(
+                    _fn(layer), x, use_reentrant=False
+                )
+                cache.append(None)
+            else:
+                x, p = layer(x, attention_mask, past, use_cache)
+                cache.append(p)
+
+        x      = self.norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits[..., :-1, :].contiguous().view(-1, logits.size(-1)),
+                labels[..., 1:].contiguous().view(-1),
+                ignore_index=-100,
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=cache if use_cache else None,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids,
+        max_new_tokens=300,
+        temperature=0.75,
+        top_k=50,
+        top_p=0.92,
+        rep_penalty=1.1,
+        eos_token_id=None,
+        pad_token_id=None,
+        **kwargs,
+    ):
+        self.eval()
+        gen  = input_ids.clone()
+        past = None
+
+        for _ in range(max_new_tokens):
+            inp = gen if past is None else gen[:, -1:]
+            out = self(inp, use_cache=True, past_key_values=past)
+            past   = out.past_key_values
+            logits = out.logits[:, -1, :].float()
+
+            # Repetition penalty
+            if rep_penalty != 1.0:
+                for tok in set(gen[0].tolist()):
+                    if logits[0, tok] > 0:
+                        logits[0, tok] /= rep_penalty
+                    else:
+                        logits[0, tok] *= rep_penalty
+
+            logits = logits / max(temperature, 1e-6)
+
+            # Top-k
+            if top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, -1:]] = float('-inf')
+
+            # Top-p (nucleus)
+            if top_p < 1.0:
+                sl, si = torch.sort(logits, descending=True)
+                cp = torch.cumsum(F.softmax(sl, dim=-1), dim=-1)
+                sl[cp - F.softmax(sl, dim=-1) > top_p] = float('-inf')
+                logits.scatter_(1, si, sl)
+
+            next_tok = torch.multinomial(F.softmax(logits, dim=-1), 1)
+            gen = torch.cat([gen, next_tok], dim=1)
+
+            if eos_token_id is not None and (next_tok == eos_token_id).all():
+                break
+
+        return gen