S-exp GPU
All pages
Docs · Models and optimizersMarkdown

Recomputing a loop body

A looped model keeps every iteration's intermediate values until the backward reads them, so its memory grows with the number of iterations. checkpoint keeps only what crosses an iteration's boundary and computes the rest again when the backward reaches it: activation checkpointing. It has nothing to do with the checkpoints a run saves.

(dotimes (k rounds)
  (setq state (checkpoint (step state))))

(checkpoint body...) returns what its body returns. Each time it is evaluated, the tensor operations its body records become one region:

  • Kept are the values the region reads from outside it, such as the carried state and the parameters, and the values read after it, such as its result.
  • Recomputed is everything else. The backward computes the region's forward again right before the region's own backward, and its gradients read the recomputed values.

The recomputation is the same computation: random draws in the body repeat their values, and the losses and gradients are those of the same file without checkpoint, up to the order a device's plan sums in (determinism). The cost is about one more forward pass of each region per step, paid in time; what it saves is the region's intermediate values, held only while its backward runs.

A body may be any code. One iteration of a loop is the common region:

(dotimes (k rounds)
  ;; one region per round
  (checkpoint
    (dotimes (j updates)
      (setq state (update state))))
  (setq logits (append logits (list (head state)))))

A round is a coarser region than an update: fewer values cross its boundary, so less is kept, and each region holds more while its backward runs. The finest region is not always the smallest plan; a CUDA run reports its planned memory before the first step (devices).

A checkpoint inside another belongs to the outer region. Where no backward runs, in evaluation passes, generators, optimizer bodies and forward-query training, the form only returns its body's value. Values a metric or diagnostic reads inside the region are kept.

A gradient taken inside the body is part of the region, so a loop body that descends a gradient of its own is recomputed whole. A gradient taken outside it reads the region's values as the forward computed them, and they are kept.

A body that records no tensor operation, such as (checkpoint state) or (checkpoint (lambda (s) ...)), is E-GRAD-005: there is nothing to recompute. checkpoint outside a traced graph is E-FLOW-003.