Ocannl.CheckpointStopping a training loop and resuming it so that it continues exactly as the uninterrupted run would have takes two kinds of state:
training_state;progress.save publishes both in one checkpoint file (a Persistence checkpoint whose header carries the progress), in one atomic write, so a crash mid-save leaves the previous checkpoint whole and never a state from one step with the progress of another. restore puts the device state back into a freshly built program's context and returns the progress.
A checkpoint holds values, not the program. The resumed process rebuilds the same model and the same routines -- the same construction order, so the tensor nodes get the same ids, which is how the checkpoint matches them -- then initializes the parameters as usual (Train.init_params), restores, and continues the loop from the returned progress: it re-arms the data position, the step-dependent inputs (a scheduled learning rate is set from progress.step, see Train.scheduled_learning_rate) and the shuffling generator.
Host-side helpers that keep state of their own carry it in progress.extra, under keys the loop chooses (gh-ocannl-1726): a mixed-precision run's Mixed_prec.Loss_scaler (its current scale and good-step count, which no device state holds) through Loss_scaler.save_state and restore_state, and each Train.Outlier_detector (its sliding window) through its own pair. The resumed program creates them afresh with the same parameters, and the restores refuse a state saved under other ones. A checkpoint saved without those keys -- any version-1 or version-2 one -- still restores: the format is unchanged, and what the loop does about the missing state is its own decision.
Device-side randomness needs no generator state saved: OCANNL's random operations are counter-based (threefry over the random seed, the node and a counter), so a resumed program with the same seed (fixed_state_for_init, or Ocannl_tensor.Tensor.set_random_seed) whose counters derive from the restored step draws the same numbers as the uninterrupted run. Under another seed it would not, so the checkpoint records the seeds (Ocannl_tensor.Tensor.random_seeds at save: every one the session created, a graph built under a scoped seed included) and restore refuses a program built under different ones.
let state = Checkpoint.training_state ~routines:[ sgd_routine ] loss in
(* ... every [k] steps: *)
Checkpoint.save ~ctx
~progress:{ step; epoch; position; host_rng = Some rng; extra = [] }
state path;
(* ... in the resumed process, after rebuilding the program and [Train.init_params]: *)
let ctx, progress = Checkpoint.restore ~ctx state path in
let rng = Option.value_exn progress.host_rng in
(* continue from progress.step / progress.epoch / progress.position *)type progress = {step : int;Optimizer steps completed.
*)epoch : int;The epoch in progress.
*)position : int;Where in the epoch's data the next batch starts, in the loop's own units.
*)host_rng : Stdlib.Random.State.t option;The host generator the loop draws from (to shuffle the data, say), captured at save through Stdlib.Random.State.to_binary_string; restore returns a fresh generator in the captured state. To checkpoint the global Stdlib.Random generator, pass Stdlib.Random.get_state () and Stdlib.Random.set_state the restored one. (Base's Random.State.t is abstract and cannot be captured: draw from a Stdlib one.)
extra : (string * string) list;Anything else the loop keeps on the host -- an early-stopping counter, a best loss, a loss scaler's or an outlier detector's save_state -- as key-value pairs, returned by restore as saved.
}val training_state :
?extra:Ir.Tnode.t list ->
routines:Context.routine list ->
Ocannl_tensor.Tensor.t ->
Ir.Tnode.t_settraining_state ~routines loss is the device state a resumable checkpoint has to hold: the values of loss.params (the state set, frozen parameters included), plus what each of the loop's routines carries from one run to the next (carried_state), plus ?extra.
The routines are the compiled ones the loop runs, so the optimizer's moments are found however the optimizer was written: an update that reads its own state without the clearing =: prefix declares it persistent (see docs/syntax_extensions.md), and that is what is read here. Nothing is matched by label. A differentiable leaf that is not a parameter but is optimized (an input stepped by Train.sgd_update ~params) is found the same way, through the routine that updates it; ?extra lists the nodes no passed routine carries (it is unrelated to progress.extra, the host-side key-value metadata) (gh-ocannl-1743).
val carried_state : Context.routine -> Ir.Tnode.t_setThe nodes a compiled routine carries between runs: those it writes whose value on entry it also reads -- in the lowered code, the materialized nodes that are both read before written and written (an optimizer moment, a parameter it updates, an accumulator). A node it only reads (a dataset, a host-set learning rate) is an input the program re-establishes; a node it writes before any read (a gradient zeroed each step) carries nothing. Raises Invalid_argument for a routine recording no kernel segments.
val save :
ctx:Context.t ->
progress:progress ->
Ir.Tnode.t_set ->
string ->
unitsave ~ctx ~progress state path writes the state nodes' current values from ctx, with progress, to the checkpoint at path, replacing any file there. The write is atomic (see Persistence.save): a failed save leaves the previous checkpoint intact. The checkpoint also records which nodes the state was, for restore to check. Raises Invalid_argument on an empty state or duplicate progress.extra keys.
val restore : ctx:Context.t -> Ir.Tnode.t_set -> string -> Context.t * progressrestore ~ctx state path uploads the checkpoint's values of the state nodes into ctx and returns the updated context with the saved progress.
Raises if the file is not a training checkpoint (a plain Persistence.save records no progress), and -- before touching ctx -- if state is not the set of nodes the checkpoint was saved from: a smaller set would restore the parameters while silently restarting the rest (say the optimizer moments) from their initialization, a different run. The message lists the nodes on either side. It also refuses, naming both sides' seeds, a program whose random seeds (Ocannl_tensor.Tensor.random_seeds) are not the ones the checkpoint was saved under. A checkpoint from before the seed was recorded (format version 1) is restored with a warning on stderr instead: its seed cannot be checked, and refusing it would strand every such checkpoint although most were saved under the seed the resume uses. Precision and shape mismatches raise as in Persistence.restore.