ag_checkpoint: Run part of the forward pass without storing its activations

In ggmlR: 'GGML' Tensor Operations for Machine Learning

Run part of the forward pass without storing its activations

Description

Executes fn(...) with the gradient tape switched off, so none of the intermediate activations inside fn are kept. During backward() the segment is run again – with the same RNG state – to rebuild what the gradient rules need.

Usage

ag_checkpoint(fn, ...)

Arguments

fn

A function taking the tensors in ... and returning a single ag_tensor. It must be reproducible: given the same inputs and RNG state it has to compute the same thing. Parameters it closes over are fine, and are given gradients as usual. One caveat, because it bites in the most natural usage: fn is called again during backward(), so it must not close over a loop variable. Wrapping layers in a for loop with lyr <- layers[[i]] in the body leaves every segment sharing one binding, and by replay time they all see the last layer – the earlier ones then get no gradient while the loss still looks perfectly normal. Capture per iteration with local(), as in the example below.

...

Input ag_tensors, passed to fn.

Details

This trades compute for memory: the segment's forward pass runs twice, and in exchange the tape never holds its activations. Use it on the deep middle of a network, not on cheap layers.

Value

An ag_tensor holding the segment's output. Its gradient flows back into both ... and any parameters used inside fn.

Examples

block <- ag_linear(8L, 8L, activation = "relu")
x <- ag_tensor(matrix(runif(8 * 4), 8, 4))
y <- ag_tensor(matrix(runif(8 * 4), 8, 4))
with_grad_tape({
  h    <- ag_checkpoint(function(inp) block$forward(inp), x)
  loss <- ag_mse_loss(h, y)
})
backward(loss)

# Checkpointing every second block of a deeper stack. local() gives each
# segment its own binding -- see the note on `fn` for why that matters.
layers <- lapply(1:4, function(i) ag_linear(8L, 8L, activation = "relu"))
with_grad_tape({
  h <- x
  for (i in seq_along(layers)) {
    seg <- local({
      lyr <- layers[[i]]
      function(inp) lyr$forward(inp)
    })
    h <- if (i %% 2L == 0L) ag_checkpoint(seg, h) else seg(h)
  }
  loss <- ag_mse_loss(h, y)
})
backward(loss)

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