Autograd Engine

ggmlR includes a PyTorch-style dynamic autograd engine built on top of ggml tensors. Gradients are computed via reverse-mode AD on a tape recorded inside with_grad_tape({}).

Tensor layout: all autograd tensors are column-major — shape is [features, batch], matching ggml’s native C layout.

library(ggmlR)

1. Core primitives

`ag_tensor` and `ag_param`

# ag_tensor — non-trainable input (e.g. data)
x <- ag_tensor(matrix(1:6 / 6, nrow = 2L))   # [2, 3]

# ag_param — trainable parameter (gradient accumulated)
W <- ag_param(matrix(rnorm(4), nrow = 2L))    # [2, 2]
b <- ag_param(matrix(0.0, 2L, 1L))

Forward pass and gradient tape

with_grad_tape({
  h    <- ag_relu(ag_add(ag_matmul(W, x), b))   # [2, 3]
  loss <- ag_mean(ag_mul(h, h))                  # scalar MSE-like
})

grads <- backward(loss)   # returns named list of gradients

cat("dL/dW:\n"); print(grads[["W"]])
cat("dL/db:\n"); print(grads[["b"]])

backward() returns gradients keyed by the parameter’s variable name.

2. Built-in operations

Function	Description
`ag_matmul(A, B)`	Matrix multiply
`ag_add(A, B)`	Element-wise add (broadcast supported)
`ag_mul(A, B)`	Element-wise multiply
`ag_relu(x)`	ReLU activation
`ag_sigmoid(x)`	Sigmoid activation
`ag_tanh(x)`	Tanh activation
`ag_softmax(x)`	Softmax (column-wise)
`ag_sum(x)`	Sum all elements → scalar
`ag_mean(x)`	Mean all elements → scalar
`ag_transpose(x)`	Transpose [m,n] → [n,m]
`ag_reshape(x, dims)`	Reshape tensor
`ag_slice_rows(x, from, n)`	Rows `from+1..from+n` (0-based `from`); gradient zero outside
`ag_concat_rows(a, b)`	`rbind(a, b)`, same number of columns; gradient split back by rows
`ag_softmax_cross_entropy_loss(logits, y)`	Fused softmax + cross-entropy

3. `ag_sequential` — module API

ag_sequential stacks layers into a callable module with .forward() and .parameters().

data(iris)
set.seed(42)

x_all <- t(scale(as.matrix(iris[, 1:4])))        # [4, 150]
y_oh  <- model.matrix(~ Species - 1, iris)
y_all <- t(y_oh)                                  # [3, 150]

idx  <- sample(150L)
x_tr <- x_all[, idx[1:120]];  x_vl <- x_all[, idx[121:150]]
y_tr <- y_all[, idx[1:120]];  y_vl <- y_all[, idx[121:150]]

model <- ag_sequential(
  ag_linear(4L,  64L, activation = "relu"),
  ag_batch_norm(64L),
  ag_dropout(0.3),
  ag_linear(64L, 32L, activation = "relu"),
  ag_linear(32L,  3L)
)

params <- model$parameters()
cat("Parameter tensors:", length(params), "\n")

ag_layer_norm() is the other normalization available. The two differ in what they average over: batch norm normalizes each feature across the batch, so its result depends on the other examples in the batch and it needs running statistics for inference; layer norm normalizes each example across its own features, so it behaves identically in training and evaluation and works with a batch of one. That independence from batch size is why transformers use it.

model_ln <- ag_sequential(
  ag_linear(4L,  64L, activation = "relu"),
  ag_layer_norm(64L),
  ag_linear(64L, 32L, activation = "relu"),
  ag_linear(32L,  3L)
)
cat("LayerNorm model parameters:", length(model_ln$parameters()), "\n")

4. Optimizers

# Adam (default lr = 1e-3)
opt <- optimizer_adam(params, lr = 1e-3)

# SGD with momentum
opt_sgd <- optimizer_sgd(params, lr = 0.05, momentum = 0.9)

Training loop

BS <- 32L
n  <- ncol(x_tr)

ag_train(model)   # set training mode (enables dropout, batch norm train)
set.seed(1)

for (ep in seq_len(150L)) {
  perm <- sample(n)
  for (b in seq_len(ceiling(n / BS))) {
    idx <- perm[seq((b-1L)*BS + 1L, min(b*BS, n))]
    xb  <- ag_tensor(x_tr[, idx, drop = FALSE])
    yb  <- y_tr[, idx, drop = FALSE]

    with_grad_tape({
      loss <- ag_softmax_cross_entropy_loss(model$forward(xb), yb)
    })
    grads <- backward(loss)
    opt$step(grads)
    opt$zero_grad()
  }

  if (ep %% 50L == 0L)
    cat(sprintf("epoch %3d  loss %.4f\n", ep, loss$data[1]))
}

5. LR schedulers

opt2 <- optimizer_adam(params, lr = 1e-3)

# Cosine annealing: lr goes from lr_max to lr_min over T_max epochs
sch_cos <- lr_scheduler_cosine(opt2, T_max = 150L, lr_min = 1e-5)

# Step decay: multiply lr by gamma every step_size epochs
sch_step <- lr_scheduler_step(opt2, step_size = 30L, gamma = 0.5)

# SGDR warm restarts: each cycle T_mult times longer than the last.
# T_mult only takes effect with restart = TRUE.
sch_sgdr <- lr_scheduler_cosine(opt2, T_max = 20L, lr_min = 1e-5,
                                restart = TRUE, T_mult = 2L)

# 1cycle: lr rises then falls, momentum (or Adam beta1) moves inversely
sch_1cyc <- lr_scheduler_onecycle(opt2, max_lr = 1e-2, total_steps = 150L)

# Cyclic: triangular by default, also "triangular2" and "exp_range"
sch_cyc <- lr_scheduler_cyclic(opt2, base_lr = 1e-4, max_lr = 1e-2,
                               step_size_up = 25L)

# Linear warmup, then cosine decay
sch_warm <- lr_scheduler_warmup_cosine(opt2, warmup_steps = 10L,
                                       total_steps = 150L)

# Call after each epoch:
# sch_cos$step()

6. Gradient clipping

with_grad_tape({
  loss <- ag_softmax_cross_entropy_loss(model$forward(ag_tensor(x_tr)), y_tr)
})
grads <- backward(loss)

# Clip global gradient norm to max_norm -- rescales every gradient by one
# factor, so their relative directions are preserved.
clip_grad_norm(params, grads, max_norm = 5.0)

# Element-wise alternative: clamp each value into [-clip_value, clip_value].
# Blunter than norm clipping, but it bounds every element individually.
clip_grad_value(params, grads, clip_value = 1.0)

# Catch NaN/Inf gradients where they appear, rather than after the weights
# have already been poisoned. action = "warn" (default), "stop" or "silent";
# max_abs additionally flags gradients that are merely oversized.
check_grad_anomaly(params, grads, action = "warn")

opt$step(grads)
opt$zero_grad()

7. Dataloader

ag_dataloader shuffles and batches column-major data matrices:

dl <- ag_dataloader(x_tr, y_tr, batch_size = BS, shuffle = TRUE)

ag_train(model)
for (ep in seq_len(100L)) {
  for (batch in dl$epoch()) {
    with_grad_tape({
      loss <- ag_softmax_cross_entropy_loss(model$forward(batch$x), batch$y$data)
    })
    grads <- backward(loss)
    opt$step(grads);  opt$zero_grad()
  }
}

8. Eval mode and inference

ag_eval(model)   # disables dropout, switches batch norm to inference stats

# Forward in chunks to avoid memory pressure
predict_cm <- function(mod, x_cm, chunk = 64L) {
  n   <- ncol(x_cm)
  out <- NULL
  for (s in seq(1L, n, by = chunk)) {
    e  <- min(s + chunk - 1L, n)
    lg <- mod$forward(ag_tensor(x_cm[, s:e, drop = FALSE]))$data
    ev <- exp(lg - apply(lg, 2, max))
    sm <- ev / colSums(ev)
    out <- if (is.null(out)) sm else cbind(out, sm)
  }
  out
}

probs <- predict_cm(model, x_vl)          # [3, 30]
preds <- apply(probs, 2, which.max)
truth <- apply(y_vl, 1, which.max)
cat(sprintf("Val accuracy: %.4f\n", mean(preds == truth)))

9. Raw `ag_param` — full manual control

For complete flexibility, build the network from scratch:

set.seed(7)
W1 <- ag_param(matrix(rnorm(64*4) * sqrt(2/4),  64, 4))
b1 <- ag_param(matrix(0.0, 64, 1))
W2 <- ag_param(matrix(rnorm(3*64) * sqrt(2/64),  3, 64))
b2 <- ag_param(matrix(0.0,  3, 1))

forward <- function(x)
  ag_add(ag_matmul(W2, ag_relu(ag_add(ag_matmul(W1, x), b1))), b2)

opt_raw <- optimizer_adam(list(W1=W1, b1=b1, W2=W2, b2=b2), lr = 1e-3)

for (ep in seq_len(200L)) {
  perm <- sample(n)
  for (b in seq_len(ceiling(n / BS))) {
    idx <- perm[seq((b-1L)*BS+1L, min(b*BS, n))]
    xb  <- ag_tensor(x_tr[, idx, drop = FALSE])
    yb  <- y_tr[, idx, drop = FALSE]
    with_grad_tape({ loss_r <- ag_softmax_cross_entropy_loss(forward(xb), yb) })
    gr <- backward(loss_r)
    opt_raw$step(gr);  opt_raw$zero_grad()
  }
}

10. Mixed precision

# f16 on GPU, f32 on CPU
device <- tryCatch({ ag_device("gpu"); "gpu" }, error = function(e) "cpu")
ag_dtype(if (device == "gpu") "f16" else "f32")

# All subsequent ag_param / ag_tensor use the selected dtype

See also vignette("gpu-vulkan", package = "ggmlR") for device management.

11. GPU residency

Tensors stay in device buffers across a whole training step instead of moving to the host after every operation. This is not a setting — it is how ag_* already works — but it changes what two things are, so it is worth knowing.

Why it exists. A naive GPU autograd re-uploads the weights on every single operation, then downloads each result. The arithmetic was never the bottleneck: uploading weights accounted for 74–98% of the forward pass. Four things now stay resident:

what	lifetime
weights	across steps
Adam moments (`m`, `v`)	across steps
gradients	within a step
forward activations	within a step

A training step went from 10 host/device crossings and 0.188 MB to 4 and 0.047 MB — forward 3, backward 1, optimizer step 0. The Adam step no longer touches the host at all. Training is numerically unchanged; this is a transport change, not a different computation.

ag_tape_memory() shows the memory side of it:

# Section 10 left the dtype at f16 on a GPU machine; the tape report is about
# bytes held, so pin f32 to keep the figures comparable across machines.
ag_dtype("f32")

x_r <- ag_tensor(matrix(rnorm(16 * 32), nrow = 16L))
W_r <- ag_param(matrix(rnorm(8 * 16) * 0.1, nrow = 8L))
y_r <- ag_tensor(matrix(0, 8L, 32L))

with_grad_tape({
  out_r  <- ag_relu(ag_matmul(W_r, x_r))
  loss_r <- ag_mse_loss(out_r, y_r)
})
backward(loss_r)

ag_tape_memory(top = 5)

Two rules this imposes

$grad may be a device handle, not a matrix. Reading it through the normal accessors is unaffected, but arithmetic directly on the slot fails — an ag_handle has no arithmetic methods. That is deliberate: failing loudly beats computing something plausible from a pointer.

A resident gradient dies at the next tape. It lives in the pass pool, which is cleared when the next with_grad_tape() opens. The ordinary loop is fine — the optimizer consumes gradients in the step that produced them — but holding one across a tape boundary does not work:

# WRONG: g1's buffer is gone by the time it is used
with_grad_tape({ l1 <- f(x1) }); backward(l1); g1 <- W$grad
with_grad_tape({ l2 <- f(x2) }); backward(l2); g2 <- W$grad
total <- g1 + g2

Materialise inside the pass that produced it, before the next tape opens.

Switching it off

Both parts are on by default and restore the previous behaviour exactly:

ggmlR:::ag_backward_graph(FALSE)     # backward as closures, not one graph
ggmlR:::ag_backward_resident(FALSE)  # download gradients into R matrices

Both also respond to an environment variable, which is the right way to switch them off for a whole session without touching code:

GGMLR_AG_BACKWARD_GRAPH=0 GGMLR_AG_RESIDENT_GRADS=0 Rscript your-script.R

A third part, a deferred forward that computes the whole forward as one graph, is off by default and enabled with GGMLR_AG_DEFER=1. It is gated because its benefit is uneven — substantial on small and mid-sized models, close to nothing where the weight upload dominates.

Residency lowers the size at which the GPU beats the CPU; it does not remove it. Below roughly d = 128 the dispatch cost still outweighs the arithmetic.

Device buffers

A table written a few columns at a time and read in random batches (a replay buffer in off-policy RL) is kept on the device with ag_buffer(): a write sends only the new columns, and ag_get_rows() sends only the indices. The same calls run on an R matrix on the CPU, as here.

ag_device("cpu")
set.seed(1)
buf <- ag_buffer(rows = 6, capacity = 50)       # one column per transition
ag_buffer_write(buf, matrix(rnorm(6 * 20), 6, 20), col_offset = 0)

f <- ag_capture(function(idx) {
  x <- ag_get_rows(buf, idx)                    # 6 x B, no gradient
  list(obs = ag_slice_rows(x, 0, 3),            # rows 1..3
       act = ag_slice_rows(x, 3, 2))            # rows 4..5
}, params = list(buf))
batch <- f(idx = matrix(c(0, 7, 19), ncol = 1))
batch$act

ag_buffer_read(buf, col_offset = 0, k = 2)      # download a range
ag_buffer_free(buf)                             # release now

Indices are 0-based and select columns (the name follows ggml_get_rows). Indices given from R are checked; inside a GPU recording they are clamped to 0..capacity-1. List the buffer in params of ag_capture() so that freeing or re-uploading it records the graph anew. Switching to the CPU moves GPU buffers to the host first, which needs about twice their device size.

12. Data-parallel training

For multi-GPU or faster single-GPU training see dp_train():

# Full example: inst/examples/dp_train_demo.R

See also vignette("data-parallel-training", package = "ggmlR").