flbjcy/LM-from-scratch

TinyStories Transformer LM (From Scratch)

A small Transformer language model, implemented entirely from scratch (no torch.nn layers beyond Parameter and containers) as part of Stanford's CS336 (Language Modeling from Scratch) assignment.

Model Details

• Architecture: Decoder-only Transformer with RoPE positional embeddings, causal multi-head self-attention, SwiGLU feed-forward layers, and RMSNorm (pre-norm configuration)
• Parameters: ~22.7M
• Vocabulary: 10,000 tokens (custom-trained byte-level BPE tokenizer)
• Context length: 256 tokens
• Training data: TinyStories (full training set, ~540M tokens)
• Training steps: 5,000
• Final validation loss: ~1.80 (cross-entropy)

What this model can (and can't) do

This model was trained only on TinyStories — a dataset of simple children's stories. It is a base language model, not an instruction-tuned assistant. It can generate coherent, story-like text continuations (especially when prompted with story-like openings such as "Once upon a time..."), but it has no capability for question-answering, factual knowledge, instruction-following, or conversation. Prompting it with anything outside the TinyStories domain will typically produce fluent but nonsensical output, since it has only ever learned to continue text in that narrow style.

Usage

This model uses a custom architecture (not directly compatible with transformers' AutoModel). To use it, you'll need the original implementation from this repo.

import torch, pickle, json
from cs336_basics.model import TransformerLM, decode
from cs336_basics.tokenizer import Tokenizer

with open("config.json") as f:
    config = json.load(f)
with open("vocab.pkl", "rb") as f:
    vocab = pickle.load(f)
with open("merges.pkl", "rb") as f:
    merges = pickle.load(f)

tokenizer = Tokenizer(vocab, merges, special_tokens=["<|endoftext|>"])
model = TransformerLM(**config)
model.load_state_dict(torch.load("model_weights.pt"))

output = decode(model, tokenizer, "Once upon a time", max_tokens=100, temperature=0.8, top_p=0.9)
print(output)

Training

Trained from scratch on a MacBook Air (Apple Silicon, MPS), using a custom implementation of AdamW, cosine learning rate scheduling with warmup, and gradient clipping. See the source repository for full implementation details.