train: --benchmark N runs the same step with no Trainer, to attribute the P2 throughput gap

train/train_ounce100m.py

@@ -160,6 +160,68 @@ class HubPush(TrainerCallback):
         return control
 
 
+def run_benchmark(args, model, dataset, steps, world, seqs_per_step):
+    """The same data path, optimizer and fp16 recipe as the real run, with no Trainer in the loop.
+
+    Why this exists: P2 measured 5,105 tok/s for 22L/seq-1024 where probe p0c measured 11,062 tok/s for a
+    near-identical shape with a hand-rolled loop. The gap is either (a) the attention kernel, (b) per-step
+    framework overhead, or (c) a measurement artefact, and which one it is decides whether the main run
+    needs 27 h or 68 h. Guessing wrong in either direction is expensive, so measure both paths.
+    """
+    import torch.distributed as dist
+    from torch.nn.parallel import DistributedDataParallel as DDP
+    from torch.utils.data import DataLoader
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    device = torch.device("cuda", local_rank)
+    torch.cuda.set_device(device)
+    if world > 1 and not dist.is_initialized():
+        dist.init_process_group("nccl")
+    model = model.to(device)
+    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),
+                            weight_decay=0.1, eps=1e-8)
+    try:
+        scaler = torch.amp.GradScaler("cuda", enabled=True)     # torch >= 2.4
+    except TypeError:
+        scaler = torch.cuda.amp.GradScaler(enabled=True)        # older signature
+    # Same reader, same shuffle, no accumulation: the comparison against the Trainer cells is then
+    # per-step cost with the framework removed, and nothing else differs.
+    dl = DataLoader(dataset, batch_size=args.micro_batch, shuffle=False, num_workers=2,
+                    pin_memory=False, drop_last=True)
+    t0 = time.time()
+    toks = 0
+    for i, batch in enumerate(dl):
+        if i >= steps:
+            break
+        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+        opt.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.float16):
+            out = net(**batch)
+        scaler.scale(out.loss).backward()
+        scaler.unscale_(opt)
+        torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
+        scaler.step(opt)
+        scaler.update()
+        toks += batch["input_ids"].numel()
+        if i == 0:
+            torch.cuda.synchronize()
+            t0 = time.time()   # exclude CUDA context + first-trace cost from the rate
+            toks = 0
+    sec = time.time() - t0
+    res = {"benchmark_steps": i + 1 if steps else 0, "tok_per_s": round(toks / max(1e-6, sec)),
+           "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())}
+    print("BENCH_JSON " + json.dumps(res, sort_keys=True), flush=True)
+    if world > 1:
+        dist.barrier()
+        dist.destroy_process_group()
+    return res
+
+
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--root", default="/kaggle/working/mixroot")
@@ -187,6 +249,8 @@ def main():
     ap.add_argument("--push-every-steps", type=int, default=0, help="0 = only at 10% boundaries")
     ap.add_argument("--prune", action="store_true", help="delete local checkpoint after hub verify")
     ap.add_argument("--dry-run", action="store_true", help="build everything, train 0 steps")
+    ap.add_argument("--benchmark", type=int, default=0,
+                    help="N steps with no Trainer in the loop (isolates framework per-step cost)")
     args = ap.parse_args()
 
     rank = int(os.environ.get("RANK", "0"))
@@ -284,6 +348,14 @@ def main():
     if not (90_000_000 <= pc["sum_numel"] <= 110_000_000) and not args.layers and not args.hidden:
         raise SystemExit(f"parameter count {pc['sum_numel']:,} is outside §2's 90-110M budget")
 
+    if args.benchmark:
+        # Both ranks iterate the same samples here (no DistributedSampler): this path measures step cost,
+        # not a training run, and it must never be used to produce a checkpoint.
+        say(f"benchmark mode: {args.benchmark} steps, no Trainer, micro={args.micro_batch} "
+            f"seq={args.seq_len} attn={args.attn} world={world}")
+        run_benchmark(args, model, train_ds, args.benchmark, world, seqs_per_step)
+        return
+
     # 10 checkpoints per run is §2's cadence; --push-every-steps overrides for the preflight.
     ckpt_every = args.push_every_steps or max(1, steps_planned // 10)