lars
87e37ebe14
CI / Format (ruff format) (push) Successful in 29s
CI / Lint (ruff check) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 38s
CI / Lint (ruff check) (pull_request) Successful in 31s
CI / Format (ruff format) (pull_request) Successful in 44s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 47s
CI / Tests (push) Successful in 4m32s
CI / Tests (pull_request) Successful in 4m25s
stage{1,2}_model.active = false already trains one stage alone, but the
checkpoint it writes holds only that stage, so giant rollout refuses it --
the "retrain stage 2 alone against a fixed, known-good stage 1" experiment
the 2026-08-03 species failure calls for wasn't runnable end to end.
Adds stage{1,2}_model.init_from (a checkpoint .pt to load this stage's
weights from before training) and .freeze (never update them), symmetric
across both stages. Both stages stay active = true, so both get built and
both land in the output checkpoint -- the frozen stage is merely
initialized from disk instead of from scratch.
Decisions made during planning:
- Soft freeze: forward/backward still run every batch (loss/grad_norm stay
meaningful, no autograd special-casing), only optimizer.step() (and, for
the frozen stage, lr_sched.step()/EMA update) is skipped -- weights are
byte-identical for the whole run. This is StageTrainer._step_optimizer,
shared by the non-adversarial path and both halves (generator + critic)
of the WGAN path, so a frozen WGAN stage's critic freezes too.
- validate_config requires init_from whenever freeze = true, unless the run
is a --resume (a resumed frozen stage's weights come from the resume
checkpoint instead) -- freezing a randomly-initialized model is almost
certainly a mistake.
- CLI flags on both `giant train` and `giant new-run`
(--stage{1,2}-init-from/--stage{1,2}-freeze), matching every other
per-stage model knob's existing treatment.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
102 lines
4.1 KiB
Python
102 lines
4.1 KiB
Python
"""Checkpoint assembly and restore.
|
|
|
|
The on-disk layout is unchanged from v0.2/v0.3.0 and is read by
|
|
`giant/cli.py`, `giant/rollout.py`, `giant/sample.py` and
|
|
`giant/analysis/router_gating.py` — stage 1's weights live under `model`,
|
|
stage 2's under `sec_decoder`, with `_ema`/`critic`/`sec_critic` companions
|
|
and per-stage `optimizer_<stage>` / `optimizer_d_<stage>` / `lr_sched_<stage>`
|
|
entries.
|
|
"""
|
|
|
|
import torch
|
|
|
|
from giant.training.trainers import StageTrainer
|
|
|
|
#: Stage name -> the checkpoint key its weights live under. Historical: stage
|
|
#: 1 predates the two-stage split, so it kept the bare "model" key.
|
|
_STAGE_KEY = {"stage1": "model", "stage2": "sec_decoder"}
|
|
_CRITIC_KEY = {"stage1": "critic", "stage2": "sec_critic"}
|
|
|
|
|
|
def build_checkpoint(
|
|
trainers: dict[str, StageTrainer],
|
|
epoch: int,
|
|
global_step: int,
|
|
best_val_loss: float,
|
|
extras: dict,
|
|
) -> dict:
|
|
"""`extras` carries the dataset-level sidecars (normalizer, vocab maps,
|
|
model_config) that `train()` receives as arguments; `None` values are
|
|
omitted so an absent sidecar leaves no key behind."""
|
|
ckpt: dict = {
|
|
"epoch": epoch,
|
|
"best_val_loss": best_val_loss,
|
|
"global_step": global_step,
|
|
}
|
|
for name, trainer in trainers.items():
|
|
sd = trainer.state_dict()
|
|
key = _STAGE_KEY[name]
|
|
ckpt[key] = sd["model"]
|
|
if "model_ema" in sd:
|
|
ckpt[f"{key}_ema"] = sd["model_ema"]
|
|
if "critic" in sd:
|
|
ckpt[_CRITIC_KEY[name]] = sd["critic"]
|
|
ckpt[f"optimizer_d_{name}"] = sd["optimizer_d"]
|
|
ckpt[f"optimizer_{name}"] = sd["optimizer"]
|
|
ckpt[f"lr_sched_{name}"] = sd["lr_sched"]
|
|
ckpt.update({k: v for k, v in extras.items() if v is not None})
|
|
return ckpt
|
|
|
|
|
|
def init_stages_from_checkpoints(trainers: dict[str, StageTrainer]) -> list[str]:
|
|
"""Load each trainer's `spec.init_from` checkpoint (gitea #42) into its
|
|
model, before training starts — the partial-retrain counterpart to
|
|
`load_checkpoint`'s full-run `--resume`. Only weights move: unlike
|
|
`load_checkpoint`, this never touches optimizer/lr_sched/epoch state, so
|
|
it composes cleanly with `--resume` (call this first; a resume's own
|
|
`load_checkpoint` then overwrites whatever this loaded with the resumed
|
|
run's own weights).
|
|
|
|
A stage with no `init_from` set (`""`, the default) is left alone. The
|
|
EMA companion (`<key>_ema`) is loaded too when both the source checkpoint
|
|
and this trainer have one, so `--weights ema` at inference still sees the
|
|
source's EMA shadow rather than a copy of its raw weights. Returns one
|
|
description string per stage actually initialized, for the caller to
|
|
echo.
|
|
"""
|
|
loaded = []
|
|
for name, trainer in trainers.items():
|
|
init_from = trainer.spec.init_from
|
|
if not init_from:
|
|
continue
|
|
key = _STAGE_KEY[name]
|
|
ckpt = torch.load(init_from, map_location="cpu", weights_only=False)
|
|
trainer.model.load_state_dict(ckpt[key])
|
|
ema_key = f"{key}_ema"
|
|
if trainer.ema_model is not None and ema_key in ckpt:
|
|
trainer.ema_model.load_state_dict(ckpt[ema_key])
|
|
loaded.append(f"{name}: loaded from {init_from}" + (" (frozen)" if trainer.frozen else ""))
|
|
return loaded
|
|
|
|
|
|
def load_checkpoint(trainers: dict[str, StageTrainer], ckpt: dict, lr: float) -> None:
|
|
"""Restore every active stage, then hand `lr`'s authority back to the
|
|
config — `load_state_dict` would otherwise leave the checkpoint's own
|
|
base LR in place, silently ignoring `--lr` on resume."""
|
|
for name, trainer in trainers.items():
|
|
key = _STAGE_KEY[name]
|
|
sd = {
|
|
"model": ckpt[key],
|
|
"optimizer": ckpt[f"optimizer_{name}"],
|
|
"lr_sched": ckpt[f"lr_sched_{name}"],
|
|
}
|
|
ema_key = f"{key}_ema"
|
|
if ema_key in ckpt:
|
|
sd["model_ema"] = ckpt[ema_key]
|
|
crit_key = _CRITIC_KEY[name]
|
|
if crit_key in ckpt:
|
|
sd["critic"] = ckpt[crit_key]
|
|
sd["optimizer_d"] = ckpt[f"optimizer_d_{name}"]
|
|
trainer.load_state_dict(sd)
|
|
trainer.resume_lr(lr)
|