Repository navigation
[XNNPACK] Delegate GEMMs whose weight is computed from parameters - #24352
JakeStevens wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24352
Note: Links to docs will display an error until the docs builds have been completed. ❌ 5 New FailuresAs of commit 1a735c5 with merge base 455aa9c ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
96e5183 to
3008de5
Compare
This PR needs a
|
3008de5 to
dc58a2f
Compare
A convolution, linear, addmm or mm whose weight or bias is computed from parameters, such as a convolution under weight_norm or spectral_norm, was declined by the XNNPACK partitioner and ran on the portable kernels together with the computation of its weight. The GEMM configs now accept such a weight or bias when it is computed with pure ATen ops from parameters, unmutated buffers or constants, the computation feeds only that GEMM, and the folded tensor is no larger than the parameters it replaces. The computation is partitioned with the GEMM, and FoldComputedParamsPass folds it into a constant at the start of XNNPACK preprocess, inside the delegate, so the top-level graph is not modified. Fixes pytorch#22078 This change was authored with Claude Code.
dc58a2f to
1a735c5
Compare
|
@claude review effort:max |
|
Claude finished @JakeStevens's task in 8m 21s —— View job This PR is functionally sound and no blocker turned up, but the partitioner and the fold pass communicate through node-name tags that get written during constraint checks and are never cleared, and the main weight_norm test never checks that a fold happened. Detailed reviewPR Review: #24352SummaryThe PR lets XNNPACK GEMM configs accept a weight or bias computed from params with pure ATen ops: the computation is partitioned with the GEMM, then Code Quality
Testing
Performance
RecommendationNeeds Discussion The fold is correct and the benchmark result is compelling. Before merging, the partitioner→pass contract should be checked locally (or tags written only on accept) rather than inferred from graph-wide name matching. The weight_norm test should also assert the fold. |
digantdesai
left a comment
digantdesai
left a comment
There was a problem hiding this comment.
I will take 1000X improvement for these gemms, thank you! :)
|
@john-rocky and @mansnils can you take a look and see how well this covers your respective PRs? both were starting to get a little unwieldily and devolving with multiple rounds of review. this is not necessarily a 1:1 complete replacement, but i hope it can get most of the benefits you were hoping for with a more limited blast radius |
|
Thanks, @JakeStevens; folding inside the delegate is cleaner than mine. I lowered #22391's test graphs plus a few on main and both branches, without running them. Same delegation as #22391: weight_norm on convolutions and Linear, #22391 also delegates GEMMs whose weight comes out of split, chunk, unbind or max, and An Embedding shared by a lookup method and an I'll close #22391 once this lands. If useful, I'm happy to help with a small follow-up for the packed split |
|
@john-rocky thanks, reproduced: 526,976 vs 264,720 bytes. Main stores
On main, Thanks also for checking #22391's graphs. A follow-up for the packed |
Summary
Fixes #22078
The XNNPACK partitioner requires a GEMM's weight and bias to be static params. A convolution, linear, addmm or mm whose weight is computed from parameters is declined, and it runs on the portable kernels together with the computation of its weight.
torch.nn.utils.parametrizations.weight_normandspectral_normproduce exactly this graph, and they are common in audio models; the issue measured 941x on wav2vec2's positional convolution.This is a smaller alternative to #22391 and #23154, which fold parameter-only subgraphs in
transform_for_pre_decomposition. Folding there rewrites the shared ATen program, before functionalization and before every other partitioner and method sees it. Here the top-level graph is not modified.get_static_param_depslets a GEMM config accept a weight or bias computed with pure ATen ops from parameters, unmutated buffers or constants. It partitions that computation together with the GEMM and tags it, andFoldComputedParamsPass, the first XNNPACK preprocess pass, folds it into a constant inside the delegate. After the fold, the delegate looks exactly as if the weight had been a parameter, and the named data store's content hashes deduplicate identical folds across methods.The partitioner only accepts a computation that the fold can replace outright without growing the program. Only the folded tensor and its source parameters may be used outside the computation, and only by that GEMM. The folded tensor may be no larger than the parameters it frees, and sources tagged for an external data file are left alone. A computation shared by several GEMMs is therefore not folded, such as the packed bias split across three linears in #23154; delegating it would require merging those GEMMs into one partition, which can create a cycle. Tags are set while a GEMM is being checked, before it is accepted, and a GEMM rejected by an earlier partitioner leaves them on nodes that later become delegate inputs, so the pass folds a tagged node only when one of the GEMMs that tagged it is in the same delegate and the node's inputs are params or already folded. As with a buffer used directly as a weight today, a buffer that only another method writes is treated as constant. Exclusivity is also checked per method. If another method uses a source without delegating it, as when an embedding lookup shares its weight with
x @ W.t(), the program stores the source as a constant and the delegate stores the folded tensor.F.linear(x, emb.weight)is already stored twice like this on main.Suggested review order:
get_static_param_depsand its two call sites ingemm_configs.py, thenfold_computed_params.py.wav2vec2
pos_conv,weight_norm(Conv1d(1024, 1024, 128, padding=64, groups=16), dim=2)on its own at sequence length 499, synthetic weights, macOS arm64, median of 5 runs after 2 warm-up runs:pow,sum,convolutionThe
.pteis 1,280 bytes smaller: one folded weight replacesweight_gandweight_v, and the program has one delegate instead of two.Test plan
python -m pytest backends/xnnpack/test/passes/test_fold_computed_params.pyThe new tests lower a weight_norm Conv1d and a linear with a computed weight and bias, run them through the runtime against eager, and check that nothing is left outside the delegate and, for the linears, that the delegate's constants are exactly the folded tensors. A linear with a computed bias is also covered under static and dynamic quantization. They also check what must not fold: a source used elsewhere, a mutated buffer, a low-rank product that would grow the program, a source that stays another input of the GEMM, an intermediate that also feeds the GEMM, sources tagged for an external file, and a weight computed by a custom op. A rejected convolution whose weight is computed by a delegated matmul, a weight computed by an earlier partitioner's delegate, and an addmm that one partitioner tags and rejects and a later one takes with a dynamic weight still export and match eager. A param that already has the folded weight's name is not replaced, and a scaled addmm with a computed weight stays on portable and matches eager.
python -m pytest backends/xnnpack/test/passes backends/xnnpack/test/ops backends/xnnpack/test/serialization backends/xnnpack/test/quantizer backends/xnnpack/test/test_xnnpack_partitioner.py -o addopts=""This includes
test_linear.py,test_bmm.pyand the quantizer tests, whichpytest.iniskips. The failures are the same as on main: three qd8 int4 linear tests, and quantizer tests that fail depending on run order. I also compared against main for weight_norm on Linear, Conv2d, depthwise and transposed convolutions, spectral_norm, a dynamic sequence length, and two methods sharing a folded weight. Tied embeddings within one method, packed biases and low-rank weights produce the same.pteas on main.This PR was authored with Claude Code.