preflight: stage pargs -- construct TrainingArguments/model/schedule on the Kaggle image, cross-check 22L=106,194,240 and 20L=99,114,048

train/preflight.py

@@ -49,6 +49,97 @@ def last_json(text, begin, end):
         return None
 
 
+# ------------------------------------------------------------------------ Pargs (CPU, free)
+def p_args(args, R):
+    """Every keyword in train_ounce100m.py's TrainingArguments is a promise about a library version, and
+    E-008/E-010 already taught us that transformers 5 drops things without warning. This constructs the
+    real objects on the Kaggle image -- TrainingArguments, the config, the model, the optimizer and the
+    trapezoid schedule -- with no data and no GPU, so a kwargs skew costs a free CPU minute instead of
+    killing a billed 2xT4 session at import time."""
+    src = r'''
+import json, os, sys, math, time
+sys.path.insert(0, "/kaggle/working")
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+import torch
+import transformers
+from torch.optim import AdamW
+from transformers import TrainingArguments
+import train_ounce100m as T
+R = {"transformers": transformers.__version__, "torch": torch.__version__}
+cfg = T.build_config(2048)
+R["config"] = {k: getattr(cfg, k) for k in ("hidden_size", "num_hidden_layers", "num_attention_heads",
+                                            "num_key_value_heads", "intermediate_size", "vocab_size",
+                                            "tie_word_embeddings", "max_position_embeddings")}
+t0 = time.time()
+m = T.LlamaForCausalLM(cfg)
+R["params"] = T.count_params(m)
+R["init_seconds"] = round(time.time() - t0, 1)
+R["shapes_20L"] = T.count_params(T.LlamaForCausalLM(T.build_config(1024, layers=20)))
+c20 = T.build_config(1024, layers=20)
+R["shape20_ffn"] = c20.intermediate_size
+# The exact kwargs the trainer passes. If the image rejects one, this is where we learn it.
+kw = dict(output_dir="/kaggle/working/argcheck", per_device_train_batch_size=2,
+          gradient_accumulation_steps=32, learning_rate=6e-4, weight_decay=0.1, adam_beta1=0.9,
+          adam_beta2=0.95, adam_epsilon=1e-8, max_grad_norm=1.0, lr_scheduler_type="constant",
+          warmup_ratio=0.0, num_train_epochs=1, max_steps=10, fp16=True, bf16=False,
+          gradient_checkpointing=False, ddp_find_unused_parameters=False, dataloader_num_workers=2,
+          dataloader_pin_memory=False, remove_unused_columns=False, ignore_data_skip=True,
+          save_strategy="steps", save_steps=5, save_total_limit=1, logging_steps=5, report_to=[],
+          seed=1, data_seed=1, accelerator_config={"dispatch_batches": False},
+          average_tokens_across_devices=False)
+try:
+    ta = TrainingArguments(**kw); R["kwargs"] = "accepted"
+except TypeError as e:
+    R["kwargs"] = f"REJECTED: {str(e)[:300]}"
+    bad = [k for k in kw if k in str(e)]
+    R["suspect_keys"] = bad
+    for k in bad:
+        kw.pop(k, None)
+    ta = TrainingArguments(**kw)
+R["fp16_bf16"] = [ta.fp16, ta.bf16]
+# schedule shape, standalone: warm 2%, flat to 80%, linear to zero
+steps = 3815
+opt = AdamW(m.parameters(), lr=6e-4, betas=(0.9, 0.95), weight_decay=0.1)
+tr = T.TrapezoidTrainer.__new__(T.TrapezoidTrainer)   # no Trainer.__init__: we want the schedule only
+tr.lr_shape = {"warmup_steps": max(100, int(steps * 0.02)), "decay_start": int(steps * 0.8),
+               "total_steps": steps}
+tr.optimizer, tr.lr_scheduler = opt, None
+sch = tr.create_scheduler(steps, opt)
+mults = []
+for i in range(steps):
+    mults.append(sch.lr_lambdas[0](i))
+    sch.optimizer.param_groups[0]["lr"] = 6e-4 * mults[-1]
+R["lr_shape"] = {"peak": max(mults), "at_warm_end": mults[76], "at_80pct": mults[int(steps*0.8)],
+                 "at_90pct": mults[int(steps*0.9)], "final": mults[-1],
+                 "plateau_fraction_steps": round(sum(1 for x in mults if x > 0.999) / steps, 3)}
+# one real forward/backward on CPU with a tiny model to prove the collator output feeds the model
+cfg2 = T.build_config(64, layers=2, hidden=128)
+m2 = T.LlamaForCausalLM(cfg2)
+store_ids = torch.randint(0, 1000, (3, 65))
+out = m2(input_ids=store_ids[:, :-1], labels=store_ids[:, 1:],
+         attention_mask=torch.ones_like(store_ids[:, :-1]))
+out.loss.backward()
+R["forward_backward"] = {"loss": round(float(out.loss), 4), "finite": math.isfinite(float(out.loss)),
+                         "grads": sum(1 for p in m2.parameters() if p.grad is not None)}
+print("ARGS_JSON_BEGIN")
+print(json.dumps(R, default=str))
+print("ARGS_JSON_END")
+'''
+    r = sh([sys.executable, "-c", src], timeout=3600, label="Pargs construct on the Kaggle image")
+    R["PARGS"] = last_json(r["out"], "ARGS_JSON_BEGIN", "ARGS_JSON_END")
+    R["PARGS_rc"] = r["rc"]
+    j = R["PARGS"] or {}
+    R["PARGS_pass"] = bool(r["rc"] == 0 and j.get("kwargs") == "accepted"
+                           and j.get("params", {}).get("sum_numel") == 106194240
+                           # 20L must be 99,114,048 per docs/01-plan.md §2.3's closed-form table, and its
+                           # FFN must stay 1536; a doubled FFN would make test T1 meaningless.
+                           and j.get("shapes_20L", {}).get("sum_numel") == 99114048
+                           and j.get("shape20_ffn") == 1536
+                           and (j.get("lr_shape") or {}).get("final") == 0.0
+                           and (j.get("forward_backward") or {}).get("finite"))
+    print("VERDICT PARGS_pass=", R["PARGS_pass"], json.dumps(j.get("kwargs"))[:200], flush=True)
+
+
 # ------------------------------------------------------------------------ P0 reader (CPU, free)
 def p0_reader(args, R):
     """The reader is where an unrecoverable main run would hide: if two runs at the same cursor read
@@ -219,7 +310,7 @@ def p3_cold_resume(args, R):
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--stage", required=True,
-                    choices=["p0", "p1", "p2", "p3", "summary", "all"])
+                    choices=["pargs", "p0", "p1", "p2", "p3", "summary", "all"])
     ap.add_argument("--seq-len", type=int, default=2048)
     ap.add_argument("--accum", type=int, default=32)
     ap.add_argument("--tokens", type=int, default=1_000_000_000)
@@ -235,8 +326,10 @@ def main():
             R = json.load(open(args.out_json))
         except Exception:
             print("existing preflight.json unreadable; starting fresh", flush=True)
-    todo = ["p0", "p1", "p2", "p3"] if args.stage == "all" else [args.stage]
+    todo = ["pargs", "p0", "p1", "p2", "p3"] if args.stage == "all" else [args.stage]
     t0 = time.time()
+    if "pargs" in todo:
+        p_args(args, R)
     if "p0" in todo:
         p0_reader(args, R)
     if "p1" in todo:
@@ -247,7 +340,7 @@ def main():
         p3_cold_resume(args, R)
     R["seconds_this_invocation"] = round(time.time() - t0, 1)
     R["gpu_hours_this_invocation"] = round(R["seconds_this_invocation"] / 3600.0, 3)
-    R["PASSES"] = {k: R.get(k + "_pass") for k in ("P0", "P1", "P2", "P3")}
+    R["PASSES"] = {k: R.get(k + "_pass") for k in ("PARGS", "P0", "P1", "P2", "P3")}
     R["GATE_3_READY"] = all(v is True for v in R["PASSES"].values())
     with open(args.out_json, "w") as f:
         json.dump(R, f, indent=1, default=str)