[ExecuTorch][WebGPU] Add q4gsw training kernels (backward, weight-grad, requant) to the WebGPU backend#20931
[ExecuTorch][WebGPU] Add q4gsw training kernels (backward, weight-grad, requant) to the WebGPU backend#20931JCNTH wants to merge 8 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20931
Note: Links to docs will display an error until the docs builds have been completed. ❌ You can merge normally! (1 Unrelated Failure), 6 Unclassified FailuresAs of commit f4b2b88 with merge base 21554e5 ( UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
SS-JIA
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
Stack from ghstack (oldest at bottom):
Add the 4-bit fine-tuning weight-path custom ops —
et_vk.linear_q4gsw_backward(input grad dX),et_vk.linear_dW(weight grad dW — a generic fp32 weight-gradient matmul, not q4gsw-specific), andet_vk.q4gsw_requant(STE re-quant of fp32 latents → 4-bit codes) — that close the on-device training loop through a frozen 4-bit base.Key changes:
runtime/ops/quantized_linear/{QuantizedLinearBackward,QuantizedLinearRequant}.cpp+LinearDw.cpp+q4gsw_{backward,requant}.wgsl+linear_dW.wgsl— the three training kernelscustom_ops_lib.py— 3 custom-op defs (linear_q4gsw_backward+register_autogradonlinear_q4gsw,linear_dW,q4gsw_requant)op_registry.py— 3OpFeaturesCMakeLists.txtWEBGPU_SRCS— wire the sourcesTraining-only, WebGPU-only custom ops registered under the shared Vulkan partitioner; Vulkan has no kernels for them yet (a separate Vulkan-kernel track covers that). Guards mirror the sibling
Dw/Requantops (scales-rank +K==0).Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D111755130
Differential Revision: D111755130