reex-1

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:

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):

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:

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