diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/config.py | 1 | ||||
| -rw-r--r-- | src/train.py | 178 |
2 files changed, 106 insertions, 73 deletions
diff --git a/src/config.py b/src/config.py index 4db66b0..945c259 100644 --- a/src/config.py +++ b/src/config.py @@ -31,6 +31,7 @@ class Config: weight_decay: float = 1.0 max_train_steps: int = 200_000 eval_every: int = 200 + eval_schedule: str = "fixed" # "fixed" | "adaptive" (more frequent early, coarser later for 4M runs) early_stop_em: float = 1.0 early_stop_patience: int = 5 batch_size: int = 32 diff --git a/src/train.py b/src/train.py index ec7bb4e..d8c03de 100644 --- a/src/train.py +++ b/src/train.py @@ -27,20 +27,26 @@ def resolve_device(device_str: str) -> torch.device: return torch.device(device_str) -@torch.no_grad() -def evaluate(model, examples, cfg: Config, device: torch.device | None = None): - """Token accuracy (teacher-forced) + exact-match accuracy (batched greedy) + per-example detail. +def make_gpu_dataset(examples: list, cfg: Config, device: torch.device) -> dict[str, torch.Tensor | int]: + """Create static tensors directly in VRAM to eliminate per-step host-to-device transfers.""" + if not examples: + return {"size": 0} + batch = make_batch(examples, cfg) + return { + "x": batch["x"].to(device, non_blocking=True), + "y": batch["y"].to(device, non_blocking=True), + "y_in": batch["y_in"].to(device, non_blocking=True), + "y_mask": batch["y_mask"].to(device, non_blocking=True), + "size": len(examples), + } - 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") + +@torch.no_grad() +def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int], examples: list, cfg: Config, device: torch.device): + """Batched evaluation operating directly on pre-allocated VRAM tensors.""" model.eval() - if not examples: + N = gpu_ds["size"] + if N == 0: return 0.0, 0.0, [] eval_bs = getattr(cfg, "eval_batch_size", 512) @@ -50,22 +56,20 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None): em_correct = 0 per_example = [] - 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) + for bi in range(0, N, eval_bs): + x = gpu_ds["x"][bi: bi + eval_bs] + y = gpu_ds["y"][bi: bi + eval_bs] + y_in_eval = gpu_ds["y_in"][bi: bi + eval_bs] + mask = gpu_ds["y_mask"][bi: bi + eval_bs] + chunk_examples = examples[bi: bi + eval_bs] with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): out = model(x, y_in_eval) - logits = out["logits"] # (B,T_out,vocab) + logits = out["logits"] 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 @@ -80,7 +84,7 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None): y_in = torch.cat([y_in, nxt[:, None]], dim=1) gen = y_in[:, 1:].cpu() - for i, (xrow, target) in enumerate(chunk): + for i, (xrow, target) in enumerate(chunk_examples): pred = decode_tokens(gen[i].tolist(), cfg) target_n = decode_tokens(target, cfg) ok = pred == target_n @@ -88,7 +92,19 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None): per_example.append((tuple(xrow), target_n, pred, ok)) model.train() - return token_correct / max(1, token_total), em_correct / len(examples), per_example + return token_correct / max(1, token_total), em_correct / N, per_example + + +@torch.no_grad() +def evaluate(model, examples, cfg: Config, device: torch.device | None = None): + """Token accuracy + exact-match accuracy + per-example detail (backward compatible wrapper).""" + if device is None: + try: + device = next(model.parameters()).device + except StopIteration: + device = torch.device("cpu") + gpu_ds = make_gpu_dataset(examples, cfg, device) + return evaluate_gpu(model, gpu_ds, examples, cfg, device) def log_example_rows(log_examples, val_per_example): @@ -133,6 +149,9 @@ def main() -> None: train_in, val_in = get_splits(cfg) train_ex = build_examples(train_in, cfg) val_ex = build_examples(val_in, cfg) + train_gpu = make_gpu_dataset(train_ex, cfg, device) + val_gpu = make_gpu_dataset(val_ex, cfg, device) + model = build_model(cfg).to(device) if cfg.compile_model and hasattr(torch, "compile") and device.type == "cuda": try: @@ -163,16 +182,30 @@ def main() -> None: patience_left = cfg.early_stop_patience best_path = os.path.join(out_dir, "best.pt") - n_batches = math.ceil(len(train_ex) / cfg.batch_size) - rng = random.Random(cfg.seed) + N_train = train_gpu["size"] + B = min(cfg.batch_size, N_train) if N_train > 0 else 1 step = 0 done = False first_row = True + def should_eval(cur_step: int) -> bool: + if cur_step >= cfg.max_train_steps: + return True + if getattr(cfg, "eval_schedule", "fixed") in ("adaptive", "logarithmic"): + if cur_step < 10_000: + return cur_step % 200 == 0 + elif cur_step < 50_000: + return cur_step % 1_000 == 0 + elif cur_step < 500_000: + return cur_step % 5_000 == 0 + else: + return cur_step % 10_000 == 0 + return cur_step % cfg.eval_every == 0 + def _eval_pass(cur_loss, halt_steps): nonlocal best_val_em, patience_left, done, first_row - train_tok, train_em, _ = evaluate(model, train_ex, cfg) - val_tok, val_em, val_per = evaluate(model, val_ex, cfg) + train_tok, train_em, _ = evaluate_gpu(model, train_gpu, train_ex, cfg, device) + val_tok, val_em, val_per = evaluate_gpu(model, val_gpu, val_ex, cfg, device) mean_halt = float(halt_steps.mean()) if halt_steps is not None else float("nan") row = { "step": step, "train_loss": float(cur_loss), "train_token_acc": train_tok, @@ -200,53 +233,52 @@ def main() -> None: f"val_em {val_em:.3f} val_tok {val_tok:.3f} halt {mean_halt:.2f}") while step < cfg.max_train_steps and not done: - rng.shuffle(train_ex) - for bi in range(n_batches): - sl = train_ex[bi * cfg.batch_size: (bi + 1) * cfg.batch_size] - if not sl: - continue - batch = make_batch(sl, cfg) - x = batch["x"].to(device) - y = batch["y"].to(device) - y_in = batch["y_in"].to(device) - mask = batch["y_mask"].to(device) - - with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): - out = model(x, y_in) - logits = out["logits"] - 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) * 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 - else: - frac = (step - cfg.halt_warmup_steps) / max(1, cfg.halt_ramp_end_steps - cfg.halt_warmup_steps) - lam = cfg.halt_penalty * frac + if B >= N_train: + x = train_gpu["x"] + y = train_gpu["y"] + y_in = train_gpu["y_in"] + mask = train_gpu["y_mask"] + else: + idx = torch.randint(0, N_train, (B,), device=device) + x = train_gpu["x"][idx] + y = train_gpu["y"][idx] + y_in = train_gpu["y_in"][idx] + mask = train_gpu["y_mask"][idx] + + with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): + out = model(x, y_in) + logits = out["logits"] + 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) * 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 else: - lam = 0.0 - halt = out["halt_steps"] - penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0 - loss = loss_tokens + penalty - - if cfg.lr_decay: - # cosine 1e-3 -> 1e-4 over the full run (E3, locked in Addendum 4) - frac = min(1.0, step / cfg.max_train_steps) - lr_t = cfg.lr * 0.1 + 0.5 * (cfg.lr - cfg.lr * 0.1) * (1 + math.cos(math.pi * frac)) - opt.param_groups[0]["lr"] = lr_t - - opt.zero_grad() - scaler.scale(loss).backward() - scaler.step(opt) - scaler.update() - - step += 1 - if step % cfg.eval_every == 0 or step >= cfg.max_train_steps: - _eval_pass(loss.detach(), halt.detach() if halt is not None else None) - if done: - break - if step >= cfg.max_train_steps: + frac = (step - cfg.halt_warmup_steps) / max(1, cfg.halt_ramp_end_steps - cfg.halt_warmup_steps) + lam = cfg.halt_penalty * frac + else: + lam = 0.0 + halt = out["halt_steps"] + penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0 + loss = loss_tokens + penalty + + if cfg.lr_decay: + # cosine 1e-3 -> 1e-4 over the full run (E3, locked in Addendum 4) + frac = min(1.0, step / cfg.max_train_steps) + lr_t = cfg.lr * 0.1 + 0.5 * (cfg.lr - cfg.lr * 0.1) * (1 + math.cos(math.pi * frac)) + opt.param_groups[0]["lr"] = lr_t + + opt.zero_grad() + scaler.scale(loss).backward() + scaler.step(opt) + scaler.update() + + step += 1 + if should_eval(step): + _eval_pass(loss.detach(), halt.detach() if halt is not None else None) + if done: break torch.save(model.state_dict(), os.path.join(out_dir, "last.pt")) |