Phase 0: project plan, pinned environment, data download + MPS smoke test
- README with full pipeline plan (tokenizer -> data -> model -> training -> eval) - pip/venv setup pinned to torch 2.13.0, numpy 2.5.1 (Python 3.13) - scripts/download_data.py: one-time TinyStoriesV2 fetch with SHA-256 recording - scripts/check_mps.py: verifies MPS backend, backward pass, GPU speedup
This commit is contained in:
@@ -0,0 +1,6 @@
|
||||
.venv/
|
||||
__pycache__/
|
||||
*.pyc
|
||||
data/
|
||||
checkpoints/
|
||||
.DS_Store
|
||||
@@ -0,0 +1,161 @@
|
||||
# experiment-llm
|
||||
|
||||
Training a small language model completely from scratch on a MacBook Pro (M1):
|
||||
own tokenizer, own transformer implementation, own training loop. No
|
||||
pretrained weights, no model libraries, no cloud compute.
|
||||
|
||||
The goal of **iteration 1** is *completeness and clarity*, not model quality:
|
||||
a small model (~4–5M parameters) taken through the entire pipeline end to end,
|
||||
with every component written to be read and understood.
|
||||
|
||||
## Principles
|
||||
|
||||
- **From scratch.** The tokenizer, model, training loop, and sampler are
|
||||
implemented in this repo. Dependencies are limited to PyTorch (tensor ops +
|
||||
autograd + MPS backend) and NumPy. No HuggingFace libraries, no nanoGPT
|
||||
imports — reference implementations are for reading, not importing.
|
||||
- **Self-sufficient.** After a one-time dataset download, everything runs
|
||||
offline on the local machine. No API calls, no cloud GPUs, no external
|
||||
evaluation services.
|
||||
- **Clarity over speed.** Straightforward code beats clever code. Every phase
|
||||
produces an inspectable artifact (tokenizer vocab, token dumps, loss curves,
|
||||
sample outputs).
|
||||
- **Config-driven.** One experiment = one config. Model size, context length,
|
||||
and training schedule live in a config file, so iteration 2 is a config
|
||||
change, not a rewrite.
|
||||
|
||||
## Setup
|
||||
|
||||
Requires Python 3.13 (Homebrew) on Apple Silicon.
|
||||
|
||||
```sh
|
||||
python3 -m venv .venv
|
||||
source .venv/bin/activate
|
||||
pip install -r requirements.txt
|
||||
|
||||
python scripts/check_mps.py # verify the Apple GPU is usable
|
||||
python scripts/download_data.py # one-time, ~2.2 GB into data/raw/
|
||||
```
|
||||
|
||||
After the download, everything runs offline.
|
||||
|
||||
## Why TinyStories
|
||||
|
||||
General web text is hopeless at this scale — a 5M-parameter model trained on
|
||||
it produces word salad. [TinyStories](https://arxiv.org/abs/2305.07759)
|
||||
(Eldan & Li, 2023) is a ~2GB corpus of synthetic children's stories with a
|
||||
deliberately restricted vocabulary, built precisely so that very small models
|
||||
can learn coherent English. It is the smallest known regime where a
|
||||
from-scratch model produces output that *makes sense* — grammatical sentences,
|
||||
consistent characters, simple plots. Our v1 model sits below the paper's
|
||||
coherence sweet spot (~30M params); some incoherence is expected and fine.
|
||||
|
||||
## Pipeline
|
||||
|
||||
```
|
||||
raw text ──► tokenizer training ──► tokenized corpus ──► pretraining ──► sampling / eval
|
||||
(phase 1) (phase 2) (phase 2) (phase 3) (phase 4)
|
||||
```
|
||||
|
||||
### Phase 0 — Setup
|
||||
- Standard `venv` + `pip` environment; dependencies (`torch`, `numpy`) pinned
|
||||
to exact versions in `requirements.txt`.
|
||||
- One-time download of TinyStoriesV2 (plain `.txt` files fetched over HTTPS
|
||||
into `data/raw/`, which is gitignored). This is the only network access in
|
||||
the whole project.
|
||||
- Deliverable: `scripts/download_data.py`, pinned environment, smoke test that
|
||||
MPS is available.
|
||||
|
||||
### Phase 1 — Tokenizer
|
||||
- Byte-pair encoding implemented from scratch: trainer, encoder, decoder.
|
||||
- Trained on a subsample of the corpus (BPE training is quadratic-ish in pure
|
||||
Python; a ~50MB sample is plenty for a 4k vocab on this data).
|
||||
- Vocab size 4,096 — TinyStories' restricted vocabulary doesn't need more,
|
||||
and a small vocab keeps the embedding table (and the model) small.
|
||||
- Deliverables: `src/tokenizer.py`, saved vocab/merges file, round-trip tests
|
||||
(`decode(encode(s)) == s`), a short report on compression ratio
|
||||
(bytes per token).
|
||||
|
||||
### Phase 2 — Data pipeline
|
||||
- Encode the full corpus once with the trained tokenizer into flat binary
|
||||
token files (`data/tokens/train.bin`, `val.bin`) — memory-mapped at training
|
||||
time, so the 2GB corpus never has to fit in RAM as Python objects.
|
||||
- Batching: random crops of `context_length + 1` tokens; inputs/targets are
|
||||
shifted views of the same crop.
|
||||
- Deliverables: `src/dataset.py`, `scripts/encode_corpus.py`, token counts
|
||||
logged for train/val.
|
||||
|
||||
### Phase 3 — Model
|
||||
- Decoder-only transformer, GPT-style, implemented module by module:
|
||||
token embeddings, causal self-attention, MLP block, RMSNorm, residual
|
||||
stream, weight-tied output head.
|
||||
- v1 configuration (~4.3M params):
|
||||
|
||||
| knob | value |
|
||||
|---|---|
|
||||
| layers | 4 |
|
||||
| d_model | 256 |
|
||||
| heads | 4 |
|
||||
| context length | 256 |
|
||||
| vocab | 4,096 |
|
||||
| positions | learned embeddings (v1 simplicity; RoPE is an iteration-2 candidate) |
|
||||
|
||||
- Deliverables: `src/model.py`, a parameter-count printout that matches the
|
||||
hand-computed estimate, a forward-pass shape test.
|
||||
|
||||
### Phase 4 — Training
|
||||
- Plain training loop: AdamW, cosine LR schedule with warmup, gradient
|
||||
clipping, periodic validation loss, checkpointing (resumable).
|
||||
- Runs on MPS; expected wall-clock for v1 is a few hours for ~100–200M tokens
|
||||
(roughly Chinchilla-optimal for this size — we don't need to see the whole
|
||||
corpus).
|
||||
- Loss curve logged to a plain CSV and plotted locally.
|
||||
- Deliverables: `src/train.py`, `configs/v1.yaml` (or `.py`), checkpoints in
|
||||
`checkpoints/` (gitignored), loss curves in `reports/`.
|
||||
|
||||
### Phase 5 — Sampling & evaluation
|
||||
- Autoregressive sampler with temperature and top-k.
|
||||
- Evaluation is fully local:
|
||||
- validation loss / perplexity,
|
||||
- a fixed prompt suite (story openings) sampled at every checkpoint so
|
||||
progress is visible qualitatively over training,
|
||||
- a sample gallery in `reports/` for the final model.
|
||||
- Deliverables: `src/sample.py`, `reports/samples-v1.md`.
|
||||
|
||||
## Project layout
|
||||
|
||||
```
|
||||
experiment-llm/
|
||||
├── README.md
|
||||
├── requirements.txt
|
||||
├── configs/ # one file per experiment (v1, v2, ...)
|
||||
├── scripts/ # one-shot entry points: download, encode, plots
|
||||
├── src/
|
||||
│ ├── tokenizer.py # BPE: train / encode / decode
|
||||
│ ├── dataset.py # memmapped token files, batch sampling
|
||||
│ ├── model.py # the transformer
|
||||
│ ├── train.py # training loop
|
||||
│ └── sample.py # generation
|
||||
├── tests/ # round-trip + shape + overfit-one-batch tests
|
||||
├── data/ # gitignored: raw/ (txt), tokens/ (bin)
|
||||
├── checkpoints/ # gitignored
|
||||
└── reports/ # loss curves, sample galleries (committed)
|
||||
```
|
||||
|
||||
## Definition of done (iteration 1)
|
||||
|
||||
- [ ] Environment reproducible with `pip install -r requirements.txt`; dataset downloaded and hashed
|
||||
- [ ] BPE tokenizer trains, round-trips, and is saved to disk
|
||||
- [ ] Corpus encoded to binary token files
|
||||
- [ ] Model builds; parameter count matches the estimate
|
||||
- [ ] Overfit-one-batch sanity check passes (loss → ~0 on a single batch)
|
||||
- [ ] Full training run completes on MPS with checkpoints + loss curve
|
||||
- [ ] Sampler produces stories from prompts; gallery committed to `reports/`
|
||||
- [ ] Every phase re-runnable from a clean clone with documented commands
|
||||
|
||||
## Iteration 2 candidates (out of scope for now)
|
||||
|
||||
- Scale toward the coherence sweet spot (~15–30M params, longer training)
|
||||
- RoPE positions, SwiGLU MLP, other modern architecture upgrades
|
||||
- Instruction tuning on TinyStories-Instruct (pretrain → SFT arc)
|
||||
- MLX port for faster training on Apple Silicon
|
||||
@@ -0,0 +1,2 @@
|
||||
torch==2.13.0
|
||||
numpy==2.5.1
|
||||
@@ -0,0 +1,50 @@
|
||||
"""Smoke test: PyTorch can see the Apple GPU and run a training step on it.
|
||||
|
||||
Checks three things:
|
||||
1. the MPS backend is available,
|
||||
2. forward + backward passes produce finite gradients on MPS,
|
||||
3. MPS matmul is actually faster than CPU (i.e., the GPU is really used).
|
||||
|
||||
Usage: python scripts/check_mps.py
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def bench_matmul(device: str, n: int = 2048, iters: int = 20) -> float:
|
||||
"""Mean seconds per (n x n) matmul; .item() forces device sync."""
|
||||
x = torch.randn(n, n, device=device)
|
||||
y = torch.randn(n, n, device=device)
|
||||
for _ in range(3): # warmup
|
||||
(x @ y).sum().item()
|
||||
start = time.perf_counter()
|
||||
for _ in range(iters):
|
||||
(x @ y).sum().item()
|
||||
return (time.perf_counter() - start) / iters
|
||||
|
||||
|
||||
def main() -> None:
|
||||
print(f"torch {torch.__version__}")
|
||||
assert torch.backends.mps.is_available(), "MPS backend not available"
|
||||
assert torch.backends.mps.is_built(), "torch was built without MPS support"
|
||||
print("MPS backend: available")
|
||||
|
||||
# A miniature training step: forward, loss, backward, finite grads.
|
||||
w = torch.randn(64, 64, device="mps", requires_grad=True)
|
||||
x = torch.randn(128, 64, device="mps")
|
||||
loss = ((x @ w) ** 2).mean()
|
||||
loss.backward()
|
||||
assert w.grad is not None and torch.isfinite(w.grad).all()
|
||||
print(f"backward pass on MPS: ok (loss={loss.item():.4f})")
|
||||
|
||||
cpu = bench_matmul("cpu")
|
||||
mps = bench_matmul("mps")
|
||||
print(f"2048x2048 matmul: cpu {cpu * 1e3:.1f} ms | mps {mps * 1e3:.1f} ms"
|
||||
f" | speedup {cpu / mps:.1f}x")
|
||||
print("smoke test passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,71 @@
|
||||
"""One-time download of the TinyStoriesV2 corpus.
|
||||
|
||||
This is the only network access in the entire project. Files land in
|
||||
data/raw/ (gitignored). SHA-256 hashes are recorded in data/raw/SHA256SUMS
|
||||
so the corpus can be verified later without re-downloading.
|
||||
|
||||
Usage: python scripts/download_data.py
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
BASE_URL = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main"
|
||||
FILES = [
|
||||
"TinyStoriesV2-GPT4-train.txt", # ~2.2 GB
|
||||
"TinyStoriesV2-GPT4-valid.txt", # ~22 MB
|
||||
]
|
||||
RAW_DIR = Path(__file__).resolve().parent.parent / "data" / "raw"
|
||||
CHUNK = 1 << 20 # 1 MB
|
||||
|
||||
|
||||
def download(name: str) -> Path:
|
||||
dest = RAW_DIR / name
|
||||
if dest.exists():
|
||||
print(f"{name}: already present, skipping download")
|
||||
return dest
|
||||
tmp = dest.with_suffix(dest.suffix + ".part")
|
||||
url = f"{BASE_URL}/{name}"
|
||||
print(f"{name}: downloading")
|
||||
with urllib.request.urlopen(url) as response, open(tmp, "wb") as f:
|
||||
total = int(response.headers.get("Content-Length") or 0)
|
||||
done = 0
|
||||
while chunk := response.read(CHUNK):
|
||||
f.write(chunk)
|
||||
done += len(chunk)
|
||||
if total:
|
||||
print(
|
||||
f"\r {done / 1e6:,.0f} / {total / 1e6:,.0f} MB"
|
||||
f" ({done * 100 // total}%)",
|
||||
end="",
|
||||
flush=True,
|
||||
)
|
||||
print()
|
||||
tmp.rename(dest) # only a fully written file ever gets the real name
|
||||
return dest
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
while chunk := f.read(CHUNK):
|
||||
h.update(chunk)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
RAW_DIR.mkdir(parents=True, exist_ok=True)
|
||||
lines = []
|
||||
for name in FILES:
|
||||
path = download(name)
|
||||
digest = sha256(path)
|
||||
print(f"{name}: {path.stat().st_size / 1e6:,.0f} MB sha256={digest}")
|
||||
lines.append(f"{digest} {name}\n")
|
||||
sums = RAW_DIR / "SHA256SUMS"
|
||||
sums.write_text("".join(lines))
|
||||
print(f"hashes recorded in {sums}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user