2 Commits

Author	SHA1	Message	Date
lars	87e37ebe14	Add per-stage init_from/freeze (gitea #42) ... 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>	2026-08-17 14:23:20 +02:00
lars	8019a80563	Refactor train.py into giant/training/ around a metrics collector ... Every metric name used to exist in four places: the dict keys each StageTrainer returned, the hardcoded _metrics_fields() column list, the ~110-line metrics_row assembly in train(), and the tqdm/summary formatting. The two had to be kept in exact correspondence by hand or csv.DictWriter would raise. Each metric is now declared once, as a MetricSpec on the trainer that computes it. MetricsCollector derives the CSV header and W&B payload from those declarations and owns all accumulation, so train() no longer carries a running sum, and every isinstance(tr, WGANStageTrainer) branch is gone — replaced by four trainer hooks (batch_loss, summary, val_objective, supports_val_loss). giant/train.py (1875 lines) becomes giant/training/: trainers.py StageSpec + shared StageTrainer base + the two subclasses metrics.py MetricSpec, MetricsCollector stage2_inputs.py the pure AR/teacher-forcing tensor helpers, moved verbatim loop.py train() (225 lines, was ~514) + graceful shutdown checkpoint.py build/load, lifted out of train()'s closures The trainers shared ~15 identical constructor arguments and copy-pasted their cosine-warmup lambda, EMA setup, state_dict/load_state_dict, resume_lr and train_mode/eval_mode. StageSpec resolves one stage's config once (constructors go from 24 and 22 keyword arguments to (spec, model, device)), the base class holds the rest, and build_stage_trainers drops from ~100 lines to 15. Metric columns are renamed to a uniform stage/split/metric scheme (stage1/train/loss, stage2/train/d_loss, stage1/lr, stage1/router/entropy, val/loss, ...). Old metrics.csv files and W&B history are not comparable. The checkpoint format is unchanged. BEHAVIOR CHANGE — WGAN best-checkpoint selection. The old code meant to score a WGAN stage on its marginal KL, but the guard `{n: kl for n in wgan_names if n not in val_loss_per_stage}` could never fire: val_loss_per_stage was pre-seeded with 0.0 for every stage, so a WGAN stage contributed a flat 0.0 and the KL was written to metrics.csv without ever influencing best.pt. val_objective now returns it as intended. On the test harness's default flow+wgan config val_loss went from 2.182 (stage 1 only) to 15.137 (stage 1 + KL 12.954), and which epoch won changed. Runs before this commit picked their best checkpoint on the non-adversarial stages alone. Written up in docs/v0.3.0-followups.md. Verified: 699 tests pass; ruff, ruff format and ty clean. Baseline-vs- refactor metrics.csv compared across five configs (flow+wgan, AR+onehot, routed, both-flow, AR-flow) — every comparable value bit-identical except val/loss where the fix applies. Resume appends without a duplicate header and reproduces a HEAD worktree's per-epoch losses and LRs exactly across the resume boundary. A refactored last.pt loads through cli.py:_load_model_weights in both raw and ema modes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-07 17:03:20 +02:00