Repository navigation
Only rewrite addmm/mm to linear when the weight is a constant - #22402
msluszniak wants to merge 12 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22402
Note: Links to docs will display an error until the docs builds have been completed.
|
meta-cla
Bot
added
the
CLA Signed
Backport of pytorch/executorch#22402. replace_addmm_mm_with_linear rewrote mm(x, transpose(w)) into linear(x, w) without checking that w is a constant, so a matmul against a runtime-computed matrix became a linear whose weight the Vulkan runtime then tried to prepack at delegate build time: (isTensorRef()) is false! Expected value to have type TensorRef, got TENSOR instead. Guard both rewrites on the operand being a parameter, buffer or lifted constant.
This PR needs a
|
SS-JIA
left a comment
SS-JIA
left a comment
There was a problem hiding this comment.
AI-assisted review. Inline findings below.
| return True | ||
| return ( | ||
| is_param(exported_program, node) | ||
| or is_buffer(exported_program, node) |
There was a problem hiding this comment.
[High] is_buffer() also accepts buffers listed in graph_signature.buffers_to_mutate. Those values are state, not build-time constants; rewriting a matmul that reads one into Vulkan linear can prepack an old value and return wrong results after the buffer changes. Exclude mutable buffers here and add a multi-invocation regression.
AI reviewer note: this comment was generated by an AI reviewer.
| # Without the owning program a lifted parameter cannot be told apart | ||
| # from a user input. Placeholders were always rewritten before, so keep | ||
| # accepting them rather than regressing callers that pass no program. | ||
| return True |
There was a problem hiding this comment.
[Medium] Returning true here leaves user-input weights unprotected for callers that omit the program. Samsung currently constructs AddmmToLinearTransform() without one, and XNNPACK calls the helper without one, so this pass still rewrites runtime placeholder weights there. Thread the owning program through those call sites, or conservatively reject placeholders when it is unavailable.
AI reviewer note: this comment was generated by an AI reviewer.
|
Mutated buffers are now rejected, and the owning program is threaded through Vulkan, Samsung and XNNPACK so the no-program case can refuse placeholders instead of accepting them; @nil-is-all could you help find a reviewer? |
|
@pytorchbot label "module: backend" |
pytorch-bot
Bot
added
the
module: backend
`replace_addmm_mm_with_linear` turns `mm(x, transpose(w))` into `linear(x, w)` whenever the second operand is fed through a transpose. It never checks where `w` comes from, so a matmul against a tensor computed at runtime is rewritten into a linear too. That rewrite is not sound. `mm`/`addmm` place no constraint on their second operand, but backends prepack a linear's weight while building their delegate graph, which is only possible for a constant. The Vulkan runtime aborts when it gets anything else: Exception raised from toTensorRef at backends/vulkan/runtime/graph/containers/Value.h:266: (isTensorRef()) is false! Expected value to have type TensorRef, got TENSOR instead. reached through linear_packed_weight -> prepack_fp_linear_weight -> PrepackNode, at delegate build time. A model that multiplies an activation by a runtime-derived matrix fails to load with no indication of which op is responsible. Guard both rewrites on the operand actually being a parameter, buffer or lifted constant, and thread the owning program through so that check can be made. Callers that pass no program keep the previous behaviour for placeholders and only lose the rewrite for runtime-computed operands, which is the case that was broken anyway. The pass is used by the Vulkan and Samsung ENN backends.
is_buffer() also accepts a buffer the program writes back to. That is state, not a build-time constant: a delegate that prepacked it would keep returning the value it was built with. Reject buffers named in buffers_to_mutate. Accepting placeholders when no program was passed left user-input weights unprotected for the callers that omit it. Vulkan, Samsung and XNNPACK all have the owning program in scope, so pass it (ConvertToLinearPass becomes an XNNPACKPass, which the pass manager already hands the program), and make the no-program case reject placeholders instead.
a8d79c8 to
3ebf9e2
Compare
A quantized weight reaches mm/addmm through a dequantize node, which is a call_function, so is_constant_tensor read it as a runtime value and the rewrite to linear was skipped. That silently disables downstream quantized-matmul fusion, which matches on linear. Exporting all-mpnet-base-v2 int8 for Vulkan produced 72 dequantize_per_channel plus 72 float addmm instead of 72 _weight_int8pack_mm, dequantizing the weights at runtime and aborting at load because the Vulkan runtime has no dequantize_per_channel kernel. Accept a dequantize node when everything it dequantizes is itself constant. Both mpnet int8 Vulkan exports are then byte-identical to the published artifacts again.
|
Still a WIP? |
ConvertToLinearPass became an XNNPACKPass in this stack, so its constructor now takes the exported program. This call site was still constructing it with no arguments and raised TypeError: XNNPACKPass.__init__() missing 1 required positional argument: 'exported_program' Every llama-family export goes through here, which is why it took out the llama runner, lora, qnn and phi_4_mini jobs along with the export_llama_lib unit tests.
|
@digantdesai No, it is ready for review :)) |
|
LGTM, CI is still red, not sure if its your PR or main is red? Please take a look? |
|
I think the reason why gitlab fails is different than my code changes ;p |
|
However, this one check is on my side: https://github.com/pytorch/executorch/actions/runs/35065274307/job/104849342287?pr=22402, I'm fixing it now. |
EdgeProgramManager.transform() runs its passes over every method, so passing a single program bound one pass to every graph, and exported_program() with no argument looks up "forward". A multimethod export has no such method, so test-lora-multimethod failed with KeyError: 'forward' before any rewriting happened. Use the per-method dict form of transform() so each pass gets the program of the graph it rewrites.
|
@digantdesai it should be fine now. |
executorch-triage
Bot
added
the
community: contribution
|
@claude review this PR |
|
rebased to retrigger the CI |
|
Claude finished @digantdesai's task in 3m 47s —— View job The constant-weight guard is sound, but the local dequant allowlist leaves out Detailed reviewPR Review: #22402SummaryThe PR rewrites Backward Compatibility
Testing
Code Quality
RecommendationRequest Changes. Add |
dequantize_per_channel_group (4-bit groupwise) was missing from the dequant allowlist, so those mm/addmm weights stopped being rewritten.
An unfolded weight reaches mm as dq(q(param)); a quantize node whose inputs are all constant is constant too.
Without the ExportedProgram the pass leaves parameter weights alone; log it so out-of-tree callers notice the missing rewrite.
Rename the computed-weight addmm test, cover a dequant of runtime data, and drop an assertion that only tested EdgeProgramManager's default method.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
@msluszniak want to fix the thing Claude suggested? I can merge it then :) |
|
@digantdesai All the suggestions from claude review had been addressed a that moment ;) |
Summary
replace_addmm_mm_with_linearrewritesmm(x, transpose(w))intolinear(x, w)whenever the second operand is fed through a transpose. It never checks wherewcomes from, so a matmul against a tensor computed at runtime becomes a linear too.That rewrite is not sound.
mm/addmmplace no constraint on their second operand, but backends prepack a linear's weight while building their delegate graph, which is only possible for a constant. The Vulkan runtime aborts when it gets anything else:reached through
linear_packed_weight->prepack_fp_linear_weight->PrepackNode, at delegate build time. The model fails to load, and the message names neither the op nor the operand, so there is nothing to work back from.This shows up in any model that multiplies an activation by a runtime-derived matrix. It was found in a TTS duration predictor whose
mmsecond operand ispermute_copy(squeeze_copy(...)), a 64x64 style-conditioned matrix produced during execution. Two otheraddmmnodes in the same graph have real parameter weights and rewrite correctly; only the computed one is affected.Fix
Guard both rewrites on the operand actually being a parameter, buffer or lifted constant, and thread the owning program through so that check can be made. Callers that pass no program get no rewrite for placeholder weights, parameters included; every in-tree caller passes it, and the pass warns once when it skips one.
The pass is used by the Vulkan and Samsung ENN backends.
Test plan
New
backends/transforms/test/test_addmm_mm_to_linear.py, four cases: a constant weight still rewrites tolinear, a computedmmoperand staysmm, a computedaddmmoperand staysaddmm, and a user-input weight is left alone.Verified against the reported model on a Galaxy S26 Ultra (Adreno 840). Before, the delegate aborts at build time. After, it loads and runs, and all four sub-models of that TTS pipeline match their CPU references:
Confirmed the guard is what changes the outcome: on the unpatched pass the same graph yields
linear=1, mm=0; patched it yieldslinear=0, mm=1.cc @cccclai