Upload train/train_ounce100m.py with huggingface_hub

train/train_ounce100m.py

@@ -125,14 +125,26 @@ class HubPush(TrainerCallback):
     """
 
     def __init__(self, repo, api, token, seq_len, seed, fingerprint, seqs_per_step,
-                 prune=True, min_free_gb=6.0, repo_type="dataset"):
+                 prune=True, min_free_gb=6.0, repo_type="dataset", stop_after_steps=0):
         self.repo, self.api, self.token = repo, api, token
         self.seq_len, self.seed, self.fingerprint = seq_len, seed, fingerprint
         self.seqs_per_step, self.repo_type = seqs_per_step, repo_type
         self.prune, self.min_free_gb = prune, min_free_gb
+        self.stop_after = int(stop_after_steps or 0)
+        self.stopped_at = None
         self.n = 0
         self.results = []
 
+    def on_step_end(self, args, state, control, **kw):
+        # A session that ends ON a verified checkpoint loses nothing when the platform takes the disk.
+        # That is how the 10 % cadence is made compatible with an unknown session cap: instead of pushing
+        # extra intermediate copies (which multiplies Hub LFS storage by the number of sessions), the run
+        # is cut into segments whose length is a whole number of checkpoint intervals.
+        if self.stop_after and state.global_step >= self.stop_after:
+            self.stopped_at = state.global_step
+            control.should_training_stop = True
+        return control
+
     def on_save(self, args, state, control, **kw):
         if not state.is_world_process_zero:
             return control
@@ -257,6 +269,10 @@ def main():
     ap.add_argument("--val-tokens", type=int, default=2_000_000)
     ap.add_argument("--log-every", type=int, default=10)
     ap.add_argument("--push-every-steps", type=int, default=0, help="0 = only at 10% boundaries")
+    ap.add_argument("--stop-after-steps", type=int, default=0,
+                    help="absolute global step at which THIS session stops training, so it ends on a "
+                         "verified checkpoint instead of being killed between two of them. Must be a "
+                         "whole multiple of the checkpoint interval; 0 = run to the horizon.")
     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,
@@ -368,6 +384,14 @@ def main():
 
     # 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)
+    if args.stop_after_steps:
+        if args.stop_after_steps % ckpt_every:
+            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} is not a whole multiple of the "
+                             f"checkpoint interval {ckpt_every}: the session would end between two "
+                             f"checkpoints, which is exactly what this flag exists to prevent")
+        if args.stop_after_steps > steps_planned:
+            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} exceeds the horizon "
+                             f"{steps_planned}; use 0 to run to the end")
 
     api = None
     if args.hub_repo:
@@ -420,10 +444,13 @@ def main():
             "decay_start": int(steps_planned * args.decay_frac),
             "total_steps": steps_planned},
     )
+    hub_cb = None
     if api is not None:
-        tr.add_callback(HubPush(args.hub_repo, api, os.environ.get("HF_TOKEN"), args.seq_len,
-                                args.data_seed, fingerprint, seqs_per_step,
-                                prune=args.prune, repo_type=args.hub_repo_type))
+        hub_cb = HubPush(args.hub_repo, api, os.environ.get("HF_TOKEN"), args.seq_len,
+                         args.data_seed, fingerprint, seqs_per_step,
+                         prune=args.prune, repo_type=args.hub_repo_type,
+                         stop_after_steps=args.stop_after_steps)
+        tr.add_callback(hub_cb)
 
     # Precision must be ASSERTED, not configured: bf16 "works" silently on Turing and fp32 master
     # weights can be lost in a refactor (docs/00-platform-notes.md §8, memory/ERRORS.md E-006). The
@@ -494,7 +521,8 @@ def main():
         say(f"validation loss={ev['eval_loss']:.4f} ppl={ppl:.2f} over {n_val} held-out windows")
     except Exception as e:
         say(f"validation eval failed: {type(e).__name__}: {str(e)[:200]}")
-    if api is not None:
+    segment = bool(hub_cb and hub_cb.stopped_at and step < steps_planned)
+    if api is not None and not segment:
         r = hubckpt.push_and_prune(args.hub_repo, out, "final", api, repo_type=args.hub_repo_type,
                                    token=os.environ.get("HF_TOKEN"), prune=args.prune)
         # .get, not [k]: `pruned` and `free_after_gb` are only present on some paths through
@@ -502,7 +530,14 @@ def main():
         say("final push:", json.dumps({k: r.get(k) for k in
                                        ("verify", "pruned", "free_before_gb", "free_after_gb",
                                         "prune_blocked_reason")}, default=str)[:900])
+    elif segment:
+        # The session stopped on a checkpoint interval, and HubPush already pushed AND verified that
+        # checkpoint and rolled latest.json to it. Re-pushing it as "final" would add a 1.7 GB LFS
+        # revision per session for a copy the Hub already holds.
+        say(f"segment boundary reached at step {step} of {steps_planned}: the checkpoint HubPush just "
+            "verified is the resume point, so no separate final push")
     say("RUN_JSON " + json.dumps({"final_step": step, "val_ppl": ppl, "tokens_consumed": consumed,
+                                 "segment_stop": segment, "steps_planned": steps_planned,
                                  "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,