CUDA Adam Checkpoints

The eager trainer records parameter leaves before input leaves on every step. Their tape identifiers are the stable zero-based parameter slots used by this internal format. The reader checks every slot against the shared optimizer schema before installing any state.

structure Runtime.Autograd.Torch.Internal.EagerSession.CudaAdamParamState : Type

Device-side Adam moment buffers for one parameter leaf.

• m : Cuda.Buffer

  First moment buffer.

• v : Cuda.Buffer

  Second moment buffer.

• t : ℕ

  Adam step counter for this parameter.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerSession.CudaAdamState : Type

Adam moment state keyed by the eager trainer's parameter-leaf slot.

inductive Runtime.Autograd.Torch.Internal.EagerSession.CudaAdamKind : Type

The update rule associated with a CUDA Adam-family state.

• adam : CudaAdamKind
• adamW : CudaAdamKind

@[implicit_reducible]

instance Runtime.Autograd.Torch.Internal.EagerSession.instBEqCudaAdamKind : BEq CudaAdamKind

def Runtime.Autograd.Torch.Internal.EagerSession.instBEqCudaAdamKind.beq : CudaAdamKind → CudaAdamKind → Bool

def Runtime.Autograd.Torch.Internal.EagerSession.instReprCudaAdamKind.repr : CudaAdamKind → ℕ → Std.Format

@[implicit_reducible]

instance Runtime.Autograd.Torch.Internal.EagerSession.instReprCudaAdamKind : Repr CudaAdamKind

structure Runtime.Autograd.Torch.Internal.EagerSession.CudaAdamConfig : Type

Hyperparameters that determine the meaning of retained Adam-family moments.

• kind : CudaAdamKind
• beta1 : Float
• beta2 : Float
• epsilon : Float
• weightDecay : Float

def Runtime.Autograd.Torch.Internal.EagerSession.checkpointName : String

Human-readable name used in checkpoint diagnostics.

def Runtime.Autograd.Torch.Internal.EagerSession.CudaAdamConfig.sameBits (left right : CudaAdamConfig) : Bool

Compare configurations by exact floating-point bit pattern.

def Runtime.Autograd.Torch.Internal.EagerSession.CudaAdamConfig.validate (config : CudaAdamConfig) : Except String Unit

Reject non-finite or mathematically invalid Adam-family hyperparameters.

def Runtime.Autograd.Torch.Internal.EagerSession.cudaAdamCheckpointFormat : CheckpointIO.Format

Versioned header for CUDA Adam and AdamW checkpoints.

def Runtime.Autograd.Torch.Internal.EagerSession.releaseCudaAdamState (state : CudaAdamState) : IO Unit

Release every device buffer owned by an Adam state map.

def Runtime.Autograd.Torch.Internal.EagerSession.ensureCudaAdamConfig (configRef : IO.Ref (Option CudaAdamConfig)) (expected : CudaAdamConfig) : IO Unit

Bind an in-memory moment map to one optimizer configuration.

Changing Adam to AdamW, or changing a moment-defining hyperparameter, requires a fresh state rather than silently reinterpreting existing moments.

def Runtime.Autograd.Torch.Internal.EagerSession.writeConfig (handle : IO.FS.Handle) (config : CudaAdamConfig) : IO Unit

def Runtime.Autograd.Torch.Internal.EagerSession.readConfig (handle : IO.FS.Handle) : IO CudaAdamConfig

def Runtime.Autograd.Torch.Internal.EagerSession.writeCudaAdamStateFloat32 (path : System.FilePath) (schema : OptimizerCheckpoint.ParameterSchema) (configRef : IO.Ref (Option CudaAdamConfig)) (stateRef : IO.Ref CudaAdamState) : IO Unit

Stream CUDA Adam or AdamW moments to disk.

The file includes the optimizer configuration and the complete ordered parameter schema. Saving before the first Adam-family update is rejected because no optimizer state has yet been defined.

def Runtime.Autograd.Torch.Internal.EagerSession.readCheckpoint (handle : IO.FS.Handle) (schema : OptimizerCheckpoint.ParameterSchema) : IO (CudaAdamConfig × CudaAdamState)

def Runtime.Autograd.Torch.Internal.EagerSession.readCudaAdamStateFloat32 (path : System.FilePath) (schema : OptimizerCheckpoint.ParameterSchema) (configRef : IO.Ref (Option CudaAdamConfig)) (stateRef : IO.Ref CudaAdamState) : IO Unit

Restore CUDA Adam or AdamW moments after validating their optimizer and parameter metadata.

The old state remains live until the replacement has been parsed completely. Duplicate ids, frozen or out-of-range parameters, zero step counters, trailing bytes, and truncated payloads are rejected.