Update README.md

README.md

@@ -12,67 +12,141 @@ tags:
 - rlvr
 datasets:
 - HuggingFaceFW/fineweb-edu
+- rasbt/math_distill
 metrics:
 - accuracy
+pipeline_tag: text-generation
 ---
 
 # 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:
+**Rune-R1** is a ~351M parameter decoder-only transformer trained from scratch and
+aligned for math reasoning via a three-stage 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.
+```
+Pretrain (FineWeb-Edu)  →  SFT (distilled CoT format)  →  GRPO (RLVR on math correctness)
+```
 
-This checkpoint is the final GRPO policy from that last stage.
+This repository holds the final checkpoint: the GRPO-tuned policy, starting from
+[Rune-R1-SFT](https://huggingface.co/samueljayasingh/Rune-R1-sft) and optimized
+with Group Relative Policy Optimization against a verifiable, rule-based reward for
+math answer correctness. See [Rune-R1-Base](https://huggingface.co/samueljayasingh/Rune-R1-base)
+and [Rune-R1-SFT](https://huggingface.co/samueljayasingh/Rune-R1-sft) for the earlier
+pipeline stages.
 
-## Model Details
+## Model Description
 
-- **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
+- **Developed by**: samueljayasingh
+- **Model type**: Causal language model (text-only)
+- **Base model**: [Rune-R1-SFT](https://huggingface.co/samueljayasingh/Rune-R1-sft) (351M, chain-of-thought SFT on top of Rune-R1-Base)
+- **Fine-tuning method**: GRPO (Group Relative Policy Optimization) with PPO-style clipping and a KL penalty to a frozen reference (RLVR — reinforcement learning from verifiable rewards)
+- **Dataset**: `data/math_train.json` (math word problems with verifiable final answers), evaluated on a 50-example MATH-500 held-out subset
+- **Language**: English
 - **Tokenizer**: GPT-2 (`tiktoken`)
+- **License**: Apache 2.0
 
-## Training (GRPO stage)
+### Architecture details
 
-- **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
+| Parameter | Value |
+|---|---|
+| Layers | 22 |
+| Embedding dimension | 1024 |
+| Attention heads / KV groups | 16 / 4 (GQA) |
+| Feed-forward hidden dim | 2816 (SwiGLU) |
+| Context length | 1024 tokens |
+| Position embeddings | RoPE (base 10,000) |
+| Normalization | RMSNorm, with QK normalization |
 
-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).
+## Intended Uses & Limitations
 
-## Usage
+**Intended use:**
+- Research into RLVR / GRPO-style reasoning fine-tuning at small model scale.
+- Reference implementation for reward-verified RL post-training pipelines (pretrain → SFT → RL).
+- Studying reward hacking, KL-regularization tradeoffs, and reasoning-accuracy dynamics under a small RL step budget.
 
-You can load and generate text with this model using the `rune` package in this repository:
+**Limitations:**
+- Small model (351M) with a limited RL budget (2,000 steps) — MATH-500 accuracy
+  remains low (0–4% across evaluation checkpoints; see table below) and should not
+  be compared to production-scale reasoning models.
+- Reward signal is a rule-based correctness check (`\boxed{}` extraction + symbolic
+  grading), so the model may still learn to produce well-formatted but incorrect
+  reasoning that occasionally reward-hacks the verifier.
+- Inherits base/SFT limitations: 1024-token context, ~5B pretraining tokens, no
+  broad safety/RLHF alignment beyond the math-correctness reward.
+- Not suitable for production or user-facing deployment — this is a research
+  artifact demonstrating the training pipeline, not a competitive reasoning model.
+
+## How to Use
 
 ```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"
+prompt = "What is 12 * 15?"
 tokens = torch.tensor([enc.encode(prompt)], dtype=torch.long)
 
-# Generation logic using model(tokens)
+# Model responds in "<think>...reasoning...</think>\n\n\\boxed{final_answer}" format.
+# See rune/generate.py in the source repo for full sampling / KV-cache generation code.
 ```
+
+The `rune` package (model definition + generation utilities) is available at the
+[Rune-R1 GitHub repository](https://github.com/samueljayasingh/Rune-R1).
+
+## Training & Evaluation
+
+### Training procedure
+
+| Parameter | Value |
+|---|---|
+| Starting checkpoint | Rune-R1-SFT |
+| Reference model | Frozen copy of the SFT checkpoint (KL penalty target) |
+| Training steps | 2,000 |
+| Rollouts per prompt (group size) | 8 |
+| Inner epochs per rollout batch | 2 |
+| Max new tokens (rollout) | 512 |
+| Sampling temperature / top-p | 0.8 / 0.9 |
+| PPO clip epsilon | 10.0 |
+| KL coefficient | 0.001 |
+| Learning rate | 1e-6 |
+| Reward function | Rule-based: extract `\boxed{}` answer, symbolically grade vs. ground truth (1.0 / 0.0) |
+| Eval cadence | MATH-500 (50-example subset), every 100 steps |
+| Hardware | 1x AMD MI300X GPU |
+
+### Evaluation results
+
+| Metric | Value |
+|---|---|
+| Final MATH-500 accuracy (step 2000) | 0% (50 examples) |
+| Peak MATH-500 accuracy | 4% (steps 1600, 1900) |
+| Mean reward per step (over training) | ~0.016 |
+| Max single-step average reward | 0.75 |
+| Steps with nonzero reward | 152 / 2001 |
+
+MATH-500 accuracy fluctuated in the 0–4% range throughout training rather than
+improving monotonically, reflecting the small model size and limited RL budget
+rather than a fully converged reasoning model.
+
+## Citation
+
+```bibtex
+@misc{RuneR12026,
+  author = {Samuel Jayasingh},
+  title = {Rune-R1: A 351M Transformer Reasoning Model Trained via Pretrain-SFT-GRPO from Scratch},
+  year = {2026},
+  publisher = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/samueljayasingh/Rune-R1}}
+}
+```
+
+## Acknowledgements
+
+- [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) — pretraining corpus.
+- [rasbt/math_distill](https://huggingface.co/datasets/rasbt/math_distill) — distilled chain-of-thought SFT data.
+- [rasbt/LLMs-from-scratch](https://github.com/rasbt/LLMs-from-scratch) — architecture and the pretrain→SFT→GRPO reasoning-from-scratch recipe this pipeline is adapted from.