Upload train/train_ounce100m.py with huggingface_hub

train/train_ounce100m.py

@@ -495,25 +495,7 @@ def main():
     if final_loss is None or not math.isfinite(final_loss):
         raise SystemExit(f"final loss is not finite ({final_loss}) -- not exporting anything")
 
-    out = os.path.join(args.out, "final")
-    tr.save_model(out)
-    # The tokenizer travels with the run so `final/` is loadable by AutoTokenizer on its own (§8). If it
-    # is not in the mix root, say so in the summary instead of skipping quietly and being surprised later.
-    shipped = []
-    import shutil as _sh
-    for fn in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json",
-               "merges.txt"):
-        src = os.path.join(args.root, fn)
-        if os.path.exists(src):
-            _sh.copy(src, out)
-            shipped.append(fn)
-    say(f"tokenizer files shipped with the export: {shipped or 'NONE -- Phase 5 must add them'}")
-    SD.save_cursor(out, SD.cursor_dict(args.seq_len, samples, args.data_seed, fingerprint, step=step))
-    json.dump({"step": step, "start_step": start_step, "tokens_consumed": consumed,
-               "tokens_this_run": consumed_this_run, "seconds": round(elapsed, 1),
-               "params": pc, "fingerprint": fingerprint, "final_loss": final_loss,
-               "prior_loss": prior, "samples_consumed": samples, "steps_planned": steps_planned},
-              open(os.path.join(out, "run_summary.json"), "w"), indent=1, sort_keys=True)
+    # Validation first, and on every rank: `evaluate()` is collective, so gating it would deadlock.
     ppl = None
     try:
         ev = tr.evaluate()
@@ -521,9 +503,39 @@ 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]}")
+
+    # Everything below writes files, so it is rank 0's alone (E-029). Two ranks building the same export
+    # directory is a write race: `shutil.copy(src, out)` with `out` not yet created writes a regular FILE
+    # called `final` where rank 0 is making a directory of the same name, and the observed result was
+    # `training_args.bin` vanishing mid-hash -- which killed a fully trained, fully checkpointed leg on
+    # its last line. The variable is also no longer called `out`, which until now meant both the run
+    # directory and the export inside it.
+    exp = os.path.join(args.out, "final")
+    if rank == 0:
+        tr.save_model(exp)
+        # The tokenizer travels with the run so `final/` is loadable by AutoTokenizer on its own (§8). If
+        # it is not in the mix root, say so in the summary instead of skipping quietly and being surprised.
+        shipped = []
+        import shutil as _sh
+        os.makedirs(exp, exist_ok=True)
+        for fn in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json",
+                   "merges.txt"):
+            src = os.path.join(args.root, fn)
+            if os.path.exists(src):
+                _sh.copy(src, exp)
+                shipped.append(fn)
+        say(f"tokenizer files shipped with the export: {shipped or 'NONE -- Phase 5 must add them'}")
+        SD.save_cursor(exp, SD.cursor_dict(args.seq_len, samples, args.data_seed, fingerprint,
+                                           step=step))
+        json.dump({"step": step, "start_step": start_step, "tokens_consumed": consumed,
+                   "tokens_this_run": consumed_this_run, "seconds": round(elapsed, 1),
+                   "params": pc, "fingerprint": fingerprint, "final_loss": final_loss,
+                   "prior_loss": prior, "samples_consumed": samples,
+                   "steps_planned": steps_planned},
+                  open(os.path.join(exp, "run_summary.json"), "w"), indent=1, sort_keys=True)
     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,
+    if api is not None and not segment and rank == 0:
+        r = hubckpt.push_and_prune(args.hub_repo, exp, "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
         # push_and_prune, and a KeyError here would throw away a completed run (E-023).