ag_estimate_training_memory: Estimate the memory a training run needs

In ggmlR: 'GGML' Tensor Operations for Machine Learning

Estimate the memory a training run needs

Description

Adds up the four things that are alive at the same time during training on the autograd (ag_*) path: the weights, their gradients, the optimizer's per-parameter state, and the activations held by the tape until zero_grad() clears it.

Usage

ag_estimate_training_memory(
  shapes,
  batch_size = 32L,
  optimizer = c("adam", "sgd", "sgd_plain"),
  activation_frac = 0.01,
  quiet = FALSE
)

Arguments

shapes	List of parameter shapes, each a numeric vector of dimensions (e.g. list(c(256, 128), c(128))), or a single such vector.
batch_size	Batch size the activations are sized for. Activations scale with it; weights and optimizer state do not.
optimizer	One of "adam" (keeps m and v), "sgd" (keeps a velocity, i.e. momentum > 0) or "sgd_plain" (no state).
activation_frac	Activations as a fraction of the weight bytes, per unit of batch. The default 0.01 reproduces the measured tape composition for a dense stack (inst/scripts/measure_ag_tape_memory.R: activations were 19% of a 4-layer 128-wide tape at batch 32). Pass your own if you have measured your model with ag_tape_memory, which is always more accurate than this scaling.
quiet	Suppress the printed report and return the figures only. Note that the double-vs-f32 caveat is part of that report, so a caller passing quiet = TRUE takes on reading it from the note field.

Details

Everything is counted at 8 bytes per scalar, because that is what the path actually uses: gradients are accumulated by R closures in double, and optimizer moments are R matrices. ag_dtype affects only forward tensors uploaded into a ggml buffer and does not reach either. An estimate made in f32 would be half the real figure, so the returned list carries a note saying so.

Value

Invisibly, a list of byte counts: params, weights, gradients, optimizer, activations, total, plus bytes_per_scalar and note. Printed as a report.

See Also

ag_tape_memory measures a real tape instead of estimating one; ggml_estimate_memory sizes a single tensor in a ggml buffer, where dtype does apply.

Examples

# a 3-layer MLP, 512 wide, trained with Adam at batch 64
ag_estimate_training_memory(
  shapes = list(c(512, 512), c(512, 512), c(512, 512)),
  batch_size = 64)

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