perf(cuda): add AMP fp16 autocast with GradScaler and torch.compile reduce-overhead mode

2 files changed, 46 insertions, 25 deletions

diff --git a/src/config.py b/src/config.py
index 25dcdda..4db66b0 100644
--- a/src/config.py
+++ b/src/config.py
@@ -39,6 +39,8 @@ class Config:
     log_n_examples: int = 10
     out_dir: str = "runs"
     device: str = "auto"             # "auto" | "cuda" | "cpu" | "mps"
+    use_amp: bool = True             # FP16 mixed precision on CUDA (Turing Tensor Cores)
+    compile_model: bool = False      # torch.compile (reduce-overhead / CUDA Graphs)
 
     @property
     def vocab(self) -> int:
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)