Break the sibling cycle when an array is released by assignment - #4453
Break the sibling cycle when an array is released by assignment#4453tudalex wants to merge 4 commits into
Conversation
The outputs of a multi-output primitive hold each other in `siblings`, a reference cycle that `~array()` breaks once it sees the last external reference go. Assignment did not run that check: it replaced the descriptor directly, so an array whose last external reference went away by assignment left the cycle intact, and with it every input below it in the graph. `compile_dfs` does exactly that. It rewires the outputs' inputs to the copied tape by assignment, which drops the trace's own consumers of a Split (or any multi-output primitive), and then rewrites the parents map by assignment, which drops the last reference to the trace's outputs of that Split. From then on the cycle owns the whole trace beneath it, captured constants included, and nothing ever frees it: `compile_erase` and `compile_clear_cache` never reach it. Compiling a function that splits a matmul of a captured weight leaked that weight once per compiled function. In mlx-swift-lm, whose Qwen3.5 decode step is compiled and splits the fused q/k/v projection in every layer, that was 0.5-1.7 GB of a 4B model left resident after the model was dropped, which killed an app that loaded another model next. Route copy assignment, move assignment and `overwrite_descriptor` through the destructor by moving the previous value into a temporary, so the cycle check runs on every way an array reference can end. `a[...] = v` on the last surviving output of a lazy split reaches the same bug with no compile involved: it goes through `overwrite_descriptor`, and with the other output already gone the two held only each other, and the split's input with them. There is a test for each path.
zcbenz
left a comment
zcbenz
left a comment
There was a problem hiding this comment.
Thanks for the fix and I think you found the root cause, but I need to think for a while about whether there is a better fix that does not requiring patching every place copying array_desc.
About the tests though, using get_active_memory and gc.collect() in tests is very unreliable and was causing a lot of flakes. Can you turn them to C++ tests like this one
Lines 675 to 692 in 90846ad
get_active_memory() and gc.collect() are unreliable in a test and have caused flakes, so track the ArrayDesc's lifetime with a weak_ptr to its data instead, in the style of the existing "test siblings circular references without eval" case. One test drops both outputs of a lazy split by assignment, covering the copy and the move operator; the other leaves a single output alive and replaces it through overwrite_descriptor, which is the path __setitem__ takes and where no compile is involved at all. Both fail on main and pass with the fix, and neither measures memory or touches the garbage collector.
f034f8f to
78de35e
Compare
zcbenz
left a comment
zcbenz
left a comment
There was a problem hiding this comment.
Thanks for updating the tests, I think the PR is ready to go with just a little refactoring.
Like having the array_desc_ destructor trigger this? I'm trying to see if I can make that work, would be even cleaner and prevent future bugs. |
The destructor body becomes array::reset(desc): it swaps desc into array_desc_ and runs the check on the descriptor it lets go of. The destructor, both assignment operators and overwrite_descriptor call it instead of releasing through a temporary array. reset() holds the new descriptor before checking the old one. The order matters when a sibling is assigned over the last outside reference to its own cycle, because the sibling is only reachable through that reference: checked first, the pair looks unreferenced, the cycle is broken and the sibling list the new value lives in is freed before the assignment reads it. Held first, the new descriptor counts as an outside reference and the cycle stays intact, as it did when the check ran after the assignment. The added test segfaults with the other order.
|
Here is an alternative in which we wrap the array_desc_ in a class, and it's destructor triggers the reset logic tudalex@3602800 |
|
Thanks for experimenting with the ideas and sorry I'm being nitpicking here, I would like to keep the code self-explanatory so anyone even without too much C++ expertise can still understand and maintain the code, that's why I don't like the original What do you think about zcbenz@a249d99? I think as long as we are not doing a self-assignment there shouldn't be need to keep an additional copy of assignee? |
reset(desc) becomes a static release(desc): it takes the descriptor the array is letting go of and runs the sibling-cycle check on it, touching nothing else. Callers hand it the old descriptor with std::exchange, which is also what takes the new one first: release(std::exchange(array_desc_, other.array_desc_)); The order is the fix, not the naming. A reset() reading array_desc_ cannot provide it: the member still holds the old descriptor while the check runs, so nothing holds the new one, and when the new one is a sibling of the old the pair looks unreferenced. The cycle is then broken and the sibling list the new value lives in is cleared out from under it, leaving an output with no siblings for outputs() to index and for Split::eval to write. Comparing array_desc_ instead of id() does not catch this, since the two descriptors genuinely differ: they are different outputs of the same split. Self-assignment is a separate case and both guards already cover it. std::exchange copies the new descriptor once and moves the old one out, so this is the single atomic increment a plain copy-assign already does. The assignment operators also use the array_desc_ comparison now, which is cheaper and clearer than comparing id().
2 participants
This PR fixes a memory leak issue that I encountered on iOS when unloading Qwen models.
With the help of AI I tracked it down to the compile function and further down to a corner case in reference counting in MLX's Array class.
When reassigning an array that is produced from a multi output (e.g. the result of a split call), the destructor code is not called, so there is no check to see if it is the last reference to a family.
This fix just makes sure that the destructor is actually called on the previous value when reassigning an array, thus fixing the memory leak.
I think this should also fix issue #3932.
A 4B params Qwen model would leak ~1Gb each time it is unloaded which is quite a lot for a mobile device.
Used to track down the bug and write tests that could reproduce it. All the tests are generated by AI and fail on the main branch. Commit description also generated by AI. Multiple solutions also brainstormed with AI, I considered this one to be the cleanest potential fix.