Upload Rune-R1 GRPO (RLVR) 351M reasoning model checkpoint

README.md

@@ -0,0 +1,78 @@
+---
+language:
+- en
+license: apache-2.0
+tags:
+- llm
+- pytorch
+- causal-lm
+- rune-r1
+- reasoning
+- grpo
+- rlvr
+datasets:
+- HuggingFaceFW/fineweb-edu
+metrics:
+- accuracy
+---
+
+# Rune-R1 (351M) — GRPO Reasoning Model
+
+**Rune-R1** is a ~351M parameter decoder-only transformer trained from scratch, then
+aligned for math reasoning via a Pretrain -> SFT -> GRPO pipeline:
+
+1. **Pretraining**: 5.05B tokens of [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) (final loss 2.999).
+2. **SFT**: supervised fine-tuning on math reasoning traces.
+3. **GRPO (RLVR)**: Group Relative Policy Optimization with PPO-style clipping and
+   KL-to-reference regularization, using a verifiable reward on math answer correctness
+   (reasoning-from-scratch style recipe). Trained for 2000 steps.
+
+This checkpoint is the final GRPO policy from that last stage.
+
+## Model Details
+
+- **Architecture**: Decoder-only Transformer
+- **Parameters**: ~351M
+- **Layers**: 22
+- **Embedding Dimension**: 1024
+- **Attention Heads / KV Groups**: 16 / 4 (Grouped-Query Attention)
+- **Feed-Forward Hidden Dim**: 2816 (SwiGLU)
+- **Position Embeddings**: RoPE (Rotary Position Embeddings)
+- **Normalization**: RMSNorm (with QK Normalization)
+- **Context Length**: 1024 tokens
+- **Tokenizer**: GPT-2 (`tiktoken`)
+
+## Training (GRPO stage)
+
+- **Steps**: 2000
+- **Reward signal**: rule-based verifier reward on math answer correctness (RLVR)
+- **Reference model**: frozen SFT checkpoint (KL penalty against drift)
+- **Eval**: MATH-500 held-out set (50-example subset), evaluated every 100 steps
+
+MATH-500 accuracy fluctuated in the 0-4% range over training (peak 4% at steps 1600
+and 1900), reflecting the small model size and limited RL budget rather than a fully
+converged reasoning model. Mean reward per step across training was ~0.016, with
+occasional higher-reward rollouts (max single-step average 0.75).
+
+## Usage
+
+You can load and generate text with this model using the `rune` package in this repository:
+
+```python
+import torch
+import tiktoken
+from rune.model import CONFIG_350M, RuneModel
+
+# Load model weights
+ckpt = torch.load("pytorch_model.bin", map_location="cpu")
+model = RuneModel(CONFIG_350M)
+model.load_state_dict(ckpt)
+model.eval()
+
+# Encode prompt
+enc = tiktoken.get_encoding("gpt2")
+prompt = "The key to machine learning is"
+tokens = torch.tensor([enc.encode(prompt)], dtype=torch.long)
+
+# Generation logic using model(tokens)
+```

config.json

@@ -0,0 +1,13 @@
+{
+  "vocab_size": 50257,
+  "context_length": 1024,
+  "emb_dim": 1024,
+  "n_heads": 16,
+  "n_layers": 22,
+  "hidden_dim": 2816,
+  "head_dim": null,
+  "qk_norm": true,
+  "n_kv_groups": 4,
+  "rope_base": 10000.0,
+  "dtype": "torch.float32"
+}

model.py

@@ -0,0 +1,215 @@
+# Architecture adapted from rasbt/LLMs-from-scratch pkg/llms_from_scratch/qwen3.py (Apache 2.0):
+# RoPE + RMSNorm + SwiGLU + grouped-query attention, trimmed to a dense ~350M config
+# with a GPT-2 (tiktoken) vocab instead of Qwen's tokenizer/MoE variants.
+import torch
+import torch.nn as nn
+
+CONFIG_350M = {
+    "vocab_size": 50257,        # tiktoken gpt2
+    "context_length": 1024,
+    "emb_dim": 1024,
+    "n_heads": 16,
+    "n_layers": 22,
+    "hidden_dim": 2816,
+    "head_dim": None,           # defaults to emb_dim // n_heads
+    "qk_norm": True,
+    "n_kv_groups": 4,
+    "rope_base": 10_000.0,
+    "dtype": torch.float32,     # fp32 master weights; train.py autocasts to bf16 for compute
+}
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, emb_dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.ones(emb_dim))
+
+    def forward(self, x):
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        norm_x = x * torch.rsqrt(variance + self.eps) * self.scale
+        return norm_x.to(input_dtype)
+
+
+def compute_rope_params(head_dim, theta_base, context_length, dtype=torch.float32):
+    assert head_dim % 2 == 0, "Head dimension must be even"
+    inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2, dtype=dtype) / head_dim))
+    positions = torch.arange(context_length, dtype=dtype)
+    angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)
+    angles = torch.cat([angles, angles], dim=1)
+    return torch.cos(angles), torch.sin(angles)
+
+
+def apply_rope(x, cos, sin, offset=0):
+    # x: (batch, heads, seq_len, head_dim). `offset` is the absolute position
+    # of x[..., 0, :] — nonzero when x is a new chunk appended after cached
+    # positions, so rotation angles pick up where the cache left off.
+    head_dim = x.shape[-1]
+    x1, x2 = x[..., : head_dim // 2], x[..., head_dim // 2:]
+    seq_len = x.shape[2]
+    max_pos = cos.shape[0]
+    if offset + seq_len > max_pos:
+        offset = max(0, max_pos - seq_len)
+    cos = cos[offset:offset + seq_len].unsqueeze(0).unsqueeze(0)
+    sin = sin[offset:offset + seq_len].unsqueeze(0).unsqueeze(0)
+    rotated = torch.cat((-x2, x1), dim=-1)
+    return ((x * cos) + (rotated * sin)).to(dtype=x.dtype)
+
+
+def new_kv_cache(n_layers):
+    """One mutable dict per layer; GroupedQueryAttention fills in 'k'/'v' and
+    grows them in place across calls sharing the same cache list."""
+    return [dict() for _ in range(n_layers)]
+
+
+class GroupedQueryAttention(nn.Module):
+    def __init__(self, d_in, num_heads, num_kv_groups, head_dim=None, qk_norm=False, dtype=None):
+        super().__init__()
+        assert num_heads % num_kv_groups == 0, "num_heads must be divisible by num_kv_groups"
+        if head_dim is None:
+            assert d_in % num_heads == 0
+            head_dim = d_in // num_heads
+
+        self.num_heads = num_heads
+        self.num_kv_groups = num_kv_groups
+        self.group_size = num_heads // num_kv_groups
+        self.head_dim = head_dim
+        self.d_out = num_heads * head_dim
+
+        self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)
+        self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)
+        self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)
+        self.out_proj = nn.Linear(self.d_out, d_in, bias=False, dtype=dtype)
+
+        self.q_norm = RMSNorm(head_dim) if qk_norm else None
+        self.k_norm = RMSNorm(head_dim) if qk_norm else None
+
+    def forward(self, x, mask, cos, sin, cache=None):
+        b, num_tokens, _ = x.shape
+
+        queries = self.W_query(x).view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
+        keys = self.W_key(x).view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
+        values = self.W_value(x).view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
+
+        if self.q_norm:
+            queries = self.q_norm(queries)
+        if self.k_norm:
+            keys = self.k_norm(keys)
+
+        past_len = 0 if cache is None or cache.get("k") is None else cache["k"].shape[2]
+        queries = apply_rope(queries, cos, sin, offset=past_len)
+        keys = apply_rope(keys, cos, sin, offset=past_len)
+
+        if cache is not None:
+            if cache.get("k") is not None:
+                keys = torch.cat([cache["k"], keys], dim=2)
+                values = torch.cat([cache["v"], values], dim=2)
+            cache["k"], cache["v"] = keys, values
+
+        keys = keys.repeat_interleave(self.group_size, dim=1)
+        values = values.repeat_interleave(self.group_size, dim=1)
+
+        if past_len == 0:
+            # No cache, or first (prefill) call on an empty cache: query and
+            # key spans are identical, standard causal mask applies.
+            context = nn.functional.scaled_dot_product_attention(
+                queries, keys, values, attn_mask=None, is_causal=True
+            )
+        elif num_tokens == 1:
+            # Single-token decode step: this query is always the newest
+            # position, so it may attend to every cached key — no mask needed.
+            context = nn.functional.scaled_dot_product_attention(
+                queries, keys, values, attn_mask=None, is_causal=False
+            )
+        else:
+            raise NotImplementedError("cache only supports prefill-then-single-token decode")
+
+        context = context.transpose(1, 2).reshape(b, num_tokens, self.d_out)
+        return self.out_proj(context)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.fc1 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], dtype=cfg["dtype"], bias=False)
+        self.fc2 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], dtype=cfg["dtype"], bias=False)
+        self.fc3 = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], dtype=cfg["dtype"], bias=False)
+
+    def forward(self, x):
+        return self.fc3(nn.functional.silu(self.fc1(x)) * self.fc2(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.att = GroupedQueryAttention(
+            d_in=cfg["emb_dim"], num_heads=cfg["n_heads"], head_dim=cfg["head_dim"],
+            num_kv_groups=cfg["n_kv_groups"], qk_norm=cfg["qk_norm"], dtype=cfg["dtype"],
+        )
+        self.ff = FeedForward(cfg)
+        self.norm1 = RMSNorm(cfg["emb_dim"])
+        self.norm2 = RMSNorm(cfg["emb_dim"])
+
+    def forward(self, x, mask, cos, sin, cache=None):
+        x = x + self.att(self.norm1(x), mask, cos, sin, cache)
+        x = x + self.ff(self.norm2(x))
+        return x
+
+
+class RuneModel(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"], dtype=cfg["dtype"])
+        self.trf_blocks = nn.ModuleList(TransformerBlock(cfg) for _ in range(cfg["n_layers"]))
+        self.final_norm = RMSNorm(cfg["emb_dim"])
+        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False, dtype=cfg["dtype"])
+
+        head_dim = cfg["head_dim"] or cfg["emb_dim"] // cfg["n_heads"]
+        cos, sin = compute_rope_params(head_dim, cfg["rope_base"], cfg["context_length"])
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+
+    def forward(self, in_idx, cache=None):
+        x = self.tok_emb(in_idx)
+        for i, block in enumerate(self.trf_blocks):
+            x = block(x, None, self.cos, self.sin, cache[i] if cache is not None else None)
+        x = self.final_norm(x)
+        return self.out_head(x.to(self.cfg["dtype"]))
+
+
+def _test_kv_cache_matches_full_forward(cfg):
+    torch.manual_seed(0)
+    model = RuneModel(cfg).eval()
+    seq = torch.randint(0, cfg["vocab_size"], (2, 12))
+
+    with torch.no_grad():
+        full_logits = model(seq)
+
+        cache = new_kv_cache(cfg["n_layers"])
+        chunks = [model(seq[:, :5], cache=cache)]
+        for i in range(5, 12):
+            chunks.append(model(seq[:, i:i + 1], cache=cache))
+        cached_logits = torch.cat(chunks, dim=1)
+
+    assert cached_logits.shape == full_logits.shape
+    max_diff = (full_logits - cached_logits).abs().max().item()
+    assert torch.allclose(full_logits, cached_logits, atol=1e-4), f"max diff {max_diff}"
+    print(f"kv-cache self-test ok (max diff vs full forward: {max_diff:.2e})")
+
+
+if __name__ == "__main__":
+    cfg = CONFIG_350M
+    model = RuneModel(cfg)
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"params: {n_params:,} ({n_params / 1e6:.1f}M)")
+
+    x = torch.randint(0, cfg["vocab_size"], (2, 16))
+    logits = model(x)
+    assert logits.shape == (2, 16, cfg["vocab_size"]), logits.shape
+    assert torch.isfinite(logits).all()
+    print("forward pass ok:", logits.shape)
+
+    _test_kv_cache_matches_full_forward(dict(cfg, n_layers=2, context_length=64))

pytorch_model.bin

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