Do not propagate donatable marks onto results - #1814
Merged
Conversation
`Nx.Defn.Expr.expr/4` builds every node by copying the shape, type and names of an existing tensor, which also copied its `donatable?` mark. Any value computed from a donated argument therefore came back marked, even though its buffer was freshly written. The mark then escaped the call and donated buffers the caller never offered. The mark describes a concrete buffer, so it is only meaningful on a parameter, which stands for an argument the caller handed us. EXLA reads it off those parameter nodes to decide what to alias, so they keep it and every other node clears it. That leaves the case where a parameter is returned as is, where the output template is the parameter itself. EXLA now clears the marks when building the output template, so results are never donatable regardless of how they were produced. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
polvalente
approved these changes
Aug 15, 2026
seanmor5
added a commit
to elixir-nx/axon
that referenced
this pull request
Aug 15, 2026
elixir-nx/nx#1814 stops donatable marks from propagating onto results, so a donating step no longer returns marked tensors and the loop does not have to clear them itself. Repins nx and exla to that commit and adds a test asserting the state a donating loop returns is not donatable, since that is the property the removal depends on. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
seanmor5
added a commit
to elixir-nx/axon
that referenced
this pull request
Aug 18, 2026
* Donate step state buffers during training Adds a `:donate_state?` option to `Axon.Loop.run/4`. When set, the loop marks the `:model_state`, `:optimizer_state`, and `:loss_scale_state` entries of the step state with `Nx.donatable/1` before each step, so supporting compilers can write the new values into the old buffers instead of allocating alongside them. The marks are applied before the strict compile, since Nx requires the compile-time templates and the runtime arguments to agree on what is donated. Only those three entries are donated because EXLA raises when a donated argument is not consumed by the computation, and the loop overwrites `:y_true` and `:y_pred` without ever reading them. A donatable mark propagates onto every result derived from a donated argument, so the loop clears the marks off each step's output and reapplies them itself. Otherwise it would donate the metrics threaded into the next iteration, and the mark would follow the caller out of the loop. On a 74k parameter model trained with Adam for 3 epochs, peak memory growth drops from ~8.6MB to 0.5MB, against 2.3MB for `force_garbage_collection?: true`. Requires the unreleased buffer donation support in Nx, so nx and exla are pinned to main for now. That also makes the deprecated `Nx.Defn.Kernel.hook/3` fatal under `--warnings-as-errors`, so the two call sites move to `io_call/3`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Drop the donatable mark sweep now that Nx clears it elixir-nx/nx#1814 stops donatable marks from propagating onto results, so a donating step no longer returns marked tensors and the loop does not have to clear them itself. Repins nx and exla to that commit and adds a test asserting the state a donating loop returns is not donatable, since that is the property the removal depends on. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When you mark a tensor with
Nx.donatable/1, you are saying one thing: "I am done with this buffer, feel free to write on top of it." That is a promise about that one tensor, and nothing else.Right now the mark spreads. Every expression node is built by copying the shape, type, and names off a tensor that already exists, and the donatable flag was riding along in that copy. So if you donate
xand then computeNx.add(x, 1), the answer comes back marked too, even though it lives in a brand new buffer that nobody offered up. The mark then follows that tensor around, and the next time you jit something it tries to donate a buffer you never meant to give away.I ran into this while adding donation to Axon's training loops. After one donating training step, every parameter and every piece of optimizer state came back marked. Axon had to sweep the marks off by hand after each step, otherwise it would donate its own metrics on the following step, and the marks would follow you out of the loop into whatever you did with your trained model afterward.
The fix is small. The mark is about a real buffer, so it only makes sense on a parameter, which is the stand in for an argument you handed in. EXLA reads the flag off those parameter nodes to work out what it can alias, so parameters keep it and everything else clears it.
There is one leftover case. If a function just hands an argument straight back, the output is the parameter itself, so the mark was still riding out with it. EXLA now clears the marks when it builds the output template, so results are never donatable no matter how they were made. The argument is still really donated, its buffer just backs the result now.
Tests cover both, on the Nx side and the EXLA side, and I added a line to the
Nx.donatable/1docs saying results are never donatable.One test in
exla/test/exla/device_memory_sharing_test.exsfails on my machine, but it fails the same way on main without this change, so it is not from here.🤖 Generated with Claude Code