xnn.common.models.recompute#
Recomputation blocks whose memory stays bounded through a double backward.
torch.utils.checkpoint.checkpoint() drops a block’s intermediates after
the forward pass and recomputes them in the backward pass. That bounds memory
for a single backward. Under create_graph=True – which
ForceStressOutput uses in training mode so
that a force loss can be differentiated – the recomputation is recorded and
the intermediates of every block are retained again until the second
backward; the same happens when the backward is wrapped in a checkpoint,
because the autograd engine evaluates the derivative formulas on its own
thread, outside the checkpoint’s saved-tensor hooks. On a 1500-atom cell the
checkpointed three-body term needed 1.6 GB in eval mode and 47 GB in training
mode either way.
recompute() closes that gap with two explicit autograd functions:
_RecomputeBlockevaluates the block without recording and saves only its inputs. Its backward is the second function._RecomputeGradcomputes the block’s gradient in a local autograd graph that is freed on exit (forward), and, when a second backward asks for it, re-runs block and gradient once more withcreate_graph=Trueand contracts with the incoming cotangents, again locally (backward).
Memory is one block at every stage; the price is one extra evaluation of the block per order of differentiation. Third derivatives are not supported. Eager only: custom autograd functions do not script.
Functions
|
Evaluate |