es-ude / es-ude/OnDeviceTraining

training-state checkpointing: optimizer/scheduler/RNG state is not serializable — resume silently corrupts training

Open
#350 0 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.