Repository navigation
Release the previous value through ~array() on assignment - #4710
Closed
yangzichao wants to merge 1 commit into
Closed
yangzichao wants to merge 1 commit into
yangzichao wants to merge 1 commit into
Conversation
array::operator= only reassigned the shared_ptr, so when the last external reference to an output of an unevaluated multi-output primitive was overwritten in place, the sibling cycle that ~array() knows how to break stayed alive and the whole graph behind it leaked. This is the path mlx-swift's Module.update(parameters:) takes, so loading a checkpoint into a quantized model leaked the quantize graph of every layer on every load. Move the old value into a temporary that dies before the reassignment, for both copy and move assignment, and add a test that fails without it.
yangzichao
added a commit
to yangzichao/mlx-audio-swift
that referenced
this pull request
Oct 16, 2026
…t graph quantize(model:) ran MLX.quantized on every random-init weight; the three sibling outputs keep each other alive and update(parameters:) overwrites them in place, which bypasses the sibling-cycle release in mlx::core::array::~array (ml-explore/mlx#4710). Every load of a quantized model left the unevaluated graph of every quantized layer behind, and the never-evaluated random draws kept growing the global key chain. Build the quantized layers from zero placeholders instead (quantizeWithPlaceholderArrays) and construct models under a throwaway random key. QuantizedEmbeddingFromArrays exists only because mlx-swift 0.31 has no explicit-array QuantizedEmbedding init; drop it at 0.32.2.
This was referenced Oct 16, 2026
ivan-digital
pushed a commit
to soniqo/speech-swift
that referenced
this pull request
Oct 16, 2026
QuantizedTextDecoder built its projections with the random-init QuantizedLinear(in, out, ...) convenience init, then fromPretrained replaced weight/scales/biases through update(parameters:). Nothing in that chain is ever evaluated, and the lazy graph outlived the model: - the three outputs of the lazy quantized(...) op are siblings holding strong references to each other; ~array() breaks that cycle, but update(parameters:) overwrites the arrays in place and array::operator= did not (ml-explore/mlx#4710); - every never-evaluated random draw keeps its link of the global MLXRandom key chain alive. One load/unload of Qwen3-ASR-1.7B-MLX-4bit added 196 fast::Quantize, 697 RandomBits and 5728 ArrayDesc to the heap, plus the scalar Metal buffers feeding them, for the lifetime of the process. Add QuantizedLinear.preQuantized, which builds the layer from lazy zero placeholders through the explicit init (the QuantizedLinear counterpart of PreQuantizedEmbedding), and withThrowawayRandomState, which scopes the remaining init draws to a key that dies with the call. Use them in QuantizedTextDecoder, makeMaybeQuantizedLinear and the Qwen3ASRModel, Qwen3ForcedAligner and SileroVAD loaders. The same heap probe now shows zero growth per reload. Add PreQuantizedLinearTests: placeholder shapes and dtypes match the convenience init, a loaded placeholder computes the same output as a layer built from the loaded arrays, and a regression test that reloads a checkpoint 100 times and asserts Memory.activeMemory does not move. Pointed at the convenience init, that last test fails.
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
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.
Proposed changes
array::operator=only reassignsarray_desc_.~array()knows how to break the reference cycle between the outputs of an unevaluated multi-output primitive once the last external reference goes away, but an array that is overwritten in place never gets there, so the whole unevaluated graph behind those outputs (inputs included) is leaked.That is the path mlx-swift's
Module.update(parameters:)takes (mlx_array_set→*ptr = other). With the common "construct the quantized layer, then load the checkpoint into it" pattern, every quantized layer leaks itsfast::Quantizetriplet plus the random-init graph feeding it, on every model load. In a long-running macOS app that reloads a 1.7B 4-bit model after idle,heapshowed 244kArrayDesc(110 MB) and 33k never-freed Metal buffers after ~40 loads; the counts decompose exactly as loads × quantized layers.The fix moves the old value into a temporary that is destroyed before the reassignment, for both copy and move assignment, so the existing sibling check in
~array()runs. The added test overwrites both outputs of asplitand checks that the graph is gone; it fails onmainand passes with this change. The full C++ suite passes (248/248, CPU build).Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes (formatted the two touched files with clang-format 21.1.8, the version pinned in.pre-commit-config.yaml)