transformer-pretraining-from-scratch

transformer-pretraining-from-scratch

A ~57M-non-embedding-parameter decoder-only transformer, pretrained from scratch in PyTorch. Every nn.Module (RMSNorm, rotary position embeddings, causal self-attention, SwiGLU MLP, weight-tied embedding/output head) is authored in the training repo: no AutoModel, no fine-tune of an existing checkpoint. Trained on ~1B tokens of FineWeb-Edu with a custom byte-level BPE tokenizer, also trained from scratch rather than reusing GPT-2's.

Full training code, tokenizer, evaluation, calibration study, interpretability pass, and writeups: GitHub repo.

Try it live, no install needed: the model runs entirely in your browser via ONNX Runtime Web.

Why this exists

This model's companion project, a cost-aware LLM router, calls a pretrained transformer through a high-level API and asks when to trust its confidence. This project builds the model itself and asks where that confidence actually comes from, mechanistically: attention patterns, calibration behavior, what a causal ablation reveals about which heads matter. It isn't trying to be state of the art, instruction-tuned, or chat-capable.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "hari-8/transformer-pretraining-from-scratch"
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(repo)

ids = tokenizer("The history of the Roman Empire", return_tensors="pt").input_ids
out = model.generate_simple(ids, max_new_tokens=40, temperature=0.8, top_k=40)
print(tokenizer.decode(out[0].tolist()))

trust_remote_code=True is required: the architecture (modeling_decoder_transformer.py) ships in this repo rather than being a stock transformers model class. model.generate_simple(...) is a minimal temperature/top-k sampling loop, not HF's KV-cached .generate(). It recomputes the full sequence each step, which is fine for short demo completions but not optimized for long generations or serving at scale.

Architecture

Training

• ~1B tokens, FineWeb-Edu sample
• Single NVIDIA L4 GPU, bf16, torch.compile, gradient accumulation
• 7,630 steps, cosine LR schedule with warmup
• Final training loss: 3.16

Evaluation

Held-out loss, perplexity, and bits-per-byte (BPB), measured on ~497K held-out tokens never seen during training. BPB, not raw perplexity, is the metric that's actually comparable across models with different tokenizers, since it normalizes by the byte length of the original text rather than by token count.

This model scores lower BPB than Pythia-70M on this held-out set. The likely reason is domain match rather than general capability: it was trained exclusively on FineWeb-Edu, a narrower and more predictable distribution than Pythia's Pile-trained generalist scope, and the two models are close in parameter count. Full methodology: src/eval/ in the GitHub repo.

Calibration

Token-level calibration (confidence is the max softmax probability, the label is whether that top-1 prediction matches the actual next token, the standard Guo et al. 2017 definition), measured on ~248K held-out tokens disjoint from the temperature-fitting split.

The model's raw confidence is already close to perfectly calibrated: the fitted temperature is barely different from 1. Full write-up: results/calibration_report.md in the GitHub repo.

Interpretability

Induction-head probing (Olsson et al. 2022 methodology) found two clear induction heads concentrated in layers 6-7, with prefix-matching scores of 0.79 and 0.74, far above every other head. A full causal ablation sweep across every head in every layer found that these aren't necessarily the most causally important heads on natural text, which points to a real limit of attention-pattern-only probing. Full write-up: results/interpretability_report.md in the GitHub repo.

Limitations

• Not instruction-tuned. This is a raw next-token-prediction language model, not a chat assistant, so expect continuation-style completions rather than question-answering or dialogue.
• Small (~57M params) and trained on ~1B tokens, so short-range text stays coherent but long-form factual accuracy or reasoning isn't reliable.
• generate_simple's decoding has no KV cache, so it's O(n²) and slow for long generations.
• Custom architecture, so it requires trust_remote_code=True to load.

Citation

If this is useful, please link back to the GitHub repo instead of citing it formally. This is a learning project, not a paper.