es-ude / es-ude/OnDeviceTraining
training-state checkpointing: optimizer/scheduler/RNG state is not serializable — resume silently corrupts training
- Dominant language
- C
- Stars
- 1
- Forks
- 3
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 8
Description
## Verified gap (2026-07-13)
`StateDictApi` (`modelLoadStateDict`) covers **weights/biases only**. Nothing else of the training state is serializable or restorable:
- **Optimizer states**: SGD momentum buffers; after #328 also AdamW's m/v moment tensors **and the `stepCount` scalar**.
- **Scheduler state** (#327): `lastEpoch`, `baseLr` — a resumed run silently restarts the LR trajectory at `baseLr`.
- **Stochastic-rounding RNG streams** (SR_HALF_AWAY, #279): xorshift state is process-lifetime; a resume replays or diverges the dead-zone escape statistics.
- **DataLoader position / shuffle epoch context.**
## Failure modes (all silent — everything "runs")
- Momentum reset: transient, usually survivable.
- LR reset: the schedule restarts fast-start on a late-training model — silently wrong trajectory, corrupted experiment comparability.
- **AdamW `t=0` on resume is the dangerous one**: bias correction `1/(1-beta^t)` re-amplifies the first steps massively; applied to a converged model this can blow it up. Nobody sees an error.
## Why this matters in months, not years
On-device training on MCUs practically implies intermittent power — checkpoint/resume is a core scenario, not a nice-to-have (energy harvesting, duty-cycled devices). The first long-running or intermittent-power demo will hit every one of these at once, and retrofitting serialization across `optimizer_t`/`states_t`/`lrScheduler_t`/RNG is far more expensive after consumers exist.
## Proposal
1. Treat training-state serialization as a **design requirement** gating the first intermittent/long-running example — not an afterthought.
2. Scope sketch: extend the state-dict concept to (optimizer type, per-parameter state buffers, impl scalars incl. `stepCount`, scheduler `{type, baseLr, lastEpoch, params}`, RNG stream state). PyTorch's `optimizer.state_dict()` / `scheduler.state_dict()` split is the parity reference.
3. Near-term (PR C, #328): document `stepCount` as non-persisted in `AdamW.h` so the limitation is at least stated where it hurts most.
Related: #327 (scheduler state), #328 (AdamW state), #279 (SR streams), serialization module in `src/serial/`.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Contributor guide
Research direction
Start with the serialization module in src/serial/, then read AdamW.h and the related issues #327, #328, and #279 to map the existing state gaps. Define the required persisted optimizer, scheduler, RNG, and data-loader state, with completion demonstrated by resume behavior that preserves training state rather than silently restarting it.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- c
- Domain
- embedded-iot, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100