From 9bfab722ddadbec8e228deeea1fad9bf1fa08b70 Mon Sep 17 00:00:00 2001 From: hugary1995 Date: Thu, 25 Jun 2026 18:26:10 -0500 Subject: [PATCH] Disable AOTAutograd donated buffers for the adjoint; bump to 1.1.3 The adjoint reuses the autograd graph via torch.autograd.grad(..., retain_graph=True) (RecursiveNonlinearEquationSolver.accumulate). On torch>=2.12 AOTAutograd collects "donated buffers" for a torch.compile'd model, and a backward compiled with non-empty donated buffers requires retain_graph=False -- so a compiled model differentiated through the adjoint raises "compiled with non-empty donated buffers". donated_buffer is a ContextVar-backed torch config: AOTAutograd compiles/runs the backward under contextvars contexts where a normal `config.donated_buffer = False` (or config.patch) override is not visible -- those contexts read the config *default*. So this lowers the *default* (the only cross-context lever), scoped to RecursiveNonlinearEquationSolver.__init__ (runs before any solve compiles a backward; merely importing pyzag/neml2 does not touch the global). It is process-global and emits a one-time UserWarning explaining the change and how to revert it. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/conf.py | 2 +- pyproject.toml | 2 +- pyzag/nonlinear.py | 61 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 63 insertions(+), 2 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index ee2af64..2eba94f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -33,7 +33,7 @@ project = "pyzag" copyright = "2025, Argonne National Laboratory" author = "Argonne National Laboratory" -release = "1.1.2" +release = "1.1.3" # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration diff --git a/pyproject.toml b/pyproject.toml index 70aa43a..2bde2b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pyzag" -version = "1.1.2" +version = "1.1.3" authors = [ { name="Mark Messner", email="messner@anl.gov" }, ] diff --git a/pyzag/nonlinear.py b/pyzag/nonlinear.py index a55ff67..fe6d809 100644 --- a/pyzag/nonlinear.py +++ b/pyzag/nonlinear.py @@ -24,11 +24,67 @@ """Basic functionality for solving recursive nonlinear equations and calculating senstivities using the adjoint method""" +import warnings + import torch from pyzag import chunktime +def _disable_donated_buffers() -> None: + """Lower torch's AOTAutograd ``donated_buffer`` default so pyzag's adjoint works + with ``torch.compile``. + + pyzag's adjoint reuses the autograd graph across the reverse sweep via + ``torch.autograd.grad(..., retain_graph=True)`` (see + :meth:`RecursiveNonlinearEquationSolver.accumulate`). On torch>=2.12 AOTAutograd + collects "donated buffers" for a ``torch.compile``'d graph, and a backward compiled + with non-empty donated buffers *requires* ``retain_graph=False``. So a compiled model + differentiated through pyzag's adjoint raises + ``RuntimeError: ... compiled with non-empty donated buffers ...``. + + ``donated_buffer`` is a ``ContextVar``-backed torch config. AOTAutograd compiles and + runs the backward under its own contextvars contexts, where a normal + ``torch._functorch.config.donated_buffer = False`` (or ``config.patch``) override is + not visible -- those contexts read the config *default*. So the only setting that + actually reaches them is the default, which we lower here, process-wide. + + This is deliberately blunt: it lowers a global torch default the first time a + :class:`RecursiveNonlinearEquationSolver` is constructed. It is scoped to actual + pyzag use -- merely importing pyzag (or neml2) does not touch the default, so code + that never runs the adjoint is unaffected. A ``UserWarning`` is emitted once, when it + takes effect. To keep donated buffers enabled, restore the default:: + + import torch._functorch.config as c + c._config["donated_buffer"].default = True + + The consequence of restoring it: any ``torch.compile``'d model differentiated through + pyzag's adjoint will raise the donated-buffer error again. (Donation saves a little + backward memory, but is fundamentally incompatible with ``retain_graph=True``.) + """ + try: + import torch._functorch.config as functorch_config + + # pylint: disable=protected-access + entry = functorch_config._config["donated_buffer"] + except (ImportError, KeyError, AttributeError): + return # older/other torch without the flag -- nothing to disable + + if getattr(entry, "default", False): + entry.default = False + warnings.warn( + "pyzag lowered torch._functorch.config.donated_buffer's default to False " + "process-wide. AOTAutograd 'donated buffers' (torch>=2.12) are incompatible " + "with pyzag's retain_graph=True adjoint over torch.compile'd models, which " + "otherwise raises 'compiled with non-empty donated buffers'. To keep donated " + "buffers enabled, restore " + "torch._functorch.config._config['donated_buffer'].default = True -- " + "compiled models differentiated through the adjoint will then error.", + UserWarning, + stacklevel=3, + ) + + class NonlinearRecursiveFunction(torch.nn.Module): # pylint: disable=W0223,W0246 """Basic structure of a nonlinear recursive function @@ -350,6 +406,11 @@ def __init__( convert_nan_gradients=True, ): super().__init__() + # Lower torch's AOTAutograd donated-buffer default so the adjoint (which uses + # retain_graph=True) works with torch.compile'd models. Scoped here to actual + # solver use, and runs before any solve compiles a backward. Process-global and + # emits a one-time warning; see _disable_donated_buffers. + _disable_donated_buffers() # Store basic information self.func = func