Upload TinyStories Transformer LM (from-scratch implementation)

bfafe8d verified 2 months ago

2.61 kB

---
license: mit
language: en
tags:
- transformer
- language-model
- from-scratch
- tinystories
---
# 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](https://huggingface.co/datasets/roneneldan/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](https://github.com/flbjcy/LM-from-scratch).

```python
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](https://github.com/flbjcy/LM-from-scratch) for full implementation details.

	---
	license: mit
	language: en
	tags:
	- transformer
	- language-model
	- from-scratch
	- tinystories
	---

	# 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](https://huggingface.co/datasets/roneneldan/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](https://github.com/flbjcy/LM-from-scratch).

	```python
	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](https://github.com/flbjcy/LM-from-scratch) for full implementation details.