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>
32 lines
902 B
Python
32 lines
902 B
Python
"""Training: per-stage trainers, metric collection, checkpointing, the loop.
|
|
|
|
Split out of the former single-module `giant/train.py`. The public surface is
|
|
`train` (the entry point `giant.pipeline` calls) plus the trainer/spec types
|
|
that tests and tooling construct directly.
|
|
"""
|
|
|
|
from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint
|
|
from giant.training.metrics import MetricsCollector, MetricSpec
|
|
from giant.training.loop import train
|
|
from giant.training.trainers import (
|
|
FlowDDPMStageTrainer,
|
|
StageSpec,
|
|
StageTrainer,
|
|
WGANStageTrainer,
|
|
build_stage_trainers,
|
|
)
|
|
|
|
__all__ = [
|
|
"FlowDDPMStageTrainer",
|
|
"MetricSpec",
|
|
"MetricsCollector",
|
|
"StageSpec",
|
|
"StageTrainer",
|
|
"WGANStageTrainer",
|
|
"build_checkpoint",
|
|
"build_stage_trainers",
|
|
"init_stages_from_checkpoints",
|
|
"load_checkpoint",
|
|
"train",
|
|
]
|