Runtime Checkpoints

TorchLean examples often want the same simple workflow:

1. train a model for a few steps,
2. save its parameters, and
3. reload those parameters later to sample / run inference.

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

What Is Supported

Both checkpoint paths preserve the model's shape-indexed parameter layout:

• CPU modules use exact Float.toBits (UInt64) values stored in JSON.
• CUDA modules stream the float32 values used by the runtime, four bytes per scalar, without first materializing the full model in host tensors.

This is enough to checkpoint any TorchLean runtime model implemented as a TorchLean.Module.ScalarModule over Float, independent of architecture.

def TorchLean.Checkpoint.saveParamBits {paramShapes : List Shape} (path : System.FilePath) (ps : Runtime.Autograd.Torch.TList Float paramShapes) (pretty : Bool := true) : IO Unit

Write a parameter pack to a JSON bits checkpoint.

Use this when the model's shape-indexed parameter tensors are already available and only need to be persisted.

def TorchLean.Checkpoint.saveModuleParams {paramShapes inputShapes natInputShapes : List Shape} (m : Runtime.Autograd.TorchLean.ScalarModule Float paramShapes inputShapes natInputShapes) (path : System.FilePath) : IO Unit

Save the current values of a TorchLean runtime module.

CPU modules use the exact Float.toBits JSON format. CUDA modules use the streamed float32 format so large models remain on the device while their parameters are written. This is architecture-agnostic: it works for any ScalarModule Float ….

def TorchLean.Checkpoint.loadModuleParams {paramShapes inputShapes natInputShapes : List Shape} (m : Runtime.Autograd.TorchLean.ScalarModule Float paramShapes inputShapes natInputShapes) (path : System.FilePath) : IO Unit

Load a JSON-bits or streamed-float32 checkpoint into a module.

The format is detected from its header. Both readers check the parameter count, every tensor shape, and every payload length before accepting the checkpoint.

def TorchLean.Checkpoint.saveOptimizerState {paramShapes inputShapes natInputShapes : List Shape} (m : Runtime.Autograd.TorchLean.ScalarModule Float paramShapes inputShapes natInputShapes) (path : System.FilePath) : IO Unit

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

This currently applies to eager CUDA Adam and AdamW, whose moment buffers live on the device. The operation fails explicitly when the selected trainer has no backend-owned optimizer state instead of writing an incomplete resume checkpoint.

def TorchLean.Checkpoint.loadOptimizerState {paramShapes inputShapes natInputShapes : List Shape} (m : Runtime.Autograd.TorchLean.ScalarModule Float paramShapes inputShapes natInputShapes) (path : System.FilePath) : IO Unit

Restore optimizer state retained by the module's runtime backend.

def TorchLean.Checkpoint.loadParamBits {paramShapes : List Shape} (path : System.FilePath) : IO (Runtime.Autograd.Torch.TList Float paramShapes)

Load a JSON bits checkpoint as a parameter list (without mutating a module).

This is useful when you want to run compiled inference directly and never instantiate a trainer.

def TorchLean.Checkpoint.readParamBits {paramShapes : List Shape} (path : System.FilePath) : IO (Except String (Runtime.Autograd.Torch.TList Float paramShapes))

Read a JSON bits checkpoint, returning an error string instead of throwing an exception.

This is useful in batch tools or CI-style runs where you want to keep going and report failures.

def TorchLean.Checkpoint.loadModelParamBits {σ τ : Shape} (model : nn.Sequential σ τ) (path : System.FilePath) : IO (TensorPack Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model))

Load a bits checkpoint whose parameter layout is determined by model.

def TorchLean.Checkpoint.toRuntimeParams {paramShapes : List Shape} (ps : TensorPack Float paramShapes) : IO (Runtime.Autograd.Torch.ParamList Float paramShapes)

Allocate runtime parameter handles from a checked parameter pack.

def TorchLean.Checkpoint.loadModelIntoModule {σ τ : Shape} (model : nn.Sequential σ τ) (m : Runtime.Autograd.TorchLean.ScalarModule Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model) [σ, τ]) (path : System.FilePath) : IO Unit

Load a checkpoint into a runtime module attached to model.

def TorchLean.Checkpoint.loadModelIntoModuleIfSome {σ τ : Shape} (model : nn.Sequential σ τ) (m : Runtime.Autograd.TorchLean.ScalarModule Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model) [σ, τ]) (path? : Option System.FilePath) : IO Unit

Load a model checkpoint when the caller supplied a path.

def TorchLean.Checkpoint.saveModelIntoPathIfSome {σ τ : Shape} (model : nn.Sequential σ τ) (m : Runtime.Autograd.TorchLean.ScalarModule Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model) [σ, τ]) (path? : Option System.FilePath) : IO Unit

Save a model's runtime parameters when the caller supplied a path.