rzen
5a7f2177bc
- 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
51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
"""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()
|