Adam Optimizer

Adam is the project's parameter-update algorithm. It is separate from the decoder-only transformer: the model produces logits, cross-entropy and autograd produce gradients, and Adam consumes those gradients to replace the trainable leaf values.

One optimizer step has this boundary:

textexisting parameter values
        ↓ forward
logits → cross-entropy → backward
        ↓ parameter gradients
global gradient-norm clipping
        ↓ clipped gradients
bias-corrected Adam update
        ↓
new parameter values

Adam does not choose batches, run the forward or backward passes, repeat steps, write metrics, or decide when training is finished. Those responsibilities belong to the training loop described in TRAINING.md.

Public interface#

The optimizer is declared in:

textinclude/riftco_transformer/optim/adam.hpp

AdamOptions contains:

A typical caller constructs the optimizer from the model's named parameters:

cppriftco_transformer::AdamOptions options;
options.learning_rate = config.learning_rate;
options.beta1 = config.adam_beta1;
options.beta2 = config.adam_beta2;
options.epsilon = config.adam_epsilon;
options.maximum_gradient_norm = config.gradient_clip;
// Optional bounded-page layout; ordinary Adam defaults to Contiguous.
options.state_storage =
    riftco_transformer::AdamStateStorageKind::Paged;
options.page_size = 4096;

riftco_transformer::Adam optimizer(model.parameters(), options);

The optimizer copies the ParameterList. Each NamedParameter owns a ParameterHandle to canonical shared parameter state and exposes .parameter as a raw-pointer compatibility view. Adam therefore updates the same state as the model rather than copies of its tensors, and its copied list keeps that state alive independently of the caller's list or wrapper lifetime.

The remaining public operations are:

textoptions()                 configured hyperparameters
backend()                 intrinsic backend of parameters and moment state
step_count()              number of successful updates
parameter_tensor_count()  number of registered parameter tensors
state_storage_kind()      contiguous or bounded-page moment layout
state_page_size()         configured maximum elements per state page
state_page_count()        total paired state pages; zero in contiguous mode
state_payload_bytes()     logical bytes in both FP32 moment vectors
step()                    clip gradients and apply one update
zero_gradients()          explicitly clear every registered gradient
global_gradient_norm()    measure any ParameterList without updating it

parameter_tensor_count() counts parameter tensors, not the total number of scalar values. global_gradient_norm() is generic over ParameterList; it does not require a decoder model and can measure a complete registered module, an explicit LoRA-only list, or a custom parameter collection.

Per-parameter state#

For every registered parameter scalar, Adam owns two FP32 state values:

• the first moment, \(m\), is an exponential average of gradients;
• the second moment, \(v\), is an exponential average of squared gradients.

Both start at zero on the same backend as the registered parameters. Contiguous mode stores each moment as one tensor with the parameter's shape; paged mode partitions both flattened moment vectors into paired 1D tensors. The optimizer also stores:

• the number of successful steps, \(t\);
• the current power \(\beta_1^t\);
• the current power \(\beta_2^t\).

The parameter list is fixed when the optimizer is constructed. Moment tensors remain aligned with that list for the optimizer's lifetime.

Contiguous and bounded-page state#

AdamStateStorageKind::Contiguous is the general API default. It allocates one first-moment tensor and one second-moment tensor for each parameter. The page_size option remains validated and observable but does not partition state in this mode; state_page_count() returns zero.

AdamStateStorageKind::Paged flattens each parameter logically and stores paired first/second-moment pages of at most page_size elements:

textparameter values:  [──────────────────────────────]
first moment:      [ page 0 ][ page 1 ][ page 2 ]
second moment:     [ page 0 ][ page 1 ][ page 2 ]
                         each page ≤ page_size

An update still computes one global gradient norm and one pair of bias- correction factors for the complete parameter list. It prepares complete next- parameter values, dispatches one page at a time, retains every successful page only as candidate state, and commits values, moments, beta powers, and the step counter only after every page succeeds. A later-page failure therefore leaves all live optimizer and parameter state unchanged.

Paging changes allocation and dispatch granularity, not the Adam equations or logical memory size. Both layouts retain exactly

bytes for \(N\) trainable scalars. state_payload_bytes() reports that logical payload, so it is identical between layouts for the same parameter list. state_page_count() counts paired page records, not the two moment tensors inside each record.

For transactionality, step() still prepares one complete next-value tensor for each parameter before committing. The page-size bound applies to moment state and each paged backend update request, not to every temporary allocation in the optimizer.

CUDA page tensors are ordinary managed allocations and may migrate between host and device under the CUDA runtime. This implementation has no explicit eviction, host/disk spill budget, prefetch policy, out-of-core scheduler, or general OS page-fault manager. “Paged” means bounded state allocations and bounded update requests here.

Python post-training keeps ordinary Full and LoRA runs contiguous by default. QLoRA defaults to paged state because its frozen NF4 base has no Adam moments and only the floating-point adapter state is updated. PostTrainingConfig.optimizer_state="auto" makes that selection; "contiguous" and "paged" are explicit overrides. Direct C++ callers may select the same storage policy through AdamOptions.

State diagnostics#

The C++ optimizer exposes state_storage_kind(), state_page_size(), state_page_count(), and state_payload_bytes(). The stable C ABI provides the matching rt_adam_state_* queries, and Python exposes these properties:

pythonwith Adam(
    parameters,
    state_storage="paged",
    page_size=4096,
) as optimizer:
    print(optimizer.state_storage)       # "paged"
    print(optimizer.state_page_size)     # 4096
    print(optimizer.state_page_count)    # paired pages across parameters
    print(optimizer.state_payload_bytes) # two FP32 moments per scalar

The Python post-training artifact records the same values under metadata key optimizer_state. These diagnostics describe the temporary optimizer used for that run; model bundles still exclude moment contents and are not resumable training checkpoints. Use Adam.state()/Adam.load_state() through a .riftckpt TrainingCheckpoint when the moments, counters, RNG, and batch position must resume together.

Global gradient clipping#

Before updating any moment, Adam measures one norm across every scalar gradient in every registered parameter tensor. If \(g_i\) is the raw gradient of scalar element \(i\), the global norm is:

The implementation accumulates this norm with a scaled sum-of-squares algorithm in double. This avoids overflow or underflow that a direct sum of float squares could cause.

Let \(G_{\max}\) be maximum_gradient_norm. One shared scale \(s\) is chosen:

Each scalar gradient is then clipped with the same multiplier:

A zero norm falls into the “otherwise” case and uses \(s=1\). Clipping is global rather than per tensor, so it preserves the direction of the complete model gradient. The clipped gradient \(\widetilde{g}\)—not the original gradient \(g\)—is used to update both Adam moments.

Adam equations#

For successful step \(t\), Adam applies the following equations independently to each parameter scalar:

Because both moment tensors start at zero, their early values are biased toward zero. Adam corrects that bias:

Finally, it updates the parameter:

The symbols mean:

The square root covers only \(\widehat{v}_t\); \(\epsilon\) is added afterward. The stored beta powers start at \(1\). They are multiplied before each update, so the first successful call uses \(\beta_1^1\) and \(\beta_2^1\) for bias correction.

The portable reference performs intermediate update arithmetic in double and checks each stored float. CUDA performs the same candidate calculation with device double intermediates and reports a non-finite candidate through one shared device flag. Metal uses native float arithmetic for its normal fast path and similarly reports unsafe arithmetic. Accelerator parity tests use numerical tolerances rather than requiring bit equality.

Apple GPU float arithmetic may flush subnormal intermediates. The precise-math kernel checks its actual operation sequence for subnormal, non-finite, and ill-conditioned cancellation results. One shared device flag requests a retry when native float is unsafe. After the command completes, that rare batch is evaluated by the portable double reference directly over the host-visible shared candidate buffers. This preserves the optimizer contract, backend identity, and transactional commit boundary; it does not move tensors to CPU storage or add a duplicate host pass to ordinary fused updates.

All parameter values and moments are written to out-of-place candidate tensors. Dispatch rejects a candidate that aliases any live buffer or another candidate, including aliases across parameter tensors in the same batch. No live optimizer state changes while a kernel is running. Only after the complete batch succeeds are candidates moved into the parameters and moment slots.

Backend execution#

Adam captures one intrinsic backend when it is constructed and requires every parameter value, gradient, and moment to remain on it. Move a model before constructing its optimizer:

cppmodel.to(ExecutionBackend::Metal);
Adam optimizer(model.parameters(), options);

The global norm remains an overflow-safe double host reduction. Metal tensors use shared storage, so this does not require copying every gradient into a temporary buffer. On the fused fast path, a clipping scale below the scalar float range is decomposed into a mantissa and exponent and reconstructed with ldexp; scale 1 is passed through directly. If the kernel observes that any resulting Adam intermediate leaves the normal float range or suffers severe cancellation, the reference retry described above handles the entire batch.

In contiguous mode, the safe, well-conditioned Metal path uses one fused kernel per parameter tensor. Each thread computes the clipped gradient, both new moments, both bias corrections, and the new parameter value. All parameter dispatches are encoded into one command buffer and synchronized once:

textvalidate and measure all gradients
        ↓
allocate every candidate value and moment tensor
        ↓
encode fused update for parameter 0 ... parameter N
        ↓ one commit and wait
check command status + shared reference-request flag
        ├── safe: keep fused candidates
        └── unsafe: rewrite all candidates with wide reference
        ↓
move every candidate into live state

Paged mode sends one page-sized update through the same Metal capability at a time. Each page is synchronous, but the outer optimizer keeps its result as candidate state; it does not publish an earlier page if a later page fails.

CUDA tensors use managed allocations. In contiguous mode, the CUDA Adam module preflights every parameter, gradient, moment, and candidate native handle, launches one grid-stride update kernel per parameter tensor, then synchronizes once for the batch. In paged mode, each page is one bounded synchronous update request over managed page tensors; the outer optimizer preserves the same whole-step commit boundary. Each thread uses double intermediates before checked conversion to the stored float. A shared device status flag turns any non-finite candidate into an exception before the public optimizer commits live values or moments. Managed allocation permits CUDA runtime migration but does not add explicit spilling or an application-level page-fault manager.

TPU calls the same portable reference implementation over its host mirror, contiguously or one state page at a time. Neither the native Metal nor CUDA candidate update changes the separate global-norm boundary: gradient-norm clipping remains host control flow on all backends.

This is per-tensor fusion in one submission, not a flattened multi-tensor kernel. Flattened parameter arenas can be considered later if profiling shows dispatch overhead matters.

Step statistics#

step() returns an AdamStepStats value:

The norm and scale are reported as double. A successful all-zero first step still increments the step count and reports a norm of 0 and a clip scale of 1.

Gradient consumption and reset#

step() consumes the gradients currently stored in the registered Parameter objects. During commit it calls Parameter::set_value() for every new value. Replacing a leaf value preserves its identity and shape but resets its gradient tensor to zeros. Consequently, gradients inspected after a successful optimizer step are zero rather than the gradients used for that step.

zero_gradients() is also available when the caller needs to discard gradients without changing parameter values or optimizer moments.

There are two different zero-gradient cases:

1. With a fresh optimizer, zero gradients and zero moments produce a zero update, so parameters remain unchanged.
2. After an earlier nonzero gradient, a later zero gradient decays the stored moments. The remaining first moment can still produce a parameter update.

The second behavior is standard Adam momentum. “Zero gradients leave parameters unchanged” is therefore guaranteed only while the relevant optimizer state is also zero.

Validation and failure behavior#

Construction rejects:

• an empty parameter list;
• an empty or duplicate parameter name;
• a null or duplicate Parameter*;
• a non-finite initial parameter value;
• a gradient whose shape differs from its parameter value;
• parameter values or gradients split across different backends;
• a non-finite initial gradient;
• an invalid or non-finite optimizer option;
• an unrecognized state-storage kind; or
• a zero page_size, even when contiguous state is selected.

Each step rechecks gradient shapes, backend identity, finite gradients, and finite parameter values. Moving a parameter after optimizer construction is therefore rejected before an update. A step also rejects counter overflow and any moment or parameter result that cannot be represented as a finite float.

The update is transactional. Norm validation, candidate allocation, and alias validation happen before device work. In paged mode, every page dispatch must succeed before the prepared parameter tensors and candidate moment pages are committed. A backend command failure, a later-page failure, or a Metal wide- reference retry failure leaves parameter values, gradients, moments, beta powers, and the successful-step counter unchanged.

Lifetime and computation-graph rules#

Adam retains the owning handles from its copied ParameterList. The canonical parameter state therefore outlives an originating Parameter wrapper or module that leaves scope. Moving a wrapper also resolves to the same canonical identity, so aliases are still detected as duplicates.

The public .parameter field is a compatibility view, not independent ownership. Code that extracts that raw pointer must keep a corresponding NamedParameter, ParameterHandle, or optimizer alive while using it. Parameter shape and backend must not change while Adam exists. Moving a parameter to another backend after optimizer construction is rejected before an update because its moment state remains on the original backend.

Parameter values must be updated only after the current backward pass has finished. Once step() replaces the leaf values, the old computation graph is consumed; the next forward pass must build a new graph from the updated values.

Milestone 7 isolated and proved this single transaction:

textprepared gradients → clip → update moments → replace parameter values

Milestone 8 now orchestrates it repeatedly:

textchoose batch → forward → loss → backward → Adam step → record metrics → repeat

The next iteration builds a fresh graph from the newly replaced leaf values. Keeping the milestones separate let the optimizer equations and state be verified independently from batch scheduling and training-control behavior.