summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--scripts/run_sweep.py3
-rw-r--r--src/config.py2
-rw-r--r--src/data.py17
-rw-r--r--src/eval.py90
-rw-r--r--src/model_api.py8
-rw-r--r--src/train.py132
-rw-r--r--tests/test_cuda_scaling.py71
7 files changed, 245 insertions, 78 deletions
diff --git a/scripts/run_sweep.py b/scripts/run_sweep.py
index 9be5592..34cccf7 100644
--- a/scripts/run_sweep.py
+++ b/scripts/run_sweep.py
@@ -22,7 +22,8 @@ import time
from concurrent.futures import ThreadPoolExecutor, as_completed
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-PY = os.path.join(ROOT, ".venv", "bin", "python")
+_venv_py = os.path.join(ROOT, ".venv", "bin", "python")
+PY = _venv_py if os.path.exists(_venv_py) else sys.executable
def parse_jobs(path: str) -> list[dict]:
diff --git a/src/config.py b/src/config.py
index 60984db..25dcdda 100644
--- a/src/config.py
+++ b/src/config.py
@@ -34,9 +34,11 @@ class Config:
early_stop_em: float = 1.0
early_stop_patience: int = 5
batch_size: int = 32
+ eval_batch_size: int = 512 # chunked evaluation batch size to avoid GPU OOM
max_out_len: int = 6
log_n_examples: int = 10
out_dir: str = "runs"
+ device: str = "auto" # "auto" | "cuda" | "cpu" | "mps"
@property
def vocab(self) -> int:
diff --git a/src/data.py b/src/data.py
index fab4c21..26fb6b9 100644
--- a/src/data.py
+++ b/src/data.py
@@ -1,4 +1,5 @@
"""Prime dataset: n -> next prime, digit-tokenized; splits and batching."""
+import bisect
import random
import torch
@@ -20,9 +21,9 @@ def sieve_primes(limit: int) -> list[int]:
def next_prime(n: int, primes: list[int]) -> int:
- for p in primes:
- if p > n:
- return p
+ idx = bisect.bisect_right(primes, n)
+ if idx < len(primes):
+ return primes[idx]
raise ValueError(f"no prime > {n} in supplied list")
@@ -75,7 +76,8 @@ def get_splits(cfg: Config) -> tuple[list[int], list[int]]:
def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list[int]]]:
"""[(input_tokens, target_tokens+EOS), ...]. task_mode selects the target function."""
- primes = sieve_primes(cfg.range_end + 100)
+ margin = max(100, int(cfg.range_end * 0.05) + 50)
+ primes = sieve_primes(cfg.range_end + margin) if cfg.task_mode != "is_prime" else []
out = []
for n in inputs:
if cfg.task_mode == "is_prime":
@@ -88,10 +90,13 @@ def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list
def _global_lengths(cfg: Config) -> tuple[int, int]:
"""(in_max, out_max): fixed global lengths so batch layout == singleton layout (codex BLOCKER fix)."""
- primes = sieve_primes(cfg.range_end + 100)
- max_target = next_prime(cfg.range_end, primes)
if cfg.vocab_mode == "integers":
return 1, 2 # [value], [value, EOS]
+ if cfg.task_mode == "is_prime":
+ return len(str(cfg.range_end)), 2 # digit token "1"/"0" + EOS
+ margin = max(100, int(cfg.range_end * 0.05) + 50)
+ primes = sieve_primes(cfg.range_end + margin)
+ max_target = next_prime(cfg.range_end, primes)
return len(str(cfg.range_end)), len(str(max_target)) + 1 # digits + EOS
diff --git a/src/eval.py b/src/eval.py
index e93f95d..1cb1191 100644
--- a/src/eval.py
+++ b/src/eval.py
@@ -26,31 +26,46 @@ FLAGGED_SIEVE_PREDS = {121, 143, 169, 187, 209}
def probe_report(model, cfg: Config, lo: int = 101, hi: int = 200) -> dict:
- primes = sieve_primes(hi + 200)
+ try:
+ device = next(model.parameters()).device
+ except StopIteration:
+ device = torch.device("cpu")
+ margin = max(200, int(hi * 0.1) + 50)
+ primes = sieve_primes(hi + margin)
correct = 0
errors = []
easy_total = 0
easy_wrong = 0
is_prime_task = cfg.task_mode == "is_prime"
- for n in range(lo, hi + 1):
- x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
- gen = greedy_decode(model, x, cfg)[0].tolist()
- pred = decode_tokens(gen, cfg)
- if is_prime_task:
- target = 1 if is_prime_n(n) else 0
- ok = (pred == 1) == (target == 1) # any non-"1" output reads as "composite"
- else:
- target = next_prime(n, primes)
- ok = pred == target
- is_easy = (n % 2 == 0) or (n % 5 == 0) # trivial composites (skip-evens / skip-5s)
- if is_easy:
- easy_total += 1
- if ok:
- correct += 1
- else:
- errors.append({"n": n, "target": target, "pred": pred})
+ eval_bs = getattr(cfg, "eval_batch_size", 512)
+ inputs = list(range(lo, hi + 1))
+
+ for bi in range(0, len(inputs), eval_bs):
+ chunk = inputs[bi: bi + eval_bs]
+ xs = [encode_int(n, cfg) for n in chunk]
+ in_max = max(len(xi) for xi in xs)
+ x_tensor = torch.full((len(chunk), in_max), cfg.pad_id, dtype=torch.long, device=device)
+ for i, xi in enumerate(xs):
+ x_tensor[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long, device=device)
+
+ gen = greedy_decode(model, x_tensor, cfg).cpu()
+ for i, n in enumerate(chunk):
+ pred = decode_tokens(gen[i].tolist(), cfg)
+ if is_prime_task:
+ target = 1 if is_prime_n(n) else 0
+ ok = (pred == 1) == (target == 1) # any non-"1" output reads as "composite"
+ else:
+ target = next_prime(n, primes)
+ ok = pred == target
+ is_easy = (n % 2 == 0) or (n % 5 == 0) # trivial composites (skip-evens / skip-5s)
if is_easy:
- easy_wrong += 1
+ easy_total += 1
+ if ok:
+ correct += 1
+ else:
+ errors.append({"n": n, "target": target, "pred": pred})
+ if is_easy:
+ easy_wrong += 1
total = hi - lo + 1
acc = correct / total
if is_prime_task:
@@ -124,14 +139,30 @@ def grokking_signature(metrics_path: str) -> dict:
@torch.no_grad()
def halting_report(model, cfg: Config) -> dict:
"""RNN halting structure: mean steps at run end + correlation with gap-to-next-prime (H1-H4)."""
- primes = sieve_primes(300)
+ try:
+ device = next(model.parameters()).device
+ except StopIteration:
+ device = torch.device("cpu")
+ margin = max(100, int(cfg.range_end * 0.05) + 50)
+ primes = sieve_primes(max(300, cfg.range_end + margin))
gaps, steps = [], []
- for n in range(cfg.range_start, cfg.range_end + 1):
- x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0)
- h = model._encode(x)
+ inputs = list(range(cfg.range_start, cfg.range_end + 1))
+ eval_bs = getattr(cfg, "eval_batch_size", 512)
+ from src.data import pad_inputs
+ for bi in range(0, len(inputs), eval_bs):
+ chunk = inputs[bi: bi + eval_bs]
+ xs = [encode_int(n, cfg) for n in chunk]
+ in_max = max(len(xi) for xi in xs)
+ x_tensor = torch.full((len(chunk), in_max), cfg.pad_id, dtype=torch.long, device=device)
+ for i, xi in enumerate(xs):
+ x_tensor[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long, device=device)
+ x_padded = pad_inputs(x_tensor, cfg)
+ h = model._encode(x_padded)
_, s = model._run_compute(h)
- gaps.append(next_prime(n, primes) - n)
- steps.append(float(s.mean()))
+ s_cpu = s.cpu().tolist()
+ for i, n in enumerate(chunk):
+ gaps.append(next_prime(n, primes) - n)
+ steps.append(float(s_cpu[i]))
mean = float(np.mean(steps))
rho = float(np.corrcoef(gaps, steps)[0, 1]) if len(set(gaps)) > 1 else 0.0
lo, hi = cfg.min_steps + 0.5, cfg.max_steps - 0.5
@@ -148,9 +179,12 @@ def halting_report(model, cfg: Config) -> dict:
"gaps": gaps, "steps_by_n": steps}
-def _load(out_dir: str, ckpt: str, cfg: Config):
- model = build_model(cfg)
- model.load_state_dict(torch.load(os.path.join(out_dir, ckpt), map_location="cpu"))
+def _load(out_dir: str, ckpt: str, cfg: Config, device: torch.device | None = None):
+ if device is None:
+ device_str = getattr(cfg, "device", "auto")
+ device = torch.device("cuda" if (device_str == "cuda" or (device_str == "auto" and torch.cuda.is_available())) else "cpu")
+ model = build_model(cfg).to(device)
+ model.load_state_dict(torch.load(os.path.join(out_dir, ckpt), map_location=device))
model.eval()
return model
diff --git a/src/model_api.py b/src/model_api.py
index 31ffe31..85b5d55 100644
--- a/src/model_api.py
+++ b/src/model_api.py
@@ -32,10 +32,16 @@ def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int
Inputs are LEFT-padded to the global layout so positions match training exactly
(batch/singleton invariance — codex BLOCKER fix)."""
+ try:
+ device = next(model.parameters()).device
+ except StopIteration:
+ device = x.device
+ if x.device != device:
+ x = x.to(device)
x = pad_inputs(x, cfg)
max_len = max_len or cfg.max_out_len
B = x.shape[0]
- y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=x.device)
+ y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
for _ in range(max_len):
out = model(x, y_in)
nxt = out["logits"][:, -1].argmax(-1)
diff --git a/src/train.py b/src/train.py
index 406defd..e8870df 100644
--- a/src/train.py
+++ b/src/train.py
@@ -21,38 +21,69 @@ def set_seed(s: int) -> None:
torch.manual_seed(s)
+def resolve_device(device_str: str) -> torch.device:
+ if device_str == "auto":
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ return torch.device(device_str)
+
+
@torch.no_grad()
-def evaluate(model, examples, cfg: Config):
+def evaluate(model, examples, cfg: Config, device: torch.device | None = None):
"""Token accuracy (teacher-forced) + exact-match accuracy (batched greedy) + per-example detail.
- BATCHED: all examples in ONE forward pass (greedy decode batch-wide). Valid because
- models are batch-invariant by construction (global fixed layout, pad no-ops — see
- tests/test_codex_fixes.py). ~50x fewer forwards than the per-example version."""
+ BATCHED & CHUNKED: examples evaluated in chunks of eval_batch_size. Valid because
+ models are batch-invariant by construction (global fixed layout, pad no-ops).
+ Chunking prevents GPU VRAM OOM on large datasets (e.g. range_end = 100k)."""
+ if device is None:
+ try:
+ device = next(model.parameters()).device
+ except StopIteration:
+ device = torch.device("cpu")
model.eval()
- batch = make_batch(examples, cfg)
- out = model(batch["x"], batch["y_in"])
- logits = out["logits"] # (B,T_out,vocab)
- y = batch["y"]
- mask = batch["y_mask"]
- pred_tok = logits.argmax(-1)
- token_correct = int((pred_tok[mask] == y[mask]).sum())
- token_total = int(mask.sum())
- # batched greedy decode
- B = batch["x"].shape[0]
- y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long)
- for _ in range(cfg.max_out_len):
- o = model(batch["x"], y_in)
- nxt = o["logits"][:, -1].argmax(-1)
- y_in = torch.cat([y_in, nxt[:, None]], dim=1)
- gen = y_in[:, 1:] # (B, max_out_len)
+ if not examples:
+ return 0.0, 0.0, []
+
+ eval_bs = getattr(cfg, "eval_batch_size", 512)
+ token_correct = 0
+ token_total = 0
em_correct = 0
per_example = []
- for i, (xrow, target) in enumerate(examples):
- pred = decode_tokens(gen[i].tolist(), cfg)
- target_n = decode_tokens(target, cfg)
- ok = pred == target_n
- em_correct += int(ok)
- per_example.append((tuple(xrow), target_n, pred, ok))
+
+ for bi in range(0, len(examples), eval_bs):
+ chunk = examples[bi: bi + eval_bs]
+ batch = make_batch(chunk, cfg)
+ x = batch["x"].to(device)
+ y = batch["y"].to(device)
+ y_in_eval = batch["y_in"].to(device)
+ mask = batch["y_mask"].to(device)
+
+ out = model(x, y_in_eval)
+ logits = out["logits"] # (B,T_out,vocab)
+ pred_tok = logits.argmax(-1)
+ token_correct += int((pred_tok[mask] == y[mask]).sum().item())
+ token_total += int(mask.sum().item())
+
+ # batched greedy decode on device
+ B = x.shape[0]
+ max_out_len = getattr(cfg, "max_out_len", 6)
+ from src.data import _global_lengths
+ _, out_max = _global_lengths(cfg)
+ max_len = max(max_out_len, out_max)
+
+ y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
+ for _ in range(max_len):
+ o = model(x, y_in)
+ nxt = o["logits"][:, -1].argmax(-1)
+ y_in = torch.cat([y_in, nxt[:, None]], dim=1)
+ gen = y_in[:, 1:].cpu()
+
+ for i, (xrow, target) in enumerate(chunk):
+ pred = decode_tokens(gen[i].tolist(), cfg)
+ target_n = decode_tokens(target, cfg)
+ ok = pred == target_n
+ em_correct += int(ok)
+ per_example.append((tuple(xrow), target_n, pred, ok))
+
model.train()
return token_correct / max(1, token_total), em_correct / len(examples), per_example
@@ -62,7 +93,7 @@ def log_example_rows(log_examples, val_per_example):
by_x = {x: (target, pred, ok) for x, target, pred, ok in val_per_example}
rows = []
for x, _target in log_examples:
- t, p, ok = by_x[tuple(x)]
+ t, p, ok = by_x.get(tuple(x), (-1, -1, False))
label = "".join(map(str, x)) # "42" in both vocab modes
rows.append(f"{label}:{'ok' if ok else f'{p}~{t}'}")
return ";".join(rows)
@@ -71,6 +102,7 @@ def log_example_rows(log_examples, val_per_example):
def main() -> None:
cfg = parse_args()
set_seed(cfg.seed)
+ device = resolve_device(cfg.device)
out_dir = os.path.join(cfg.out_dir, cfg.model, f"seed{cfg.seed}")
csv_path = os.path.join(out_dir, "metrics.csv")
if os.path.exists(csv_path):
@@ -78,24 +110,36 @@ def main() -> None:
f"(reruns would corrupt the CSV and checkpoint provenance).")
os.makedirs(out_dir, exist_ok=True)
cfg.save(os.path.join(out_dir, "config.json"))
+ meta = {
+ "python": sys.version.split()[0],
+ "torch": torch.__version__,
+ "numpy": np.__version__,
+ "device": str(device),
+ "torch_threads": torch.get_num_threads(),
+ "cmd": sys.argv,
+ }
+ if device.type == "cuda":
+ meta["gpu_name"] = torch.cuda.get_device_name(device)
+ meta["gpu_capability"] = list(torch.cuda.get_device_capability(device))
with open(os.path.join(out_dir, "run_meta.json"), "w") as fh:
- json.dump({
- "python": sys.version.split()[0],
- "torch": torch.__version__,
- "numpy": np.__version__,
- "device": "cpu",
- "torch_threads": torch.get_num_threads(),
- "cmd": sys.argv,
- }, fh, indent=2)
+ json.dump(meta, fh, indent=2)
train_in, val_in = get_splits(cfg)
train_ex = build_examples(train_in, cfg)
val_ex = build_examples(val_in, cfg)
- model = build_model(cfg)
- print(f"model={cfg.model} params={model.param_count()} train={len(train_ex)} val={len(val_ex)} "
+ model = build_model(cfg).to(device)
+ print(f"model={cfg.model} params={model.param_count()} device={device} "
+ f"train={len(train_ex)} val={len(val_ex)} "
f"vocab={cfg.vocab} eos={cfg.eos_id} pad={cfg.pad_id}")
- opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+ opt_kwargs = {"lr": cfg.lr, "weight_decay": cfg.weight_decay}
+ if device.type == "cuda":
+ try:
+ opt = torch.optim.AdamW(model.parameters(), fused=True, **opt_kwargs)
+ except Exception:
+ opt = torch.optim.AdamW(model.parameters(), **opt_kwargs)
+ else:
+ opt = torch.optim.AdamW(model.parameters(), **opt_kwargs)
ce = nn.CrossEntropyLoss(reduction="none")
log_examples = sorted(val_ex, key=lambda ex: (len(ex[0]), ex[0]))[: cfg.log_n_examples]
@@ -149,12 +193,16 @@ def main() -> None:
if not sl:
continue
batch = make_batch(sl, cfg)
- out = model(batch["x"], batch["y_in"])
+ x = batch["x"].to(device)
+ y = batch["y"].to(device)
+ y_in = batch["y_in"].to(device)
+ mask = batch["y_mask"].to(device)
+ out = model(x, y_in)
logits = out["logits"]
- y_safe = batch["y"].clamp(max=cfg.vocab - 1) # CE index guard for pad positions
+ y_safe = y.clamp(max=cfg.vocab - 1) # CE index guard for pad positions
loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(
- logits.shape[0], -1) * batch["y_mask"].float()
- loss_tokens = loss_tokens.sum() / batch["y_mask"].sum().clamp(min=1)
+ logits.shape[0], -1) * mask.float()
+ loss_tokens = loss_tokens.sum() / mask.sum().clamp(min=1)
if cfg.halting and step >= cfg.halt_warmup_steps:
if step >= cfg.halt_ramp_end_steps:
lam = cfg.halt_penalty
diff --git a/tests/test_cuda_scaling.py b/tests/test_cuda_scaling.py
new file mode 100644
index 0000000..5f0befc
--- /dev/null
+++ b/tests/test_cuda_scaling.py
@@ -0,0 +1,71 @@
+"""Tests for CUDA device handling, chunked evaluation, and scaling to 1000 and 100_000."""
+import torch
+
+from src.config import Config, parse_args
+from src.data import build_examples, get_splits, next_prime, sieve_primes
+from src.model_api import build_model
+from src.train import evaluate, resolve_device
+
+
+def test_device_resolution():
+ assert resolve_device("cpu") == torch.device("cpu")
+ auto_dev = resolve_device("auto")
+ expected = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ assert auto_dev == expected
+
+
+def test_config_device_args():
+ cfg = parse_args(["rnn", "0", "--device", "cpu", "--eval_batch_size", "256"])
+ assert cfg.device == "cpu"
+ assert cfg.eval_batch_size == 256
+
+
+def test_bisect_next_prime_large_range():
+ primes = sieve_primes(100_100)
+ # Check boundary primes
+ assert next_prime(2, primes) == 3
+ assert next_prime(97, primes) == 101
+ assert next_prime(997, primes) == 1009
+ assert next_prime(99_991, primes) == 100_003
+
+
+def test_chunked_evaluate_matches_single_batch():
+ """Verify that chunked evaluation with eval_batch_size=4 gives identical results to full batch."""
+ cfg = Config(model="rnn", eval_batch_size=4)
+ model = build_model(cfg)
+ train_in, val_in = get_splits(cfg)
+ val_ex = build_examples(val_in, cfg)
+
+ # Chunked
+ tok_chunked, em_chunked, per_chunked = evaluate(model, val_ex, cfg)
+
+ # Single full batch
+ cfg_full = Config(model="rnn", eval_batch_size=len(val_ex))
+ tok_full, em_full, per_full = evaluate(model, val_ex, cfg_full)
+
+ assert tok_chunked == tok_full
+ assert em_chunked == em_full
+ assert per_chunked == per_full
+
+
+def test_range_extension_1000():
+ cfg = Config(range_start=2, range_end=1000, model="transformer")
+ train_in, val_in = get_splits(cfg)
+ assert len(train_in) + len(val_in) == 999
+ train_ex = build_examples(train_in, cfg)
+ assert len(train_ex) == len(train_in)
+ model = build_model(cfg)
+ tok, em, per = evaluate(model, build_examples(val_in[:20], cfg), cfg)
+ assert 0.0 <= em <= 1.0
+
+
+def test_range_extension_100k_data_shapes():
+ """Test data generation and model forward at 100k range without OOM."""
+ cfg = Config(range_start=2, range_end=100_000, model="rnn")
+ # Take a small slice of inputs from 100k range
+ sample_inputs = [2, 42, 999, 50_000, 99_991]
+ examples = build_examples(sample_inputs, cfg)
+ assert len(examples) == 5
+ model = build_model(cfg)
+ tok, em, per = evaluate(model, examples, cfg)
+ assert 0.0 <= em <= 1.0
generated by cgit v1.2.3 (git 2.46.0) at 2026-09-01 05:10:22 +0000