pablovela5620's picture
chain checkpoints, project notes, model card section (part 4)
644ccd9 verified
|
Raw History Blame Contribute Delete
2.23 kB

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.