SFT on Alpaca 52k, 3 epochs, bf16 on H200

README.md

@@ -1,3 +1,56 @@
 ---
+library_name: transformers
 license: apache-2.0
+base_model: harims95/LoopLM-135M-naive
+tags:
+- generated_from_trainer
+model-index:
+- name: LoopLM-135M-naive-sft
+  results: []
 ---
+
+<!-- This model card has been generated automatically according to the information the Trainer had access to. You
+should probably proofread and complete it, then remove this comment. -->
+
+# LoopLM-135M-naive-sft
+
+This model is a fine-tuned version of [harims95/LoopLM-135M-naive](https://huggingface.co/harims95/LoopLM-135M-naive) on an unknown dataset.
+
+## Model description
+
+More information needed
+
+## Intended uses & limitations
+
+More information needed
+
+## Training and evaluation data
+
+More information needed
+
+## Training procedure
+
+### Training hyperparameters
+
+The following hyperparameters were used during training:
+- learning_rate: 2e-05
+- train_batch_size: 16
+- eval_batch_size: 8
+- seed: 42
+- gradient_accumulation_steps: 2
+- total_train_batch_size: 32
+- optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
+- lr_scheduler_type: cosine
+- lr_scheduler_warmup_steps: 100
+- num_epochs: 3
+
+### Training results
+
+
+
+### Framework versions
+
+- Transformers 5.12.1
+- Pytorch 2.12.1+cu130
+- Datasets 5.0.0
+- Tokenizers 0.22.2

configuration_looplm.py

@@ -0,0 +1,44 @@
+"""LoopLM configuration class for HuggingFace integration."""
+from transformers import PretrainedConfig
+
+
+class LoopLMConfig(PretrainedConfig):
+    model_type = "looplm"
+
+    def __init__(
+        self,
+        vocab_size: int = 50304,
+        d_model: int = 1024,
+        n_prelude: int = 4,
+        n_coda: int = 2,
+        mu_rec: int = 6,
+        n_q_heads: int = 16,
+        n_kv_heads: int = 8,
+        head_dim: int = 64,
+        qk_norm: bool = True,
+        rope_theta: float = 10000.0,
+        dense_ffn: int = 2816,
+        tie_embeddings: bool = True,
+        final_z_loss_coef: float = 1e-4,
+        use_a_matrix: bool = False,
+        use_input_norm: bool = False,
+        init_std: float = 0.02,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_prelude = n_prelude
+        self.n_coda = n_coda
+        self.mu_rec = mu_rec
+        self.n_q_heads = n_q_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = head_dim
+        self.qk_norm = qk_norm
+        self.rope_theta = rope_theta
+        self.dense_ffn = dense_ffn
+        self.tie_embeddings = tie_embeddings
+        self.final_z_loss_coef = final_z_loss_coef
+        self.use_a_matrix = use_a_matrix
+        self.use_input_norm = use_input_norm
+        self.init_std = init_std
+        super().__init__(**kwargs)

modeling_looplm.py

@@ -0,0 +1,342 @@
+"""Self-contained LoopLM model definition for HuggingFace integration."""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def rms_norm(x: Tensor, weight: Tensor | None = None, eps: float = 1e-6) -> Tensor:
+    out = F.rms_norm(x, (x.size(-1),), eps=eps)
+    return out * weight if weight is not None else out
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: Tensor) -> Tensor:
+        return rms_norm(x, self.weight, self.eps)
+
+
+def precompute_rope(head_dim: int, max_seq: int, theta: float, device) -> tuple[Tensor, Tensor]:
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(max_seq, device=device).float()
+    freqs = torch.outer(t, inv_freq)
+    return freqs.cos(), freqs.sin()
+
+
+def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+    seq_len, dim = x.shape[-2], x.shape[-1]
+    cos = cos[:seq_len].view(1, 1, seq_len, dim // 2)
+    sin = sin[:seq_len].view(1, 1, seq_len, dim // 2)
+    x1, x2 = x[..., : dim // 2], x[..., dim // 2 :]
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class Attention(nn.Module):
+    """GQA attention with per-head QK-norm before RoPE."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.nq, self.nkv, self.hd = cfg.n_q_heads, cfg.n_kv_heads, cfg.head_dim
+        assert self.nq % self.nkv == 0, "n_q_heads must be divisible by n_kv_heads"
+        self.rep = self.nq // self.nkv
+        qkv_out = (self.nq + 2 * self.nkv) * self.hd
+        self.qkv = nn.Linear(cfg.d_model, qkv_out, bias=False)
+        self.proj = nn.Linear(self.nq * self.hd, cfg.d_model, bias=False)
+        if cfg.qk_norm:
+            self.q_norm = RMSNorm(self.hd)
+            self.k_norm = RMSNorm(self.hd)
+
+    def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+        batch, seq_len, _ = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.split(
+            [self.nq * self.hd, self.nkv * self.hd, self.nkv * self.hd], dim=-1
+        )
+        q = q.view(batch, seq_len, self.nq, self.hd).transpose(1, 2)
+        k = k.view(batch, seq_len, self.nkv, self.hd).transpose(1, 2)
+        v = v.view(batch, seq_len, self.nkv, self.hd).transpose(1, 2)
+        if self.cfg.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        k = k.repeat_interleave(self.rep, dim=1)
+        v = v.repeat_interleave(self.rep, dim=1)
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        out = out.transpose(1, 2).reshape(batch, seq_len, self.nq * self.hd)
+        return self.proj(out)
+
+
+class DenseFFN(nn.Module):
+    """SwiGLU feed-forward."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.w13 = nn.Linear(cfg.d_model, 2 * cfg.dense_ffn, bias=False)
+        self.w2 = nn.Linear(cfg.dense_ffn, cfg.d_model, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        gate, up = self.w13(x).chunk(2, dim=-1)
+        return self.w2(F.silu(gate) * up)
+
+
+class TransformerBlock(nn.Module):
+    """Standard pre-norm block for Prelude / Coda (no loop, no LTI params)."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.d_model)
+        self.attn = Attention(cfg)
+        self.ffn_norm = RMSNorm(cfg.d_model)
+        self.ffn = DenseFFN(cfg)
+
+    def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class ParcaeLoopBlock(nn.Module):
+    """One iteration of the recurrent unit. Called T times by ParcaeTransformer."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.attn_norm = RMSNorm(cfg.d_model)
+        self.attn = Attention(cfg)
+        self.ffn_norm = RMSNorm(cfg.d_model)
+        self.ffn = DenseFFN(cfg)
+        d_model = cfg.d_model
+        if cfg.use_a_matrix:
+            self.A_log = nn.Parameter(torch.zeros(d_model))
+            target_decay = math.sqrt(1.0 / 5.0)
+            target_dt = -math.log(target_decay)
+            dt_bias_init = math.log(math.expm1(target_dt))
+            self.dt_bias = nn.Parameter(torch.full((d_model,), dt_bias_init))
+            self.B = nn.Parameter(torch.eye(d_model))
+
+    def _lti_step(self, h: Tensor, e: Tensor) -> Tensor:
+        dt = F.softplus(self.dt_bias)
+        A = torch.exp(self.A_log)
+        decay = torch.exp(-dt * A)
+        input_gain = dt
+        input_write = F.linear(e, self.B)
+        return decay * h + input_gain * input_write
+
+    def forward(self, h: Tensor, e: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+        if self.cfg.use_a_matrix:
+            y = self._lti_step(h, e)
+        else:
+            y = h + e
+        y = y + self.attn(self.attn_norm(y), cos, sin)
+        y = y + self.ffn(self.ffn_norm(y))
+        return y
+
+
+class ParcaeTransformer(nn.Module):
+    """Prelude + recurrent loop + coda."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        d_model = cfg.d_model
+        self.embed = nn.Embedding(cfg.vocab_size, d_model)
+        self.prelude = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_prelude)])
+        self.prelude_norm = RMSNorm(d_model) if cfg.use_input_norm else nn.Identity()
+        self.loop = ParcaeLoopBlock(cfg)
+        self.coda = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_coda)])
+        self.final_norm = RMSNorm(d_model)
+        self.lm_head = nn.Linear(d_model, cfg.vocab_size, bias=False)
+        if cfg.tie_embeddings:
+            self.lm_head.weight = self.embed.weight
+        self.h0_std = 0.02
+        self._rope_cache: dict = {}
+        self.apply(self._init)
+
+    def _init(self, module: nn.Module):
+        std = self.cfg.init_std
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+            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=std)
+
+    def rope(self, seq_len: int, device, dtype):
+        key = (seq_len, device, dtype)
+        if key not in self._rope_cache:
+            cos, sin = precompute_rope(self.cfg.head_dim, seq_len, self.cfg.rope_theta, device)
+            self._rope_cache[key] = (cos.to(dtype), sin.to(dtype))
+        return self._rope_cache[key]
+
+    def _run_prelude(self, idx: Tensor):
+        x = self.embed(idx)
+        cos, sin = self.rope(x.size(1), x.device, x.dtype)
+        for blk in self.prelude:
+            x = blk(x, cos, sin)
+        return self.prelude_norm(x), cos, sin
+
+    def _run_coda(self, h: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+        for blk in self.coda:
+            h = blk(h, cos, sin)
+        return self.final_norm(h)
+
+    def _h0(self, batch: int, seq_len: int, device, dtype) -> Tensor:
+        return torch.randn(batch, seq_len, self.cfg.d_model, device=device, dtype=dtype) * self.h0_std
+
+    def forward(
+        self,
+        idx: Tensor,
+        targets: Tensor | None = None,
+        T_per_seq: Tensor | None = None,
+        n_no_grad: int = 0,
+    ):
+        e, cos, sin = self._run_prelude(idx)
+        batch, seq_len = idx.shape
+
+        if T_per_seq is None:
+            T_per_seq = torch.full((batch,), self.cfg.mu_rec, device=idx.device, dtype=torch.long)
+        t_max = int(T_per_seq.max().item())
+        n_no_grad = min(n_no_grad, t_max)
+
+        h = self._h0(batch, seq_len, idx.device, e.dtype)
+
+        with torch.no_grad():
+            for t in range(n_no_grad):
+                active = (t < T_per_seq).view(batch, 1, 1).to(h.dtype)
+                h_new = self.loop(h, e, cos, sin)
+                h = active * h_new + (1.0 - active) * h
+
+        for t in range(n_no_grad, t_max):
+            active = (t < T_per_seq).view(batch, 1, 1).to(h.dtype)
+            h_new = self.loop(h, e, cos, sin)
+            h = active * h_new + (1.0 - active) * h
+
+        h = self._run_coda(h, cos, sin)
+
+        if targets is None:
+            return self.lm_head(h), h.new_zeros(())
+
+        logits = self.lm_head(h).float()
+        ce = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        z = (torch.logsumexp(logits, dim=-1) ** 2).mean()
+        loss = ce + self.cfg.final_z_loss_coef * z
+        return loss, {"ce": ce.detach(), "z": z.detach()}
+
+
+def count_params(model: ParcaeTransformer) -> dict:
+    cfg = model.cfg
+    total = sum(p.numel() for p in model.parameters())
+    embed = cfg.vocab_size * cfg.d_model
+    return {"total": total, "non_embed": total - embed, "embed": embed}
+
+
+# =============================================================================
+# HuggingFace wrapper
+# =============================================================================
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+try:
+    from .configuration_looplm import LoopLMConfig
+except ImportError:
+    from configuration_looplm import LoopLMConfig
+
+
+@dataclass
+class _MC:
+    """Internal ModelConfig adapter so ParcaeTransformer can be constructed from HF config."""
+
+    vocab_size: int
+    d_model: int
+    n_prelude: int
+    n_coda: int
+    mu_rec: int
+    n_q_heads: int
+    n_kv_heads: int
+    head_dim: int
+    qk_norm: bool
+    rope_theta: float
+    dense_ffn: int
+    tie_embeddings: bool
+    final_z_loss_coef: float
+    use_a_matrix: bool
+    use_input_norm: bool
+    init_std: float
+
+
+class LoopLMForCausalLM(PreTrainedModel):
+    config_class = LoopLMConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _tied_weights_keys = {"model.lm_head.weight": "model.embed.weight"}
+
+    @property
+    def all_tied_weights_keys(self):
+        return self._tied_weights_keys
+
+    def __init__(self, config: LoopLMConfig):
+        super().__init__(config)
+        mc = _MC(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            n_prelude=config.n_prelude,
+            n_coda=config.n_coda,
+            mu_rec=config.mu_rec,
+            n_q_heads=config.n_q_heads,
+            n_kv_heads=config.n_kv_heads,
+            head_dim=config.head_dim,
+            qk_norm=config.qk_norm,
+            rope_theta=config.rope_theta,
+            dense_ffn=config.dense_ffn,
+            tie_embeddings=config.tie_embeddings,
+            final_z_loss_coef=config.final_z_loss_coef,
+            use_a_matrix=config.use_a_matrix,
+            use_input_norm=config.use_input_norm,
+            init_std=config.init_std,
+        )
+        self.model = ParcaeTransformer(mc)
+
+    def tie_weights(self, *args, **kwargs):
+        """Tie input and output embeddings (called by HF on load).
+
+        Accepts *args, **kwargs to be compatible with various transformers versions
+        that may pass missing_keys, recompute_mapping, etc.
+        """
+        if hasattr(self.model, "embed") and hasattr(self.model, "lm_head"):
+            self.model.lm_head.weight = self.model.embed.weight
+
+    def forward(self, input_ids, labels=None, attention_mask=None, **kwargs):
+        # Always compute logits so HF Trainer can use them in train and eval.
+        logits, _ = self.model(input_ids, None)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)).float(),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        return CausalLMOutput(loss=loss, logits=logits)
+
+    def get_input_embeddings(self):
+        return self.model.embed
+
+    def set_input_embeddings(self, new_embeddings):
+        self.model.embed = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.model.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.model.lm_head = new_embeddings

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}