diff options
| -rw-r--r-- | src/train.py | 188 |
1 files changed, 132 insertions, 56 deletions
diff --git a/src/train.py b/src/train.py index d8c03de..a0874c4 100644 --- a/src/train.py +++ b/src/train.py @@ -27,10 +27,23 @@ def resolve_device(device_str: str) -> torch.device: return torch.device(device_str) -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.""" +def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cache_size: int = 10_000) -> dict[str, torch.Tensor | int | bool]: + """Create static tensors in VRAM for small/medium datasets. For large datasets (>max_cache_size), + signals that chunked CPU streaming should be used to prevent VRAM OOM.""" + from src.data import _global_lengths + in_max, out_max = _global_lengths(cfg) if not examples: - return {"size": 0} + return { + "x": torch.empty((0, in_max), dtype=torch.long, device=device), + "y": torch.empty((0, out_max), dtype=torch.long, device=device), + "y_in": torch.empty((0, out_max), dtype=torch.long, device=device), + "y_mask": torch.empty((0, out_max), dtype=torch.bool, device=device), + "size": 0, + "is_cached": True, + } + if len(examples) > max_cache_size: + return {"size": len(examples), "is_cached": False} + batch = make_batch(examples, cfg) return { "x": batch["x"].to(device, non_blocking=True), @@ -38,12 +51,13 @@ def make_gpu_dataset(examples: list, cfg: Config, device: torch.device) -> dict[ "y_in": batch["y_in"].to(device, non_blocking=True), "y_mask": batch["y_mask"].to(device, non_blocking=True), "size": len(examples), + "is_cached": True, } @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.""" +def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples: list, cfg: Config, device: torch.device): + """Batched evaluation: operates directly on VRAM tensors if cached, or streams in eval_batch_size chunks.""" model.eval() N = gpu_ds["size"] if N == 0: @@ -55,13 +69,21 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int], examples: list, c token_total = 0 em_correct = 0 per_example = [] + is_cached = gpu_ds.get("is_cached", False) 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] + if is_cached: + 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] + else: + batch = make_batch(chunk_examples, 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) with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): out = model(x, y_in_eval) @@ -183,11 +205,17 @@ def main() -> None: best_path = os.path.join(out_dir, "best.pt") N_train = train_gpu["size"] - B = min(cfg.batch_size, N_train) if N_train > 0 else 1 + n_batches = math.ceil(N_train / cfg.batch_size) if N_train > 0 else 0 + rng = random.Random(cfg.seed) step = 0 done = False first_row = True + if N_train == 0: + print("Warning: empty training split; skipping training.") + torch.save(model.state_dict(), os.path.join(out_dir, "last.pt")) + return + def should_eval(cur_step: int) -> bool: if cur_step >= cfg.max_train_steps: return True @@ -233,53 +261,101 @@ 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: - if B >= N_train: - x = train_gpu["x"] - y = train_gpu["y"] - y_in = train_gpu["y_in"] - mask = train_gpu["y_mask"] + if train_gpu.get("is_cached", False): + perm = torch.randperm(N_train, device=device) + for bi in range(n_batches): + idx = perm[bi * cfg.batch_size: (bi + 1) * cfg.batch_size] + 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: + 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 + if step >= cfg.max_train_steps: + break 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: - 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 + 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 + 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 + if step >= cfg.max_train_steps: + break torch.save(model.state_dict(), os.path.join(out_dir, "last.pt")) print(f"DONE steps={step} best_val_em={best_val_em:.4f}") |