[ExecuTorch][WebGPU] Add q4gsw training kernels (backward, weight-grad, requant) to the WebGPU backend#20931
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.
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
…d, requant) to the WebGPU backend Pull Request resolved: #20931 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), and `et_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 kernels - `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`) - `op_registry.py` — 3 `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Training-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`/`Requant` ops (scales-rank + `K==0`). Co-authored-with: Claude Code. ghstack-source-id: 405026076 @exported-using-ghexport Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)
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