Save periodic LoRA weight checkpoints

README.md

@@ -1154,6 +1154,16 @@ it be evaluated on official test `Sessions_1--3`.  A second controlled run may
 omit `--freeze-head` and use `--head-learning-rate 1e-5`; all other settings
 must remain identical.
 
+For a continuous five-epoch run, set
+`TOTAL_STEPS=$((5 * MAX_STEPS))` and pass
+`--max-steps "${TOTAL_STEPS}" --save-every-steps "${MAX_STEPS}"`.  The trainer
+keeps one AdamW/cosine-scheduler trajectory and writes evaluation-ready
+`adapter/` plus `turn_head.pt` pairs under `checkpoint-step-MAX_STEPS`,
+`checkpoint-step-2*MAX_STEPS`, through the fifth epoch.  The top-level adapter
+and head duplicate the final fifth checkpoint.  Evaluate all five intermediate
+checkpoints only on held-out `Session_12`, select one, then run that checkpoint
+once on `Sessions_1--3`.
+
 ### Frozen-Thinker sparse head-only control
 
 To determine whether the regression comes from transition-sparse sampling or

train_omni_turn_lora.py

@@ -324,6 +324,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--gradient-checkpointing", action="store_true")
     parser.add_argument("--seed", type=int, default=1337)
     parser.add_argument("--log-every", type=int, default=5)
+    parser.add_argument(
+        "--save-every-steps",
+        type=int,
+        default=0,
+        help=(
+            "Save an evaluation-ready adapter/head checkpoint every N optimizer "
+            "updates under checkpoint-step-N; 0 disables intermediate saves"
+        ),
+    )
     parser.add_argument("--eval-max-batches", type=int, default=4)
     parser.add_argument(
         "--dtype",
@@ -937,11 +946,44 @@ def _jsonable_args(args: argparse.Namespace) -> dict[str, Any]:
     return output
 
 
+def _save_weight_checkpoint(
+    checkpoint_dir: Path,
+    *,
+    raw_language_model: torch.nn.Module,
+    raw_head: TurnPredictionHead,
+    args: argparse.Namespace,
+    hidden_size: int,
+    training_step: int,
+) -> None:
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    if not args.freeze_thinker:
+        raw_language_model.save_pretrained(checkpoint_dir / "adapter")
+    torch.save(
+        {
+            "head": raw_head.state_dict(),
+            "head_config": {
+                "hidden_size": hidden_size,
+                "num_horizons": NUM_HORIZONS,
+                "num_classes": NUM_CLASSES,
+                "dropout": raw_head.dropout,
+            },
+            "initial_head_checkpoint": (
+                str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
+            ),
+            "head_frozen": bool(args.freeze_head),
+            "training_step": training_step,
+        },
+        checkpoint_dir / "turn_head.pt",
+    )
+
+
 def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
     if args.max_steps <= 0:
         raise ValueError("--max-steps must be positive")
     if args.batch_size <= 0 or args.gradient_accumulation_steps <= 0:
         raise ValueError("Batch size and gradient accumulation must be positive")
+    if args.save_every_steps < 0:
+        raise ValueError("--save-every-steps must be non-negative")
     if args.online_current_tail:
         if run.enabled:
             raise ValueError(
@@ -1023,7 +1065,10 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         output_dir / "turn_head.pt",
         output_dir / "train_summary.json",
     )
-    if not args.overwrite_output and any(path.exists() for path in known_outputs):
+    intermediate_outputs = list(output_dir.glob("checkpoint-step-*"))
+    if not args.overwrite_output and (
+        any(path.exists() for path in known_outputs) or intermediate_outputs
+    ):
         raise FileExistsError(
             f"Training output already exists in {output_dir}; pass --overwrite-output"
         )
@@ -1289,6 +1334,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         train_sampler.set_epoch(sampler_epoch)
     iterator = iter(train_loader) if batch_producer is None else None
     history: list[dict[str, float | int]] = []
+    saved_checkpoint_steps: list[int] = []
     progress = tqdm(
         total=args.max_steps,
         desc=("Frozen Thinker head" if args.freeze_thinker else "Thinker LoRA"),
@@ -1401,6 +1447,22 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
                 ),
                 flush=True,
             )
+        if args.save_every_steps and update_step % args.save_every_steps == 0:
+            torch.cuda.synchronize(device)
+            _distributed_barrier(run)
+            if run.is_main:
+                checkpoint_dir = output_dir / f"checkpoint-step-{update_step}"
+                _save_weight_checkpoint(
+                    checkpoint_dir,
+                    raw_language_model=raw_language_model,
+                    raw_head=raw_head,
+                    args=args,
+                    hidden_size=dataset.hidden_size,
+                    training_step=update_step,
+                )
+                saved_checkpoint_steps.append(update_step)
+                print(f"Saved intermediate checkpoint to {checkpoint_dir}", flush=True)
+            _distributed_barrier(run)
     progress.close()
     torch.cuda.synchronize(device)
     _distributed_barrier(run)
@@ -1441,26 +1503,14 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
     else:
         per_rank_peak_gib = [local_peak_gib]
 
-    adapter_dir = output_dir / "adapter"
     if run.is_main:
-        if not args.freeze_thinker:
+        _save_weight_checkpoint(
-            raw_language_model.save_pretrained(adapter_dir)
+            output_dir,
-        torch.save(
+            raw_language_model=raw_language_model,
-            {
+            raw_head=raw_head,
-                "head": raw_head.state_dict(),
+            args=args,
-                "head_config": {
+            hidden_size=dataset.hidden_size,
-                    "hidden_size": dataset.hidden_size,
+            training_step=args.max_steps,
-                    "num_horizons": NUM_HORIZONS,
-                    "num_classes": NUM_CLASSES,
-                    "dropout": raw_head.dropout,
-                },
-                "initial_head_checkpoint": (
-                    str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
-                ),
-                "head_frozen": bool(args.freeze_head),
-                "training_step": args.max_steps,
-            },
-            output_dir / "turn_head.pt",
         )
     global_batch_size = (
         run.world_size * args.batch_size * args.gradient_accumulation_steps
@@ -1523,6 +1573,7 @@ def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
         },
         "args": _jsonable_args(args),
         "history": history,
+        "intermediate_checkpoint_steps": saved_checkpoint_steps,
     }
     if run.is_main:
         save_json(output_dir / "train_summary.json", summary)