Recognise a buffer as mutated when the write is not its direct user - #21979
Recognise a buffer as mutated when the write is not its direct user#21979john-rocky wants to merge 2 commits into
Conversation
tag_constant_data decided whether a buffer is mutated by looking through its users for the node that performs the write. A cache that is first concatenated with its new values and written back further down the chain has the concat as its user, not the write, so it was taken for constant data and tagged into the delegate. A backend that compiles mutable buffers into state then produced a model that asks the runtime for state — after the partitioner had been told, by take_over_mutable_buffer=False, that this runtime has none. On the Core ML backend that is a load that succeeds and an execute that fails with 'The input feature for layers_N_conv_conv_state must be an MLState, but it was not'. The signature already records which buffers are mutated, so match on the target instead of on the mutating node's name. The direct-user check stays for parameters and lifted constants, which have no target to match.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21979
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
|
Correction to my own PR description, before anyone spends time on it. I wrote that LFM2.5-350M "runs where it previously failed" and answers So the c4w build produces wrong text. I should not have quoted a single lucky prompt as What this does and does not change. The partitioning fix is unaffected: the state error What I do not yet know is whether the wrong text is the 4-bit weights or something else Please review the change on its own terms — it fixes a misclassification with a test — and |
|
Correcting the correction — and this time in the change's favour, which is why I want to My five-prompt check ran every prompt through one loaded method. The KV cache persists Reloading the method per prompt, LFM2.5-350M through the Core ML delegate with no Exact agreement with the unquantised reference on every counted position. So the model this The 4-bit builds are a separate matter and are lossy, which is what I had been seeing Summary of what is measured, so the record is straight:
Apologies for the churn on this thread. Both wrong claims came from the same harness bug, |
The reel's acts and the LLM's backend both come from Documents now, so the same install records the Core ML take and the XNNPACK take. Before pytorch/executorch#21979 the Core ML side could not run at all. Two things the text act needed. It now waits on the model instead of reporting it missing: opening the Core ML build compiles it, the act arrives long before that finishes, and the old code read a not-yet-open model as a not-present one. And the act deadline is minutes rather than 45 seconds, because a deadline shorter than the compile turns a first run into a blank screen.
|
@claude review this code |
|
Claude finished @metascroy's task in 1m 31s —— View job Code reviewI read the change against The fix is rightThe bug is real. The new code compares on the value side instead: mutated_targets = set(buffers_to_mutate.values())
...
if buffers_map.get(node.name) in mutated_targets:
mutated_buffer.add(node)
continue
One correctness nuance worth a note (non-blocking)
The sibling function has the same latent pattern
for node_user in node.users:
if node_user.name in buffers_to_mutate:
is_mutated_buffer_node = TrueThat's the exact pattern you just fixed in The test is good
Two minor test nits (optional):
On the thread churnThe load/execute error, the measurement correction, and the final 27/27 unquantised agreement all point to this being a genuine partitioning fix independent of the separate 4-bit-weight quality issue — agreed that the 4-bit results don't belong to this PR. Good discipline posting the correction rather than leaving the wrong numbers standing. Verdict: LGTM. Correct root-cause fix, matches on the right key, keeps the old path where it's still needed, and ships with a test that fails without the change. Only real follow-up is the mirror pattern in |
|
@john-rocky did you confirm this PR fixes the "The input feature for layers_7_conv_conv_state must be an MLState, but it was not." issue in the motivation? Generally that error I have seen when you run a model with states on a binary compiled for < macOS15/iOS18. |
| # this hold for a buffer that is not mutated in one step: a cache that is first | ||
| # concatenated with the new values and written back further down the chain has | ||
| # something other than the mutation as its user, and used to be taken for a | ||
| # constant. It would then be handed to a delegate as a buffer, and a backend that |
There was a problem hiding this comment.
Can you remove references to CoreML in the comment
There was a problem hiding this comment.
Done. Splitting the function for the C901 lint moved the explanation into the helper's docstring, and the backend-specific wording — the Core ML mention and take_over_mutable_buffer — is gone.
|
The PR seems reasonable, but can you take a look as well @JacobSzwejbka since it's editing exir |
Review asked for two things here: the comment explaining the target-matching should not name a specific backend, and flake8 C901 flagged tag_constant_data at complexity 13 against the limit of 12. Moving the detection loop into _mutated_buffer_placeholders answers both — the explanation moves into the helper's docstring with the backend-specific wording dropped, and both functions sit below the complexity limit. Tagging behaviour is unchanged.
|
@metascroy Yes — and to make sure the answer doesn't rest on the LFM2.5 runs alone, I re-confirmed it today with a controlled A/B. Your diagnosis of when the error fires matches what I measure. It takes a model with state features running on a binary that does not bind states. On the same pip executorch 1.4.0 pybind runtime (macOS arm64), a genuinely stateful model — a directly written buffer, What the PR fixes is why a nominally stateless export was stateful at all. With The A/B, single variable: same script, same binary, same macOS, only the class ConvCache(torch.nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("conv_state", torch.zeros(1, 8, 4))
self.conv = torch.nn.Conv1d(8, 8, 5)
def forward(self, x):
h = torch.cat([self.conv_state, x], dim=2) # the buffer's user is the cat
self.conv_state.copy_(h[:, :, 1:]) # the write lands downstream
return self.conv(h)exported with
So the fix removes the error by removing the unintended states, not by teaching the runtime to bind them — a model that is supposed to have states still needs a macOS15/iOS18-capable binary, exactly as you said. Happy to fold this two-cause distinction into the PR description if that's useful. |
|
On |
|
This fix to tag_constant_data was landed by someone else here: #22172 (Thanks for originally finding it @john-rocky!) There is a similar fix needed for mutable buffers here: #22296 |
|
Thanks @metascroy for the pointer, and thanks @usamahz for landing the fix. #22172 fixes ご教示ありがとうございます @metascroy さん、修正を入れてくださった @usamahz さんにも感謝します。#22172 が本 PR と同じ方式(signature のミューテーション対象との照合)で |
3 participants
Fixes #21855.
What happens today
tag_constant_datadecides whether a buffer is mutated by walking its users and lookingfor the node that performs the write:
That holds only when the buffer is written in one step. A cache that is concatenated with
its new values and written back further down the chain has the
catas its user, not thewrite, so it is taken for constant data and tagged into the delegate.
A backend that compiles mutable buffers into state then produces a model that asks the
runtime for state — after the partitioner was told, by
take_over_mutable_buffer=False,that this runtime has none. On Core ML that is a load which succeeds and an execute which
fails:
Measured on LFM2.5-350M at partition time, ten of its twenty-two mutated buffers were
tagged and twelve were not. The twelve are the KV caches, written by
index_putdirectly.The ten are the short-convolution states, written a step later — and they are exactly the
buffers named in the error.
The change
The signature already records which buffers are mutated, so match on the target rather
than on the mutating node's name. The direct-user walk stays for parameters and lifted
constants, which have no buffer target to match.
Checking it
The new test in
test_partitioner.pyfails onmain('tag0' is not None: a mutated buffer must not be tagged as constant data) and passes with the change.End to end,
export_llmwith the Core ML backend now runs where it previously failed atexecute. LFM2.5-350M, c4w, macOS arm64:
"The capital of France is"" Paris", matching the unquantised modelArgmax agrees with the original model on every position where the context determines the
token; the two that differ are the first, where the distribution after
<bos>alone isnearly flat.
I also checked that nothing else moves: the other buffer-mutation and constant-data tests
in that file pass,
take_over_mutable_buffer=Truebehaves as before, and a mutable-buffermodel lowered to XNNPACK produces identical outputs with and without the change, cache
advancing correctly across calls.
One thing this does not fix, for completeness: a Qwen3.5-0.8B c4w build now runs but its
output is wrong, which looks like the 4-bit weights rather than the partitioning. I have
not shipped that one and will report it separately once I know which.