Skip to content

Persist optimizer state in checkpoints (closes #43) - #54

Merged
w4ffl35 merged 2 commits into
masterfrom
fix/nir-opt-checkpoint
Oct 14, 2026
Merged

w4ffl35 merged 2 commits into
masterfrom
fix/nir-opt-checkpoint

Conversation

@w4ffl35

@w4ffl35 w4ffl35 commented Oct 14, 2026 •

Copy link
Copy Markdown
Member

Checkpoints persist weights plus meta but drop SGD velocity and Adam m/v/step, so resuming silently restarts the optimizer (wrong for Adam bias correction).

This change adds an optional v2 optimizers section keyed by node name:

  • SGD entries carry velocity plus optional velocity_bias; Adam entries carry step/m/v plus optional step_bias/m_bias/v_bias.
  • Hyperparameters are not persisted; re-init with the same hyperparameters before load for a bit-identical resume (documented in checkpoint.h).
  • v1 files still load; caller optimizers absent from the file keep their state. Plain loads/save stay weights-only and ignore the new section.
  • Buffer shapes validate against the graph with named errors mirroring the weight-mismatch discipline; every file optimizer entry must match a caller entry.
  • tools/patch_nir_weights.py accepts v1/v2 (weights only, optimizers ignored).

Tests: test_checkpoint gains SGD and Adam save-fresh-load-N-steps bit-identical resume checks (memcmp on weights and state, step counters asserted), v1-compat, plain-ignores-optimizers, file round-trip, and save/load reject coverage. Full ctest passes 41/41 on Debug, Release, and ASan/UBSan locally.

@w4ffl35
w4ffl35 merged commit f6f7fc9 into master Oct 14, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant