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.

Code / reproduction: github.com/puranikyashaswinsharma/reex-1

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

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "puranikyashaswinsharma/reex-1", subfolder="hf_format"
)
tokenizer = AutoTokenizer.from_pretrained(
    "puranikyashaswinsharma/reex-1", subfolder="hf_format"
)

prompt = "### User:\nWhat is the capital of France?\n\n### Assistant:\n"
ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
out = model.generate(ids, max_new_tokens=80, do_sample=True, temperature=0.8, top_k=40, repetition_penalty=1.3)
print(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

• Mechanistic interpretability: circuit discovery via activation patching / EAP-IG
• TransformerLens HookedTransformer compatibility bridge
• Broader SFT data mix to improve factual consistency
• Investigate and fix LAMBADA harness scoring issue
• Larger pretraining run if compute allows

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