train: --grad-ckpt, and record attn/grad_ckpt in every result line

train/train_ounce100m.py

@@ -178,6 +178,9 @@ def run_benchmark(args, model, dataset, steps, world, seqs_per_step):
     if world > 1 and not dist.is_initialized():
         dist.init_process_group("nccl")
     model = model.to(device)
+    if args.grad_ckpt:
+        model.gradient_checkpointing_enable()
+        model.config.use_cache = False
     net = DDP(model, device_ids=[local_rank]) if world > 1 else model
 
     opt = torch.optim.AdamW(net.parameters(), lr=args.lr, betas=(0.9, 0.95),
@@ -214,7 +217,8 @@ def run_benchmark(args, model, dataset, steps, world, seqs_per_step):
            "seconds": round(sec, 1), "tokens": toks,
            "peak_gpu_gb": round(torch.cuda.max_memory_allocated() / 1e9, 2),
            "world": world, "seq_len": args.seq_len, "micro_batch": args.micro_batch,
-           "attn": args.attn, "params": sum(p.numel() for p in model.parameters())}
+           "attn": args.attn, "grad_ckpt": bool(args.grad_ckpt),
+           "params": sum(p.numel() for p in model.parameters())}
     print("BENCH_JSON " + json.dumps(res, sort_keys=True), flush=True)
     if world > 1:
         dist.barrier()
@@ -230,8 +234,11 @@ def main():
     ap.add_argument("--hub-repo-type", default="dataset", choices=["dataset", "model"])
     ap.add_argument("--seq-len", type=int, default=2048)
     ap.add_argument("--attn", default="eager", choices=["eager", "sdpa", "flash_attention_2"],
-                    help="p0c measured eager beating SDPA on Turing; flash_attention_2 needs sm80+ and "
-                         "is not available on a T4")
+                    help="p0c measured eager beating SDPA on Turing; Gate 3 P2/A-B measured +52 % for "
+                         "eager at seq 1024 and an OOM at 2048. flash_attention_2 needs sm80+ and is not "
+                         "available on a T4")
+    ap.add_argument("--grad-ckpt", action="store_true",
+                    help="trade ~30%% recompute for activation memory; the only route to eager at 2048")
     ap.add_argument("--layers", type=int, default=0)
     ap.add_argument("--hidden", type=int, default=0)
     ap.add_argument("--micro-batch", type=int, default=2)
@@ -386,7 +393,9 @@ def main():
         num_train_epochs=1,
         max_steps=0 if args.dry_run else steps_planned,
         fp16=True, bf16=False,
-        gradient_checkpointing=False,   # measure before needing it; docs/01-plan.md §6
+        # measured decision, not a default: A-B found eager attention 52 % faster at seq 1024 and OOM at
+        # 2048; checkpointing is what makes eager at 2048 possible at all, at the cost of a recompute pass
+        gradient_checkpointing=args.grad_ckpt,
         ddp_find_unused_parameters=False,
         dataloader_num_workers=2,
         dataloader_pin_memory=False,
@@ -496,6 +505,7 @@ def main():
                                  "micro_batch": args.micro_batch, "seq_len": args.seq_len,
                                  "accum": args.accum, "world": world,
                                  "params": pc["sum_numel"], "attn": args.attn,
+                                 "grad_ckpt": bool(args.grad_ckpt),
                                  "effective_attn": eff_attn,
                                  "peak_gpu_gb": peak_gb()}, sort_keys=True))