wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Save periodic LoRA weight checkpoints

Browse Source
This commit is contained in:
Xin Wang
2026-08-24 15:31:44 +08:00
parent 7acab44ea2
commit 7fdda297d3
2 changed files with 81 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
README.md
View File

@@ -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

91
train_omni_turn_lora.py
View File

@@ -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)