From 7fdda297d3ec81a7f2c8c838966de8600ba8cc3a Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Mon, 24 Aug 2026 15:31:44 +0800 Subject: [PATCH] Save periodic LoRA weight checkpoints --- README.md | 10 +++++ train_omni_turn_lora.py | 91 ++++++++++++++++++++++++++++++++--------- 2 files changed, 81 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index aea49eb..52f5705 100644 --- a/README.md +++ b/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 diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index 6b359f1..23f16c2 100644 --- a/train_omni_turn_lora.py +++ b/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: - raw_language_model.save_pretrained(adapter_dir) - torch.save( - { - "head": raw_head.state_dict(), - "head_config": { - "hidden_size": dataset.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": args.max_steps, - }, - output_dir / "turn_head.pt", + _save_weight_checkpoint( + output_dir, + raw_language_model=raw_language_model, + raw_head=raw_head, + args=args, + hidden_size=dataset.hidden_size, + training_step=args.max_steps, ) 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)