Upload train/train_ounce100m.py with huggingface_hub

train/train_ounce100m.py

@@ -259,6 +259,15 @@ def main():
                     help="4 x 1024 x 32 accum x 2 cards = 262,144 tokens/step, the frozen batch size")
     ap.add_argument("--accum", type=int, default=32)
     ap.add_argument("--lr", type=float, default=6e-4)
+    ap.add_argument("--optim", default="adamw_torch",
+                    help="only mathematically-equivalent implementations are fair game at freeze time "
+                         "(adamw_torch / adamw_torch_fused / adamw_apex_fused); 8-bit variants change the "
+                         "update, so they are excluded by rule, not by preference")
+    ap.add_argument("--gc-nonreentrant", default=False, action=argparse.BooleanOptionalAction,
+                    help="pass use_reentrant=False to gradient checkpointing: same recomputation, "
+                         "different autograd path, and reports on whether it is faster under DDP conflict")
+    ap.add_argument("--torch-compile", default=False, action=argparse.BooleanOptionalAction,
+                    help="Trainer's torch_compile hook; fused kernels, no change to the update rule")
     ap.add_argument("--tokens", type=int, default=1_000_000_000)
     ap.add_argument("--max-steps", type=int, default=0, help="0 = derive from --tokens; >0 overrides")
     ap.add_argument("--warmup-frac", type=float, default=0.02)
@@ -365,6 +374,20 @@ def main():
     vstore = SD.PackedTokenStore(args.root, man, files=vs)
     n_val = min(max(1, args.val_tokens // args.seq_len), vstore.total_tokens // args.seq_len)
     val_ds = SD.make_dataset(vstore, args.seq_len, args.data_seed, shuffle=False, count=n_val)
+    need = steps_planned * seqs_per_step
+    if len(train_ds) < need:
+        # Reading past the end of the stream makes the Trainer begin epoch 2 at the top of the same
+        # permutation: the run would then report the full token target while quietly re-training on the
+        # first windows. The mix is built with ~12 % headroom over the target for exactly this reason, so
+        # arriving here means that assumption stopped being true (review finding), and the honest response
+        # is to stop before step 1 rather than to discover it in the final report.
+        raise SystemExit(
+            f"dataset too small for the plan: {len(train_ds):,} windows available from sample "
+            f"{start_sample:,}, {need:,} needed for {steps_planned} steps x {seqs_per_step} sequences "
+            f"({need * args.seq_len:,} tokens). Lower --tokens to the 0.9 B fallback (D-012) or rebuild "
+            "the mix larger -- never let the epoch wrap.")
+    say(f"data: {len(train_ds):,} training windows available, {need:,} needed "
+        f"({100.0 * len(train_ds) / need - 100:.1f} % margin), {n_val:,} held-out val windows")
 
     cfg = build_config(args.seq_len, args.layers or None, args.hidden or None, attn=args.attn)
     set_seed(args.seed)  # before init: both ranks must start from identical weights, not two draws
@@ -410,6 +433,10 @@ def main():
     tacc = TrainingArguments(
         output_dir=args.out,
         per_device_train_batch_size=args.micro_batch,
+        # Same geometry for the validation pass: with the library default (8) the end-of-run eval measures
+        # a peak that has nothing to do with the training configuration, which is how an OOM in the no-
+        # checkpointing cell could be mistaken for a passing one (review finding).
+        per_device_eval_batch_size=args.micro_batch,
         gradient_accumulation_steps=args.accum,
         learning_rate=args.lr,
         weight_decay=0.1,
@@ -423,6 +450,12 @@ def main():
         # 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,
+        # Same recomputation, different autograd path -- offered as a cells-vs-baseline question, not a
+        # belief. Reports on use_reentrant=False under DDP contradict each other across versions.
+        gradient_checkpointing_kwargs={"use_reentrant": not args.gc_nonreentrant}
+        if args.grad_ckpt else None,
+        optim=args.optim,
+        torch_compile=args.torch_compile,
         ddp_find_unused_parameters=False,
         dataloader_num_workers=2,
         dataloader_pin_memory=False,
@@ -556,6 +589,8 @@ def main():
                                  "accum": args.accum, "world": world,
                                  "params": pc["sum_numel"], "attn": args.attn,
                                  "grad_ckpt": bool(args.grad_ckpt),
+                                 "optim": args.optim, "gc_nonreentrant": bool(args.gc_nonreentrant),
+                                 "torch_compile": bool(args.torch_compile),
                                  "effective_attn": eff_attn,
                                  "peak_gpu_gb": peak_gb()}, sort_keys=True))