Reex-116M
A ~116M parameter decoder-only transformer, pretrained and chat-tuned from scratch on free compute. This is my first from-scratch language model, trained, debugged, and released end-to-end as a learning project and a base for future mechanistic interpretability work.
What this is (and isn't)
Reex holds a coherent short conversation, gets simple factual questions right some of the time, and follows chat formatting reliably. It is
not competitive with production chatbots. It struggles with arithmetic, multi-step reasoning, and long-range coherence, and its factual grounding is inconsistent. That's the expected and honest ceiling for a 116M-parameter model trained on ~2B tokens of free GPU time, not a bug in the pipeline. See
Evaluation for real, unfiltered sample outputs and benchmark scores.
The goal of this project was to build and understand the entire pipeline (architecture, pretraining, checkpointing, fine-tuning, and release) from scratch, and to produce a model simple and standard enough to later support mechanistic interpretability work (circuit discovery, activation patching, probing).
Architecture
Llama-style decoder-only transformer:
| Component | Choice |
|---|
| Layers | 12 |
| Hidden size | 768 |
| Attention heads | 12 (query), 4 (key/value, grouped-query attention) |
| Normalization | RMSNorm |
| MLP | SwiGLU |
| Positional encoding | RoPE |
| Vocabulary | GPT-2 BPE (50,257 tokens) |
| Context length | 1024 |
| Parameters | 115,883,520 |
The GPT-2 tokenizer and standard decoder-only architecture were deliberate choices for compatibility with existing interpretability tooling (TransformerLens, nnsight). A HF-native LlamaForCausalLM-compatible checkpoint is also provided under hf_format/ in this repo, verified to produce numerically identical logits to the original training checkpoint (see Conversion Notes below).
Training
Pretraining:
- Data: FineWeb-Edu (10BT sample), streamed
- ~2B tokens, sequence length 1024, mixed precision (fp16)
- 20,000 steps, batch size 8 x grad accumulation 8 (effective batch 64)
- Hardware: free-tier T4 GPU (16GB)
- Final training loss: ~3.9-4.2
Chat fine-tuning (SFT):
- Data: OpenAssistant OASST2, English single-turn pairs, filtered
- 3 epochs, ~13,650 steps, LR 1e-5
- Final avg loss: 2.95
Evaluation
Zero-shot benchmarks
Run with
EleutherAI's lm-evaluation-harness, the same tool behind the HF Open LLM Leaderboard. Compared against published numbers for similarly-sized models (
source):
| Benchmark | Reex-116M | GPT-2 Small (124M) | Pythia-160M | OPT-125M |
|---|
| ARC-Challenge | 23.0% | 22.6% | 23.1% | 22.1% |
| ARC-Easy | 35.9% | 39.7% | 36.4% | 39.9% |
| HellaSwag | 26.1% | 31.4% | 30.3% | 31.6% |
| PIQA | 56.1% | 62.1% | 59.8% | 62.0% |
| WinoGrande | 51.5% | 50.7% | 50.8% | 51.8% |
Reex is competitive with established 100-160M baselines on ARC-Challenge and WinoGrande, and somewhat behind on ARC-Easy, HellaSwag, and PIQA, consistent with training on fewer tokens than these longer-established baselines. LAMBADA is excluded from this table due to a harness scoring anomaly under investigation; manual inspection confirms the model produces sensible next-token predictions on long-context completions (see repo for details).
Qualitative samples
Unfiltered outputs from the SFT checkpoint, shown as-is including failures:
| Prompt | Output | Notes |
|---|
| "What is the capital of France?" | "The capital of France is Paris." | Correct |
| "Hi, how are you?" | "Hello there, I'm doing well. If you have any questions or concerns about the topic then please feel free to ask!" | Coherent, well-formatted |
| "Can you help me plan a birthday party?" | Structured numbered list; content partially on-topic | Format correct, content wanders |
| "What is 15 plus 27?" | Nonsensical, unrelated to arithmetic | Known limitation |
| "Write a short poem about the ocean." | Deflects rather than completing the task | Known limitation |
| "Tell me a joke." | Produces a sentence, not an actual joke | Known limitation |
Full sample set in the
GitHub repo.
Usage
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4model = AutoModelForCausalLM.from_pretrained(
5 "puranikyashaswinsharma/reex-1", subfolder="hf_format"
6)
7tokenizer = AutoTokenizer.from_pretrained(
8 "puranikyashaswinsharma/reex-1", subfolder="hf_format"
9)
10
11prompt = "### User:\nWhat is the capital of France?\n\n### Assistant:\n"
12ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
13out = model.generate(ids, max_new_tokens=80, do_sample=True, temperature=0.8, top_k=40, repetition_penalty=1.3)
14print(tokenizer.decode(out[0], skip_special_tokens=True))
Or clone the
GitHub repo for the original custom architecture + local chat script.
Conversion notes
This repo includes both the original training-format checkpoints (checkpoint_step20000.pt, checkpoint_sft_final.pt) and an HF-native LlamaForCausalLM conversion (hf_format/). The conversion required permuting the Q/K projection weights to translate between interleaved-RoPE (the convention used during training) and split-half-RoPE (HF's Llama convention), verified numerically identical (max logit difference 0.0000 on test inputs) after the permutation fix.
Roadmap
Acknowledgments
Built on free compute. Pretraining data from
FineWeb-Edu (HuggingFaceFW). SFT data from
OASST2 (OpenAssistant/LAION). Architecture inspired by LLaMA and modern open decoder-only LM designs.
License
MIT
