Only fuse an RMS norm whose weight can be prepacked - #22777
Conversation
et_vk.rms_norm prepacks its weight, so the multiply that the fusion folds in has to be a constant the prepacker can see. The pattern folded in any multiply that followed the norm, so a multiplier computed in the graph produced an op the runtime aborted on at the first inference: prepack_standard ... (graph.val_is_tref(tensor_data)) is false! This affects any adaptive normalization - a norm whose scale is produced at inference from a conditioning signal, as in DiT-style AdaLN - and also Gemma's ordinary RMSNorm, whose scale is written `1.0 + weight` and is therefore constant valued but still an intermediate node. Reject a non-prepackable weight in two places, because the two callers see different information. The detector is given only the graph, so it can require the weight to be a placeholder, which covers every computed multiplier. Only the replacement is given the exported program, so the distinction between a constant placeholder and a user input is made there. Not fusing is correct: the norm and the multiply are both supported ops, so the pattern simply stays unfused and costs one extra dispatch. Test Plan: python -m unittest backends.vulkan.test.test_vulkan_passes -v 12 tests pass. The new test, test_rms_norm_fuses_only_prepackable_weight, fails without this change with "1 != 0 : expected 0 fused rms_norm" for both the computed-constant and the graph-input case. Verified end to end on an AMD Radeon 8060S (RADV GFX1151), lowering each module and comparing the delegate's output against eager: multiplier before after leaf parameter fused, ok fused, ok leaf buffer fused, ok fused, ok parameter through a dtype cast prepack abort unfused, ok 1.0 + parameter (Gemma) prepack abort unfused, ok graph input (adaptive norm) prepack abort unfused, ok computed by a layer (AdaLN) prepack abort unfused, ok backends.vulkan.test.test_vulkan_delegate shows an identical set of results before and after the change.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22777
Note: Links to docs will display an error until the docs builds have been completed.
|
|
Hi @giuliocorradi! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
This PR needs a
|
There was a problem hiding this comment.
🟡 Changes recommended
The replacement-level early-return can be counted as a successful replacement by the current pattern framework, potentially triggering unnecessary retrace/recompile work for models with many adaptive RMSNorms.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
This PR fixes a Vulkan RMSNorm fusion correctness issue where the fusion could fold a non-prepackable multiplier into et_vk.rms_norm, causing the Vulkan runtime to abort during weight prepacking. The change tightens the fusion criteria so the folded-in multiplier is only accepted when it can be resolved as a prepackable constant/parameter in the exported program, and adds a regression test covering the previously-crashing shapes.
Changes:
- Add a guard in the RMSNorm pattern detector to reject non-leaf (graph-computed) multipliers.
- Add a guard in the RMSNorm replacement to reject placeholder multipliers that are actually user inputs (not params/buffers/lifted constants).
- Add a unit test ensuring fusion only occurs for prepackable RMSNorm weights.
File summaries
| File | Description |
|---|---|
| backends/vulkan/patterns/rms_norm.py | Tightens RMSNorm fusion eligibility to avoid folding non-prepackable multipliers; adds a final is_param_node check during replacement. |
| backends/vulkan/test/test_vulkan_passes.py | Adds coverage for RMSNorm fusion only occurring when the folded multiplier is prepackable (and not an intermediate node or graph input). |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # The detector only sees the graph, which cannot distinguish a constant | ||
| # placeholder from a user input; both look the same there. Only a constant | ||
| # is actually prepackable, so make the final check here, where the exported | ||
| # program is available. | ||
| if not utils.is_param_node(ep, match.weight_node): | ||
| return |
Fixes #22773
Summary
et_vk.rms_normprepacks its weight, so the multiply that the fusion folds in has to be a constant the prepacker can see.RmsNormMatchfolded in whatever multiply followed the norm, without checking, so a multiplier computed in the graph produced an op the runtime aborts on at the first inference:Two common shapes hit this:
norm(x) * (1.0 + weight). The multiplier is constant valued, but it is an intermediate node, not a leaf, so it is equally unprepackable.So the deciding property is not "constant vs. computed" — it is whether the multiplier is something the prepacker can resolve to a constant tensor.
The change
Reject a non-prepackable weight. This is done in two places because the two callers see different things:
RmsNormMatch.__init__) is handed only a node —get_all_fusable_subgraphshas noExportedProgram— so it requires the weight to be a placeholder. That covers every computed multiplier, including both cases above.utils.is_param_node.Refusing the match is safe: the norm and the multiply are ordinary supported ops, so the pattern stays unfused and costs one extra dispatch. Note the weight is not part of
match.all_nodes, so declining the match does not leave the partitioner holding nodes nothing will replace.Test plan
12 tests pass. The new test,
test_rms_norm_fuses_only_prepackable_weight, fails without the source change, in both of its negative subtests:Also verified end to end on an AMD Radeon 8060S (RADV GFX1151, RDNA 3.5), lowering each module and comparing the delegate's output against eager:
1.0 + parameter(Gemma)The dtype-cast row is worth calling out: the pattern's own docstring describes the Llama-style graph where the weight reaches the multiply through a
to_copy, and that form aborts today whenever the cast is not constant-folded before the pass. Requiring the weight node itself to be a placeholder covers it, at the cost of not fusing that form — which is what already happens, except now it runs instead of aborting. If fusing through a weight cast is wanted, that seems better as a separate change that prepacks the underlying constant deliberately.backends.vulkan.test.test_vulkan_delegateproduces an identical set of results before and after (its failures in my environment are pre-existing and unrelated).Found while lowering openpi's π₀.₅ to the Vulkan delegate.