diff options
Diffstat (limited to 'src/train.py')
| -rw-r--r-- | src/train.py | 69 |
1 files changed, 44 insertions, 25 deletions
diff --git a/src/train.py b/src/train.py index e8870df..ec7bb4e 100644 --- a/src/train.py +++ b/src/train.py @@ -44,6 +44,7 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None): return 0.0, 0.0, [] eval_bs = getattr(cfg, "eval_batch_size", 512) + use_amp = getattr(cfg, "use_amp", True) and device.type == "cuda" token_correct = 0 token_total = 0 em_correct = 0 @@ -57,8 +58,9 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None): 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) + 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) pred_tok = logits.argmax(-1) token_correct += int((pred_tok[mask] == y[mask]).sum().item()) token_total += int(mask.sum().item()) @@ -71,10 +73,11 @@ def evaluate(model, examples, cfg: Config, device: torch.device | None = None): 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) + with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp): + 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): @@ -103,6 +106,7 @@ def main() -> None: cfg = parse_args() set_seed(cfg.seed) device = resolve_device(cfg.device) + use_amp = cfg.use_amp and device.type == "cuda" 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): @@ -115,6 +119,8 @@ def main() -> None: "torch": torch.__version__, "numpy": np.__version__, "device": str(device), + "use_amp": use_amp, + "compile_model": cfg.compile_model, "torch_threads": torch.get_num_threads(), "cmd": sys.argv, } @@ -128,7 +134,13 @@ def main() -> None: train_ex = build_examples(train_in, cfg) val_ex = build_examples(val_in, cfg) model = build_model(cfg).to(device) - print(f"model={cfg.model} params={model.param_count()} device={device} " + if cfg.compile_model and hasattr(torch, "compile") and device.type == "cuda": + try: + model = torch.compile(model, mode="reduce-overhead") + except Exception as e: + print(f"Warning: torch.compile failed ({e}), falling back to uncompiled model") + + print(f"model={cfg.model} params={model.param_count()} device={device} amp={use_amp} " f"train={len(train_ex)} val={len(val_ex)} " f"vocab={cfg.vocab} eos={cfg.eos_id} pad={cfg.pad_id}") @@ -140,6 +152,7 @@ def main() -> None: opt = torch.optim.AdamW(model.parameters(), **opt_kwargs) else: opt = torch.optim.AdamW(model.parameters(), **opt_kwargs) + scaler = torch.amp.GradScaler("cuda", enabled=use_amp) ce = nn.CrossEntropyLoss(reduction="none") log_examples = sorted(val_ex, key=lambda ex: (len(ex[0]), ex[0]))[: cfg.log_n_examples] @@ -197,31 +210,37 @@ def main() -> None: 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 = 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 + + 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: - 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 + 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() - loss.backward() - opt.step() + 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) |