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 : LibTorch.Buffer

  First moment buffer.

• v : LibTorch.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

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

@[instance_reducible]

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

@[instance_reducible]

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

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

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

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

• dtype : LibTorch.Dtype

  Arithmetic and storage of the parameters and moments, not their host coefficient type.

• 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 invalid Adam-family hyperparameters before updating or restoring device state.

Epsilon must remain finite and positive after conversion to the parameter dtype, as in the native kernel.

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

Version 2 header, retained for layouts with independent trainable parameters.

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

Version 3 records storage sharing while retaining the original parameter-slot identifiers.

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

Version 4 also restores the eager session's random-operation counter.

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

Version 5 retains the moment dtype and an optional random counter. Older files are binary32.

def Runtime.Autograd.Torch.Internal.EagerSession.readAdamFormat (handle : IO.FS.Handle) (schema : OptimizerCheckpoint.ParameterSchema) : IO CheckpointIO.Format

Read a supported header and reject a legacy checkpoint when the destination shares parameters.

Version 2 contains no sharing information. Its per-slot histories cannot establish which moment pair belongs to shared storage, so an aliased trainer must start a fresh history or load version 3 or later. Version 4 carries the random-operation counter as well. The check precedes config parsing and all device allocation.

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

Release every device buffer owned by an Adam state map.

def Runtime.Autograd.Torch.Internal.EagerSession.bindAdamConfig (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

Serialize the hyperparameters as a tag plus four raw bit patterns.

Bits rather than decimal text: a checkpoint has to restore the same moments bit for bit, and a decimal round trip of a Float is exactly where that guarantee would be lost.

def Runtime.Autograd.Torch.Internal.EagerSession.readConfig (handle : IO.FS.Handle) (dtype : LibTorch.Dtype := LibTorch.Dtype.float32) : IO CudaAdamConfig

Inverse of writeConfig. An unrecognized kind tag fails loudly rather than defaulting to Adam, since silently reading AdamW moments as Adam moments would corrupt training quietly.

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

Stream CUDA Adam or AdamW moments to disk.

The file includes the optimizer configuration and the complete ordered parameter schema. Shared parameters use version 3 and one moment entry per representative; independent parameters keep the version 2 encoding. For binary32, supplying the session counter selects version 4, which retains both storage sharing and the next random key. Binary64 uses version 5 to retain its dtype and eight-byte moments, with or without a random counter. 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) (expectedConfig : Option CudaAdamConfig := none) (dtype : LibTorch.Dtype := LibTorch.Dtype.float32) : IO (CudaAdamConfig × CudaAdamState × Option ℕ)

Read a whole CUDA Adam checkpoint: metadata, then one moment pair per representative.

The state is built in a local map and only installed once every parameter has been read and checked against schema, so a truncated file leaves the live optimizer untouched.

def Runtime.Autograd.Torch.Internal.EagerSession.loadAdam (path : System.FilePath) (schema : OptimizerCheckpoint.ParameterSchema) (configRef : IO.Ref (Option CudaAdamConfig)) (stateRef : IO.Ref CudaAdamState) (rngCounter? : Option (IO.Ref ℕ) := none) (dtype : LibTorch.Dtype := LibTorch.Dtype.float32) : 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.