Skip to content

Checkpoint format: extend WCHK flag bits, retire 4-variant brute-force auto-loader #154

Description

@atech-hub

The current `load_checkpoint_auto` tries 4 combinations of `(learnable_ode, use_corrector)` to guess at a checkpoint's Dims, matching by param count. It's fragile: every new flag that affects serialisation doubles the search space, and today we already have at least 5 flags affecting param count (learnable_ode, use_corrector, learnable_attn, use_dyn_harmonics, use_rk4_weights).

Current WCHK v4 flag bits (`src/common/checkpoint.rs`):

  • bit 0: learnable_ode
  • bit 1: use_layer_scale
  • bit 2: use_rk4_weights
  • bit 3: use_dyn_harmonics

Missing bits (must be inferred via brute-force today):

  • bit 4: learnable_attn
  • bit 5: use_corrector
  • bit 6: phase_native
  • bit 7: tied_embeddings

Fix:

  1. Add the missing bits to WCHK v4 flags. Backward-compat: old checkpoints without these bits get sensible defaults (use_corrector=true, phase_native=false — matches the pre-phase-native era).
  2. Rewrite `load_checkpoint_auto` to read flags from the header and construct `Dims` directly — no variant guessing.
  3. Simplify `common/generate.rs` which has its own 4-variant loop for the same reason.

Surfaces affected: generate, wave-generate, serve, analyze, ode-monitor, phase-decode, galaxy-scan, scale-checkpoint — all use `load_checkpoint_auto` or their own variant loops. One fix benefits all.

Motivation: discovered while debugging the wave+split-band training — a checkpoint with `phase_native=false` had 175,864 params (lm_head ghost), auto-loader couldn't match, panicked with `"Cannot match checkpoint param count 175864 to any model variant"`. The checkpoint structure was unambiguous from the header if we'd encoded `phase_native` in the flag bits.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions