review round: pointer written+read-back before prune, verify retried, latest rolled to final, stale-pointer scan on resume, ddp_timeout 4h, chunked batched commits, manifest self-consistency gate, full content verification after upload

train/train_ounce100m.py

@@ -133,6 +133,7 @@ class HubPush(TrainerCallback):
         self.stop_after = int(stop_after_steps or 0)
         self.stopped_at = None
         self.n = 0
+        self.steps = []
         self.results = []
 
     def on_step_end(self, args, state, control, **kw):
@@ -156,18 +157,27 @@ class HubPush(TrainerCallback):
         SD.save_cursor(d, SD.cursor_dict(self.seq_len, samples, self.seed, self.fingerprint,
                                          step=step))
         t0 = time.time()
+        # Upload and verify *without* pruning, roll the pointer and read it back, and only then delete the
+        # local copy. The old order (prune inside push_and_prune, pointer written afterwards) meant one
+        # failed pointer commit left a Hub-verified checkpoint that nothing pointed at and no local copy to
+        # recover it with -- and every guard a resume runs would still pass, because weights, cursor,
+        # fingerprint and seed are all self-consistent. Only the pointer would be wrong, and the next
+        # session would quietly re-read 100 M tokens.
         res = hubckpt.push_and_prune(self.repo, d, f"ckpt/checkpoint-{step}", self.api,
-                                     repo_type=self.repo_type, token=self.token, prune=self.prune,
-                                     min_free_gb=self.min_free_gb)
-        hubckpt.write_latest_pointer(self.repo, {
-            "path_in_repo": f"ckpt/checkpoint-{step}", "step": step, "samples_consumed": samples,
-            "tokens_consumed": samples * self.seq_len, "fingerprint": self.fingerprint,
-            "seed": self.seed}, self.api, repo_type=self.repo_type, token=self.token)
+                                     repo_type=self.repo_type, token=self.token, prune=False)
+        ptr_body = {"path_in_repo": f"ckpt/checkpoint-{step}", "step": step,
+                    "samples_consumed": samples, "tokens_consumed": samples * self.seq_len,
+                    "fingerprint": self.fingerprint, "seed": self.seed}
+        hubckpt.write_latest_pointer(self.repo, ptr_body, self.api, repo_type=self.repo_type,
+                                     token=self.token)
+        if self.prune:
+            hubckpt.prune_local(d, res, self.min_free_gb)
         loss = next((h["loss"] for h in reversed(state.log_history) if "loss" in h), None)
         print(f"CKPT {step} samples={samples:,} push+verify={round(time.time()-t0,1)}s "
               f"ok={res['verify']['ok']} pruned={res.get('pruned')} free={res['free_after_gb']}GB "
               f"loss={loss}", flush=True)
         self.n += 1
+        self.steps.append(step)
         self.results.append(res)
         return control
 
@@ -351,6 +361,29 @@ def main():
             # exactly the §3.1 violation this whole branch exists to prevent (E-018's shape).
             raise SystemExit(f"--resume auto could not read the hub pointer: {ptr['error']}")
         else:
+            # A pointer is a resume point only if nothing newer was ever pushed. This scan catches the one
+            # case write_latest_pointer's own read-back cannot see: a pointer left behind by a rollback, a
+            # deletion, or a session that died between its checkpoint push and its pointer roll. Resuming
+            # from the stale pointer would re-train an interval with weights and a cursor that agree with
+            # each other and disagree with the Hub.
+            import re as _re
+            steps_on_hub = set()
+            for attempt in (1, 2):
+                try:
+                    steps_on_hub = {int(m.group(1)) for k in
+                                    hubckpt.hub_listing(args.hub_repo, args.hub_repo_type, token=tk)
+                                    if (m := _re.match(r"ckpt/checkpoint-(\d+)/", k))}
+                    break
+                except Exception as e:
+                    say(f"   checkpoint scan attempt {attempt}: {type(e).__name__}: {str(e)[:140]}")
+                    time.sleep(15)
+            else:
+                raise SystemExit("--resume auto could not enumerate the checkpoint repo, so it cannot "
+                                 f"prove latest.json (step {ptr.get('step')}) is the newest state")
+            if steps_on_hub and max(steps_on_hub) > int(ptr.get("step") or 0):
+                raise SystemExit(f"latest.json points at step {ptr.get('step')} but {args.hub_repo} also "
+                                 f"holds checkpoint-{max(steps_on_hub)}: the pointer is behind the run. "
+                                 "Resuming here would re-train that interval, so refusing.")
             # Both ranks need the checkpoint on disk, and this block runs before TrainingArguments has
             # created a process group -- so there is no barrier to synchronise on. Downloading into a
             # per-rank directory instead removes the race entirely (two writers, one path), at the cost of
@@ -483,6 +516,11 @@ def main():
         optim=args.optim,
         torch_compile=args.torch_compile,
         ddp_find_unused_parameters=False,
+        # Rank 0 spends minutes inside on_save uploading ~1.7 GB, listing the repo and downloading one file
+        # back to hash it, while rank 1 waits at the next collective. The library's 1,800 s default would
+        # abort the run on a slow-egress Kaggle box -- precisely during the checkpoint that the other rank
+        # was protecting. Four hours makes the collective patient without hiding a real deadlock.
+        ddp_timeout=14400,
         dataloader_num_workers=2,
         dataloader_pin_memory=False,
         remove_unused_columns=False,    # our dict keys must reach the model untouched
@@ -543,6 +581,20 @@ def main():
 
     elapsed = time.time() - t0
     step = tr.state.global_step
+    # Did this segment really end on a pushed checkpoint? `stopped_at` is set when the step counter reaches
+    # the target, but performing the save at that step is the library's decision. If it ever ordered the
+    # stop ahead of the save, the session would exit 0 having pushed nothing, and every relaunch would
+    # redo the same interval while `latest` stayed put -- which looks like a stuck run, not a broken one.
+    if hub_cb is not None and args.stop_after_steps:
+        if args.stop_after_steps <= start_step:
+            raise SystemExit(f"--stop-after-steps {args.stop_after_steps} is not beyond the resume point "
+                             f"step {start_step}: this session would train nothing")
+        if step < args.stop_after_steps:
+            raise SystemExit(f"asked to stop at step {args.stop_after_steps}, stopped at {step}")
+        if args.stop_after_steps < steps_planned and args.stop_after_steps not in hub_cb.steps:
+            raise SystemExit(f"segment ended at step {step} but no checkpoint was pushed and verified "
+                             f"for it (pushed: {sorted(hub_cb.steps)}). The next session would resume "
+                             "from the previous pointer and re-read this interval.")
     samples = samples_at(step, seqs_per_step)
     consumed = samples * args.seq_len
     consumed_this_run = max(0, step - start_step) * tokens_per_step
@@ -599,8 +651,17 @@ def main():
         # .get, not [k]: `pruned` and `free_after_gb` are only present on some paths through
         # push_and_prune, and a KeyError here would throw away a completed run (E-023).
         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])
+                                      ("verify", "pruned", "free_before_gb", "free_after_gb",
+                                       "prune_blocked_reason")}, default=str)[:900])
+        # The run is over and `latest.json` has to say so. `on_save` only ever rolls it at a 381-step
+        # boundary, so without this the pointer stops at 3,810, the launcher's `start_step >= horizon`
+        # check never fires, and every relaunch after completion trains four more steps and pushes again.
+        if step >= steps_planned:
+            hubckpt.write_latest_pointer(args.hub_repo, {
+                "path_in_repo": "final", "step": step, "samples_consumed": samples,
+                "tokens_consumed": consumed, "fingerprint": fingerprint, "seed": args.data_seed},
+                api, repo_type=args.hub_repo_type, token=os.environ.get("HF_TOKEN"))
+            say(f"latest.json rolled to the completed run: final at step {step}")
     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