File size: 2,230 Bytes

644ccd9

# W9 — expose the training-speed knobs the hill-climb needs, as Tyro flags on `zipdepth-train-catalog`, on branch `zdp/8-train-fast`
# (already created from zdp/7-eval-tools and checked out; the idiomatic dataloader from W8 lands later on this branch — do not touch
# zipdepth/catalog/dataset.py or the dataset construction beyond what a flag needs). Worktree:
# /home/pablo/.paseo/worktrees/26sr02x6/dirty-swan. REPORT: /tmp/zdp-recut/apply/W9-REPORT.md. Rules: python-conventions
# (PEP 526, jaxtyping, no Any, names never carry dtype/shape), AGENTS.md, no vendored edits, loss/architecture/eval unchanged.

APPLY (each default-off or default = today's behaviour so the baseline recipe is bit-for-bit today's):
1. `--compile-mode {off,default,reduce-overhead,max-autotune}` (torch.compile of the student for training; today's behaviour
   stays the default; document that beartype dev mode disables it), `--channels-last` (model + batches to channels_last
   memory format), `--fused-adamw` (torch.optim.AdamW(fused=True)), `--amp-dtype {bfloat16,float16}` if not already a flag.
2. `--max-lr` already exists; add `--warmup-pct` and `--final-div-factor` passthroughs to the OneCycle schedule, and
   `--schedule {onecycle,constant-cooldown}` where constant-cooldown = constant lr then linear cooldown over the last
   `--cooldown-pct` of total_steps.
3. Progressive resolution: `--stage-schedule "384x512:0.3,768x1024:0.7"` (fractions of total_steps per resolution; the loader
   rebuilds at a stage boundary; checkpoints carry the stage). Skip if the idiomatic dataset cannot change resolution mid-run
   without a restart — say so; the orchestrator can then run stages as separate resumed arms.
4. `--preset {v4,fast}`: named recipes that set the above (v4 = today's; fast = filled in after the hill-climb) — the winning
   recipe lands here later.
5. Tests: each flag has a unit test proving it changes the constructed model/optimizer/scheduler and that the default equals
   today's objects (bit-equal state_dict / same class + kwargs). Gates: zipdepth-catalog-dev lint typecheck deadcode tests;
   the live smoke `zipdepth-catalog-train-smoke` must pass with `--compile-mode default --channels-last --fused-adamw`.