Initial release: LoopLM-135M-naive trained on FineWeb 4.6B tokens

README.md

@@ -1,3 +1,96 @@
 ---
 license: apache-2.0
+tags:
+- looped-transformer
+- language-model
+- research
+- from-scratch
+language:
+- en
+library_name: transformers
+datasets:
+- HuggingFaceFW/fineweb
 ---
+
+# LoopLM-135M-naive
+
+A 135M parameter **dense looped transformer** trained from scratch on FineWeb. Built as part of an exploration of looped LLM architectures inspired by the Parcae paper ([arXiv:2604.12946](https://arxiv.org/abs/2604.12946)).
+
+Trained on Modal H100s as a hobby research project. This is the **naive looped** variant — a clean baseline without Parcae's LTI stability mechanisms, which were found to underperform at this scale.
+
+## Architecture
+
+- **Type:** Dense looped transformer (recurrent reuse of one block)
+- **Total params:** 135M (134.1M unique trainable)
+- **d_model:** 1024
+- **Prelude:** 4 transformer blocks
+- **Loop:** 1 transformer block, called T ~ Poisson(μ=6) times per sequence
+- **Coda:** 2 transformer blocks
+- **Attention:** GQA (16 query heads / 8 KV heads), head_dim=64, RoPE θ=10000, QK-norm
+- **FFN:** SwiGLU, dense_ffn=2816
+- **Vocab:** 50304 (GPT-2 BPE + padding), tied embeddings
+
+## Training
+
+- **Dataset:** FineWeb (raw, ~4.6B tokens consumed)
+- **Hardware:** 2× H100 on Modal
+- **Steps:** 17,500 (stopped early at checkpoint to fit budget)
+- **Batch:** 262,144 tokens/step
+- **Optimizer:** Muon (matrices) + AdamW (norms, embeddings)
+- **LR:** Muon 0.02, AdamW 3e-4, cosine decay (60% constant + 40% decay to 0.1×)
+- **Precision:** bf16 with fp32 logits
+
+## Results
+
+| Model | Architecture | Tokens | Val Loss |
+|---|---|---|---|
+| LoopLM-30M (prior) | Dense (8 layers) | 1B | 3.91 |
+| **LoopLM-135M-naive (this)** | **Dense looped** | **4.6B** | **3.95** |
+| HobbyLM-130M-MoE (sibling) | MoE (140M total / 62M active) | 10B | 3.30 |
+
+## Honest Findings
+
+This project explored several Parcae stability mechanisms (A matrix constraints, prelude input normalization, identity-init B matrix, custom optimizer routing). **None outperformed naive looping at this scale.** The paper's improvements appear to be most relevant at 1B+ parameters and 100B+ tokens, while at 135M / 4.6B tokens, simple recurrent reuse is competitive enough.
+
+Compared to a sibling 130M MoE trained on more data, the dense looped model loses by ~0.65 in val loss — MoE wins on token-efficiency at this scale.
+
+## Usage
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+tokenizer = AutoTokenizer.from_pretrained("harims95/LoopLM-135M-naive")
+model = AutoModelForCausalLM.from_pretrained(
+    "harims95/LoopLM-135M-naive",
+    trust_remote_code=True,
+    torch_dtype=torch.float32,
+)
+model.eval()
+
+# Forward pass for next-token logits
+inputs = tokenizer("The capital of France is", return_tensors="pt")
+with torch.no_grad():
+    out = model(**inputs)
+next_token_logits = out.logits[0, -1]
+next_token = tokenizer.decode(next_token_logits.argmax())
+print(next_token)
+```
+
+Note: a HuggingFace-style `.generate()` is not wired in this release. The model returns logits and you'd write a simple sampling loop manually. Training was focused on the architecture exploration, not deployment.
+
+## Reproducibility
+
+Training code: [github.com/harims95/LoopLM](https://github.com/harims95/LoopLM)
+
+The `spec.json` in this repo captures the exact training config used for this checkpoint.
+
+## Acknowledgments
+
+- Built on top of [harishsg993010/HobbyLM](https://github.com/harishsg993010/HobbyLM) (training infrastructure, optimizer setup, data pipeline)
+- Architecture inspired by [Parcae](https://arxiv.org/abs/2604.12946) — Scaling Laws For Stable Looped Language Models
+- Trained on [FineWeb](https://huggingface.co/datasets/HuggingFaceFW/fineweb)
+
+## License
+
+Apache 2.0 — feel free to use, study, or build on. Acknowledgment appreciated but not required.

config.json

@@ -0,0 +1,26 @@
+{
+  "model_type": "looplm",
+  "architectures": ["LoopLMForCausalLM"],
+  "auto_map": {
+    "AutoConfig": "configuration_looplm.LoopLMConfig",
+    "AutoModelForCausalLM": "modeling_looplm.LoopLMForCausalLM"
+  },
+  "vocab_size": 50304,
+  "d_model": 1024,
+  "n_prelude": 4,
+  "n_coda": 2,
+  "mu_rec": 6,
+  "n_q_heads": 16,
+  "n_kv_heads": 8,
+  "head_dim": 64,
+  "qk_norm": true,
+  "rope_theta": 10000.0,
+  "dense_ffn": 2816,
+  "tie_embeddings": true,
+  "final_z_loss_coef": 0.0001,
+  "use_a_matrix": false,
+  "use_input_norm": false,
+  "init_std": 0.02,
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.0"
+}

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)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3a57a0942fd460fc31260a2cdccac52c5d31d80c7c32fdaf48394673184a6bb
+size 742462848

modeling_looplm.py

@@ -0,0 +1,315 @@
+"""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
+
+    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 forward(self, input_ids, labels=None, **kwargs):
+        if labels is not None:
+            loss, _ = self.model(input_ids, labels)
+            logits = None
+        else:
+            logits, _ = self.model(input_ids, None)
+            loss = None
+        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

spec.json

@@ -0,0 +1,78 @@
+{
+  "timestamp": "2026-06-28T11:27:05.869434+00:00",
+  "run_name": "real_naive_fineweb_5B_2gpu",
+  "git_commit": "unknown",
+  "cli_args": {
+    "preset": "135M",
+    "run_name": "real_naive_fineweb_5B_2gpu",
+    "data_dir": "/data/fineweb",
+    "train_pattern": "fineweb_train_*.bin",
+    "val_pattern": "fineweb_train_*.bin",
+    "max_steps": 20000,
+    "seq_len": 1024,
+    "batch_tokens": 262144,
+    "micro_batch_seqs": 32,
+    "val_every": 250,
+    "out_dir": "/data/runs",
+    "save_every": 2500,
+    "no_compile": false,
+    "holdout_last_for_val": true,
+    "set": [
+      "use_a_matrix=false",
+      "use_input_norm=false"
+    ]
+  },
+  "train_config": {
+    "data_dir": "/data/fineweb_edu",
+    "train_pattern": "edu_fineweb_train_*.bin",
+    "val_pattern": "edu_fineweb_val_*.bin",
+    "seq_len": 1024,
+    "batch_tokens": 262144,
+    "micro_batch_seqs": 32,
+    "max_steps": 20000,
+    "warmup_steps": 100,
+    "cooldown_frac": 0.4,
+    "final_lr_frac": 0.1,
+    "muon_lr": 0.02,
+    "muon_momentum": 0.95,
+    "muon_wd": 0.1,
+    "muon_ns_steps": 5,
+    "adam_lr": 0.0003,
+    "adam_betas": [
+      0.9,
+      0.95
+    ],
+    "adam_wd": 0.1,
+    "grad_clip": 1.0,
+    "val_every": 250,
+    "val_tokens": 10485760,
+    "log_every": 10,
+    "seed": 1337,
+    "compile": true,
+    "bf16": true,
+    "out_dir": "/data/runs",
+    "run_name": "real_naive_fineweb_5B_2gpu"
+  },
+  "model_config": {
+    "vocab_size": 50304,
+    "d_model": 1024,
+    "n_prelude": 4,
+    "n_coda": 2,
+    "mu_rec": 6,
+    "n_q_heads": 16,
+    "n_kv_heads": 8,
+    "head_dim": 64,
+    "qk_norm": true,
+    "rope_theta": 10000.0,
+    "dense_ffn": 2816,
+    "tie_embeddings": true,
+    "final_z_loss_coef": 0.0001,
+    "use_a_matrix": false,
+    "use_input_norm": false,
+    "init_std": 0.02
+  },
+  "hostname": "modal",
+  "gpu_count": 2,
+  "gpu_type": "NVIDIA H100 80GB HBM3",
+  "pytorch_version": "2.12.0+cu130"
+}

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": null,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}