larsandClaude Opus 5
8019a80563
CI / Format (ruff format) (push) Successful in 27s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 39s
CI / Type check (ty) (push) Successful in 43s
CI / Format (ruff format) (pull_request) Successful in 31s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 31s
CI / Tests (push) Successful in 2m35s
CI / Tests (pull_request) Successful in 2m36s
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>
54 lines
2.9 KiB
Markdown
54 lines
2.9 KiB
Markdown
# v0.3.0 — post-implementation audit: open discrepancies
|
||
|
||
**Status:** steps 1–7 of `docs/v0.3.0-design.md` §12 are implemented (branch
|
||
`v0.3.0-stage2-autoregressive`, commits `eb6dd27`..`200c6d2`). This document
|
||
tracked discrepancies found between that implementation and the design contract
|
||
during a 2026-08-07 audit, as concrete work items. **All items (1-9) are now
|
||
resolved** — either implemented (1-5, 7-9) or explicitly deferred into
|
||
`docs/v0.3.0-design.md` §11.2 (6: `stage2_model.generator = "ddpm"`). Step 8
|
||
(`estimate_batch_size` recalibration) was intentionally still outstanding per
|
||
§12 and was never tracked here.
|
||
|
||
---
|
||
|
||
## Confirmed correct during the audit (no action needed)
|
||
|
||
For reference — these were explicitly checked against the design doc and
|
||
match it, including two spots the doc itself flagged as likely stale that
|
||
turned out fine:
|
||
|
||
- Config schema, `DEFAULT_CONFIG`, `migrate_config` table (§4), `save_config`/
|
||
`merge_cli_overrides` recursion, `default_out_dir_name`, `Conditioning` enum,
|
||
`n_sec.mode = "stop_token"` error (§9, §11.2).
|
||
- `network.py`'s full class decomposition (§5.3), dict-returning
|
||
`build_models`/`build_critics` (§5.4), `ExpertTrunk` separate in/out dims,
|
||
ST-Gumbel wiring (§2.1), AR token layout (§6.1), Markov/Attention history
|
||
encoders (§6.2).
|
||
- Shared PDG top-N type map (one map, not two, per §8), `other_policy`
|
||
sample/modal/drop (§11.1), embedding L1-nearest decode, and the L1-distance
|
||
diagnostic surfaced in `giant analyze` (§11.3).
|
||
- `analysis/render.py`/`analysis/router_gating.py` correctly branch
|
||
old-flat vs new-nested `model_config["router"]` location — doc flagged this
|
||
as a likely stale spot (§10) but it's actually fine.
|
||
- `giant/training/`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence,
|
||
per-stage router auxiliary losses, stage-prefixed metrics, stage-2-only
|
||
training via ground-truth `x1_s1` (§7).
|
||
|
||
## Behavior change: WGAN best-checkpoint selection
|
||
|
||
The `giant/training/` split fixed a dead guard in WGAN validation scoring. A
|
||
WGAN stage was *meant* to contribute its marginal KL to the `val_loss` that
|
||
drives `best.pt`, but the guard `if n not in val_loss_per_stage` could never
|
||
fire (every stage was pre-seeded to `0.0`), so the stage contributed a flat
|
||
`0.0` and the KL was written to `metrics.csv` without ever being used.
|
||
`WGANStageTrainer.val_objective` now returns the KL as intended.
|
||
|
||
**Consequence:** any checkpoint selected before this commit under a config
|
||
with a WGAN stage — including the v0.3.0 default (`stage2_model.generator =
|
||
"wgan"`) — picked its best epoch on the non-adversarial stages alone. Measured
|
||
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. Do not compare `val/loss` or `best.pt` choice across this commit.
|
||
- `pipeline.py`'s deleted wgan+router rejection, per-stage `centers_init`
|
||
seeding, removed stale expert-size warning (§9, §10).
|