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:
2026-07-12 10:24:52 -04:00
commit 5a7f2177bc
5 changed files with 290 additions and 0 deletions
+50
View File
@@ -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()
+71
View File
@@ -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()