ag_capture_step: Capture a whole training step as reusable device graphs

In ggmlR: 'GGML' Tensor Operations for Machine Learning

Capture a whole training step as reusable device graphs

Description

Like ag_capture, for a training step: the forward pass in fn, its backward pass and the Adam update are recorded once, and every later call is upload the inputs, two graph computes, download the outputs – no R-level op dispatch. On the CPU the wrapper runs the same step as ordinary code (zero_grad, tape, backward, step).

Usage

ag_capture_step(fn, params = NULL, optimizer, max_grad_norm = NULL)

Arguments

fn	Function of named ag_tensor arguments returning the loss (an ag_tensor), or a named list of ag_tensors with an element loss.
params	Optional list of further ag_tensors the recording references (e.g. frozen weights) and ag_buffers read with ag_get_rows; the optimizer's parameters are always included.
optimizer	An optimizer from optimizer_adam, created on the GPU when the step is to run there.
max_grad_norm	NULL (no clipping) or a positive number: clip gradients to this global norm before the update.

Details

Two graphs, not one: forward with backward, then the update. Adam writes each weight in place, and a weight is also read inside the first graph (by the forward, and by the backward of the layer below) through tensors that share its memory but are not connected to the update in the graph – so a single graph would leave the order of that read and that write to the backend.

fn is called WITHOUT with_grad_tape: the wrapper records the tape itself. Everything ag_capture requires of fn applies: data only through the arguments, no reading values back, weights on the device. The backward must be one that runs as a graph (the ops the graph backward supports).

Input names and shapes are fixed by the first call; a call with other shapes is an error (pad a short last minibatch, or run it uncaptured). The step counter and the learning rate are read from optimizer on every call, so learning-rate schedules work, and optimizer$state() stays valid. Outputs are the values of the forward pass, before the update, as in an ordinary training loop. After a captured step $grad is not set.

Supported: optimizer_adam without gradient accumulation, F32 compute dtype.

Gradient clipping by global norm (max_grad_norm): the norm of all gradients is computed in the first graph and every gradient is multiplied by min(1, max_grad_norm / (norm + 1e-6)) before the update. Branch-free, unlike clip_grad_norm: for a norm within about 1e-6 below max_grad_norm the factor is slightly under 1 instead of exactly 1. The pre-clipping norm is returned as an extra output grad_norm. A non-finite norm is an error raised BEFORE the update: weights, moments and the step counter are left as they were. The moments accumulate clipped gradients, so optimizer state saved with and without clipping is not interchangeable on resume. Change the limit between calls with ag_capture_set (no re-recording).

Value

A function taking the same named arguments as numeric matrices and returning the loss as a matrix, or a named list of matrices.

See Also

ag_capture, ag_capture_set, ag_capture_free

Examples

W <- ag_param(matrix(rnorm(6, sd = 0.1), 2, 3))
opt <- optimizer_adam(list(W), lr = 1e-2)
train <- ag_capture_step(function(x, y) ag_mse_loss(ag_matmul(W, x), y),
                         optimizer = opt)
for (i in 1:5) loss <- train(x = matrix(1, 3, 4), y = matrix(0, 2, 4))

ggmlR documentation built on Oct. 23, 2026, 1:07 a.m.