fix DDP hang: single checkpoint write, NCCL timeout

scugnizz-llama.py

@@ -5,7 +5,8 @@
 # dependencies = ["torch","datasets","transformers","huggingface_hub","numpy"]
 # ///
 
-import argparse, json, math, os, random, time, warnings
+import argparse, json, math, os, random, shutil, time, warnings
+from datetime import timedelta
 from contextlib import nullcontext
 from dataclasses import dataclass
 from pathlib import Path
@@ -301,7 +302,7 @@ def init_dist():
     if "RANK" not in os.environ:
         dev = "cuda" if torch.cuda.is_available() else "cpu"
         return 0, 1, 0, dev
-    dist.init_process_group(backend="nccl")
+    dist.init_process_group(backend="nccl", timeout=timedelta(hours=2))
     rank = int(os.environ["RANK"])
     local_rank = int(os.environ["LOCAL_RANK"])
     world_size = int(os.environ["WORLD_SIZE"])
@@ -433,10 +434,19 @@ def save_ckpt(path, model, opt, step, best, rank, world_size, mirror_resume=Fals
     print("SAVED", path, f"step={step}", flush=True)
     if mirror_resume:
         resume = path.parent / RESUME_CKPT
-        tmp_resume = resume.with_suffix(resume.suffix + ".tmp")
-        torch.save(payload, tmp_resume)
-        tmp_resume.replace(resume)
-        print("SAVED", resume, f"step={step} (keep for optimizer resume)", flush=True)
+        shutil.copy2(path, resume)
+        print("SAVED", resume, f"step={step} (copy for optimizer resume)", flush=True)
+    barrier(rank, world_size)
+
+
+def copy_best_ckpt(out, rank, world_size):
+    if not is_main(rank):
+        barrier(rank, world_size)
+        return
+    last, best = out / "checkpoint_last.pt", out / "checkpoint_best.pt"
+    if last.exists():
+        shutil.copy2(last, best)
+        print("SAVED", best, "(copy of checkpoint_last)", flush=True)
     barrier(rank, world_size)
 
 
@@ -690,7 +700,7 @@ def main():
 
     log0(rank, "initializing data streams...", flush=True)
     train = Batcher(a, tok, dev, rank, log=is_main(rank))
-    val = Batcher(a, tok, dev, 10_000 + rank, log=False)
+    val = Batcher(a, tok, dev, 10_000, log=False) if is_main(rank) else None
     log0(rank, f"starting training loop from step {start}...", flush=True)
     t0 = time.time()
     roll = 0.0
@@ -730,17 +740,28 @@ def main():
         if step > 0 and a.weights_save_interval > 0 and step % a.weights_save_interval == 0:
             save_weights(weights_last, model, step, rank, world_size)
 
+        val_improved = False
         if step > 0 and step % a.eval_interval == 0:
             if is_main(rank):
                 vl = eval_model(raw_model(model), val, a, dev, dt)
                 log0(rank, f"EVAL {step:08d} | val_loss {vl:.4f} | val_ppl {math.exp(min(20, vl)):.1f}", flush=True)
                 if vl < best:
                     best = vl
-                    save_ckpt(bestp, model, opt, step, best, rank, world_size)
+                    val_improved = True
+            best_pkg = [best]
+            if world_size > 1:
+                dist.broadcast_object_list(best_pkg, src=0)
+            best = float(best_pkg[0])
+            if world_size > 1:
+                imp_pkg = [val_improved]
+                dist.broadcast_object_list(imp_pkg, src=0)
+                val_improved = imp_pkg[0]
             barrier(rank, world_size)
 
         if step > 0 and step % a.save_interval == 0:
             save_ckpt(last, model, opt, step, best, rank, world_size, mirror_resume=True)
+            if val_improved:
+                copy_best_ckpt(out, rank, world_size)
             if is_main(rank) and a.push_every_save:
                 upload_final(str(out), a.hub_repo_id, a.hub_path, f"pretrain checkpoint step {step}")