[ExecuTorch][Vulkan] Add et_vk.linear_q4gsw_backward kernel (4-bit input-grad)#20943
[ExecuTorch][Vulkan] Add et_vk.linear_q4gsw_backward kernel (4-bit input-grad)#20943JCNTH wants to merge 11 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20943
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 1 PendingAs of commit 031f2a8 with merge base 21554e5 ( 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):
Adds the Vulkan
et_vk.linear_q4gsw_backwardkernel — the input-gradient of the frozen 4-bitlinear_q4gswbase, for on-device adapter training. Computesd_x[M, K] = d_out[M, N] @ dequant(W)[N, K], contracting over N.Problem:
et_vk.linear_q4gsw_backwardis registered in the shared Vulkan partitioner (custom_ops_lib.py+op_registry.py) but Vulkan had no runtime kernel, so the op could not run.Solution: a 4M x 4K register-tiled GLSL kernel that reads the SAME W_4X8 block-packed weight the forward reads, re-deriving the nibble/scale addressing so the training loop stays consistent with the forward with no re-pack.
dequant(W[n, k]) = (code - 8) * scale.Key changes:
glsl/q4gsw_backward.{glsl,yaml}— buffer x float; W_4X8 nibble unpack mirroringglsl/q4gsw_linear_gemm__w_4x8.glsl(even-N low nibble, odd-N high,N4_paddedivec4 stride,[num_groups, N]scales).impl/QuantizedLinearBackward.cpp— reusesprepack_q4_w_4x8_nc_buffer+prepack_q4_scalesfrom the forward; 1D tile dispatchceil(M/4) * ceil(K/4)with a workgroup-count guard;group_sizespecialization constant.Constraints: buffer storage, fp32;
N % 4 == 0,K % 4 == 0,group_size % 4 == 0(matches the forward prepack). Weight/scale layout identical to the forward.@exported-using-ghexport
Differential Revision: D111797529
Differential Revision: D111797529