Runtime Checkpoints

TorchLean examples often want the same simple workflow:

1. train a model for a few steps,
2. save its state, and
3. restore that state before inference or further training.

The implementation delegates binary encoding and device transfers to Runtime.Autograd.Model.StateIO, while this module provides the checked runtime API used by trainers and data loaders.

What Is Supported

Both checkpoint formats preserve the model's shape-indexed state layout:

• Native Float32 modules use exact little-endian binary32 payloads on CPU and CUDA.
• Binary64 Float modules use exact binary64 JSON on CPU and binary32 payloads on CUDA.
• Streamed checkpoints write one tensor at a time instead of materializing a second host copy of a large device-resident model.

This is enough to checkpoint any TorchLean runtime model implemented as a TorchLean.Module.Objective over native Float32 or binary64 Float, independent of architecture.

Results of Trainer.train use Checkpoint.State: trained.save path writes the trained state as Float tensors and trainer.load path data restores it.

Explicit immutable nn.State α checkpoints use Checkpoint.Encoding α. Instances preserve exact bits for Float, Float32, and configured binary formats, including arbitrary supported widths. The complex instance stores both coordinates using their component encoding. Loading checks the format tag as well as the state layout; it never silently casts between scalar formats.

class TorchLean.Checkpoint.Checkpointable (α : Type) [Storage α] : Type

Checkpoint persistence supported by a tensor element type.

• write {shapes : List Shape} : Bool → System.FilePath → Runtime.Autograd.Torch.ParamList α shapes → IO Unit

  Write every parameter and persistent buffer in order.

• read {shapes : List Shape} : Bool → System.FilePath → Runtime.Autograd.Torch.ParamList α shapes → IO Unit

  Validate and restore every parameter and persistent buffer in order.

@[instance_reducible]

instance TorchLean.Checkpoint.instCheckpointableFloat : Checkpointable Float

Binary64 Float checkpoints keep exact binary64 values on CPU and binary32 values on CUDA.

@[instance_reducible]

instance TorchLean.Checkpoint.instCheckpointableFloat32 : Checkpointable Float32

Native Float32 checkpoints preserve exact binary32 payloads on both CPU and CUDA.

def TorchLean.Checkpoint.save {α β : Type} [Storage α] [Storage β] [Context α] [Checkpointable α] {stateShapes inputShapes dataInputShapes : List Shape} (objective : Module.Objective α β stateShapes inputShapes dataInputShapes) (path : System.FilePath) : IO Unit

Save the current values of a TorchLean runtime module.

The module's element type determines its persistence implementation. Native Float32 modules preserve exact binary32 payloads on CPU and CUDA. Binary64 Float modules preserve binary64 values on CPU and the runtime's binary32 values on CUDA.

def TorchLean.Checkpoint.load {α β : Type} [Storage α] [Storage β] [Context α] [Checkpointable α] {stateShapes inputShapes dataInputShapes : List Shape} (objective : Module.Objective α β stateShapes inputShapes dataInputShapes) (path : System.FilePath) : IO Unit

Load a compatible checkpoint into a module.

The reader checks the format header, state-tensor count, every shape, and every payload length before accepting the checkpoint.

def TorchLean.Checkpoint.Optimizer.save {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (objective : Module.Objective α β stateShapes inputShapes dataInputShapes) (path : System.FilePath) : IO Unit

Save optimizer state retained by a module's runtime backend.

This currently applies to eager CUDA Adam and AdamW, whose moment buffers live on the device. The checkpoint also preserves the eager session's random counter, so dropout continues with the next mask after a restore. The operation fails explicitly when the trainer has no backend-owned optimizer state instead of writing an incomplete resume checkpoint.

def TorchLean.Checkpoint.Optimizer.load {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (objective : Module.Objective α β stateShapes inputShapes dataInputShapes) (path : System.FilePath) : IO Unit

Restore optimizer state and its random counter. Legacy files lacking the counter load with a warning: their moments are recoverable, but they cannot reproduce a stochastic continuation.

def TorchLean.Checkpoint.State.save {σ τ : Shape} {α : Type} [Storage α] [Encoding α] (model : nn.Sequential σ τ) (state : nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes model)) (path : System.FilePath) : IO Unit

Save an immutable state pack in its scalar encoding, with the layout supplied by model.

def TorchLean.Checkpoint.State.load {σ τ : Shape} {α : Type} [Storage α] [Encoding α] (model : nn.Sequential σ τ) (path : System.FilePath) : IO (nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes model))

Load immutable model state without mutating a runtime module.

The model supplies the exact dependent state layout, so malformed, missing, extra, or incorrectly shaped tensors are rejected at the boundary.