Autograd Engine

The hardware and bandwidth for this mirror is donated by dogado GmbH, the Webhosting and Full Service-Cloud Provider. Check out our Wordpress Tutorial.
If you wish to report a bug, or if you are interested in having us mirror your free-software or open-source project, please feel free to contact us at mirror[@]dogado.de.

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.

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)

cat("dL/dW:\n"); print(W$grad)
cat("dL/db:\n"); print(b$grad)

backward() accumulates each gradient into the parameter’s $grad. It also returns an environment keyed by tensor id, which is what an optimizer’s $step() takes.

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
`ag_multihead_attention(...)`	Multi-head attention module, one head at a time
`ag_flash_attention(q, k, v, n_heads)`	All heads in one fused op, gradient in one more (section 3.1)

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")

3.1 Fused attention

ag_flash_attention() computes every head in a single ggml_flash_attn_ext() call and its gradient in a single ggml_flash_attn_back(), so the tape holds one node where a per-head loop records dozens. q, k and v are [d_model, seq], with rows split into n_heads contiguous blocks — the same head layout ag_multihead_attention() uses.

set.seed(2)
d_model <- 16L; seq_len <- 8L; n_heads <- 4L

q <- ag_param(matrix(runif(d_model * seq_len, -1, 1), d_model, seq_len))
k <- ag_param(matrix(runif(d_model * seq_len, -1, 1), d_model, seq_len))
v <- ag_param(matrix(runif(d_model * seq_len, -1, 1), d_model, seq_len))

with_grad_tape({
  att  <- ag_flash_attention(q, k, v, n_heads = n_heads, causal = TRUE)
  loss <- ag_mse_loss(att, ag_tensor(matrix(0, d_model, seq_len)))
})
g_att <- backward(loss)
cat("dL/dq:", paste(dim(q$grad), collapse = " x "), "\n")

Projections are not included: apply W_q, W_k, W_v before the call and W_o after it.
k and v may have a different sequence length from q (cross-attention), as long as they match each other.
causal = TRUE lets query j attend to keys 1..j. A custom mask is [seq_kv, seq_q] — keys index the rows, queries the columns, the transpose of the usual convention — either logical (TRUE = attend) or numeric (0 / -Inf). The same mask is applied to the gradient.
On the attention core alone (forward and backward, CPU backend) the fused op is 15–36× faster than the per-head loop, depending on size.

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, 2, 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(); a full example is in inst/examples/dp_train_demo.R.

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

These binaries (installable software) and packages are in development.
They may not be fully stable and should be used with caution. We make no claims about them.
Health stats visible at Monitor.

Autograd Engine

1. Core primitives

ag_tensor and ag_param

Forward pass and gradient tape

2. Built-in operations

3. ag_sequential — module API

3.1 Fused attention

4. Optimizers

Training loop

5. LR schedulers

6. Gradient clipping

7. Dataloader

8. Eval mode and inference

9. Raw ag_param — full manual control

10. Mixed precision

11. GPU residency

Two rules this imposes

Switching it off

Device buffers

12. Data-parallel training

`ag_tensor` and `ag_param`

3. `ag_sequential` — module API

9. Raw `ag_param` — full manual control