train: frozen-width keeps GQA 3:1; report peak_gpu_gb and per-run tok/s

train/train_ounce100m.py

@@ -47,8 +47,12 @@ def build_config(seq_len, layers=None, hidden=None):
     parameter count and makes T1 compare two different models."""
     h = hidden or BASE["hidden"]
     L = layers or BASE["layers"]
-    inter = BASE["ffn"] if hidden is None and layers is None else (int(round(8 * h / 3)) // 64) * 64
-    kv = BASE["kv_heads"] if hidden is None else _kv_dividing(BASE["heads"], h)
+    # "variant" means a different WIDTH, not merely a different layer count: the frozen GQA ratio (9 Q /
+    # 3 KV) belongs to the frozen width, and passing --hidden 576 explicitly must not quietly turn the
+    # model into MHA-9 and make test T1 measure something else.
+    frozen_width = (h == BASE["hidden"])
+    inter = BASE["ffn"] if frozen_width else (int(round(8 * h / 3)) // 64) * 64
+    kv = BASE["kv_heads"] if frozen_width else _kv_dividing(BASE["heads"], h)
     cfg = LlamaConfig(vocab_size=BASE["vocab"], hidden_size=h, intermediate_size=inter,
                       num_hidden_layers=L, num_attention_heads=BASE["heads"],
                       num_key_value_heads=kv,
@@ -331,6 +335,7 @@ def main():
     step = tr.state.global_step
     samples = samples_at(step, start_sample, seqs_per_step)
     consumed = samples * args.seq_len
+    consumed_this_run = step * tokens_per_step
     hist = [h for h in tr.state.log_history if "loss" in h]
     final_loss = hist[-1]["loss"] if hist else None
     say(f"TRAIN DONE step={step}/{steps_planned} elapsed={elapsed/60:.1f} min "
@@ -369,8 +374,22 @@ def main():
                                    token=os.environ.get("HF_TOKEN"), prune=args.prune)
         say("final push:", json.dumps({k: r[k] for k in ("verify", "pruned", "free_after_gb")})[:900])
     say("RUN_JSON " + json.dumps({"final_step": step, "val_ppl": ppl, "tokens_consumed": consumed,
-                                 "tok_per_s": round(consumed / max(1, elapsed)),
-                                 "final_loss": final_loss, "prior_loss": prior}, sort_keys=True))
+                                 "tok_per_s": round(consumed_this_run / max(1, elapsed)),
+                                 "final_loss": final_loss, "prior_loss": prior,
+                                 "micro_batch": args.micro_batch, "seq_len": args.seq_len,
+                                 "accum": args.accum, "world": world,
+                                 "params": pc["sum_numel"],
+                                 "peak_gpu_gb": peak_gb()}, sort_keys=True))
+
+
+def peak_gb():
+    """Peak CUDA allocation, in GB. The batch size at seq 2048 is the one thing p0c could not tell us
+    (it measured seq 1024), and an OOM mid-main-run is a lost session, so this number is part of every
+    preflight cell's evidence rather than something to discover at step 300 of the real run."""
+    try:
+        return round(torch.cuda.max_memory_allocated() / 1e9, 2) if torch.cuda.is_available() else None
+    except Exception:
+        return None
 
 
 if __name__ == "__main__":