[ExecuTorch][WebGPU] Tests for the q4gsw_train WebGPU ops

JCNTH · JCNTH · commit 0545abb80efc · 2026-07-21T09:50:45.000-07:00
Pull Request resolved: #20932 Export-delegation + fp64 golden tests for the q4gsw training (`backward`, `dw`, `requant`) ops. Key changes: - `test/ops/{test_quantized_linear_backward,test_linear_q4gsw_dw,test_q4gsw_requant}.py` — fp64 dequant-matmul / weight-grad / re-quant goldens + a round-trip + delegation checks — fp64 library-reference goldens + delegation checks (≥2 cases/op). Co-authored-with: Claude Code. ghstack-source-id: 405026079 @exported-using-ghexport Differential Revision: [D111755121](https://our.internmc.facebook.com/intern/diff/D111755121/)
diff --git a/backends/webgpu/test/ops/test_linear_dW.py b/backends/webgpu/test/ops/test_linear_dW.py
@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`et_vk.linear_dW` (STE weight gradient) export + fp64 golden.
+
+The weight gradient of a frozen 4-bit linear: `d_W[N, K] = d_out^T @ x`, the grad
+wrt the dequantized weight (both operands are fp32; no int4 unpack). Reached by a
+direct op call in the on-device training graph. CONFIGS reuse Llama-3.2-1B linear
+shapes plus a non-tile-aligned shape that exercises the kernel's boundary clamp.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (m tokens, k in_features, n out_features).
+CONFIGS = {
+    "q_proj_112": (112, 2048, 2048),
+    "kv_proj_112": (112, 2048, 512),
+    "boundary": (13, 18, 10),  # non-multiple-of-4: exercises the min()-clamp
+}
+
+
+class LinearDwModule(torch.nn.Module):
+    def forward(self, d_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        return torch.ops.et_vk.linear_dW(d_out, x)
+
+
+def _det_inputs(m: int, k: int, n: int):
+    """Deterministic fp32 d_out [m, n] + x [m, k] (fixed seed)."""
+    g = torch.Generator().manual_seed(0)
+    d_out = torch.randn(m, n, generator=g, dtype=torch.float32)
+    x = torch.randn(m, k, generator=g, dtype=torch.float32)
+    return d_out, x
+
+
+def _fp64_golden(d_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    """fp64 truth: d_W = d_out^T @ x, [N, M] @ [M, K] = [N, K]."""
+    return (d_out.double().t() @ x.double()).to(torch.float32)
+
+
+def _export(d_out: torch.Tensor, x: torch.Tensor):
+    ep = torch.export.export(LinearDwModule().eval(), (d_out, x))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestLinearDw(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (m, k, n) in CONFIGS.items():
+            with self.subTest(config=name):
+                d_out, x = _det_inputs(m, k, n)
+                et = _export(d_out, x)
+                self.assertTrue(
+                    _delegated(et),
+                    f"Expected a VulkanBackend delegate (linear_dW {name})",
+                )
+
+    def test_op_matches_fp64_golden(self) -> None:
+        # Op (d_out^T @ x) vs fp64 matmul truth: guards the formula.
+        for name, (m, k, n) in CONFIGS.items():
+            with self.subTest(config=name):
+                d_out, x = _det_inputs(m, k, n)
+                got = torch.ops.et_vk.linear_dW(d_out, x)
+                golden = _fp64_golden(d_out, x)
+                self.assertEqual(tuple(got.shape), (n, k))
+                torch.testing.assert_close(got, golden, atol=5e-4, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_q4gsw_requant.py b/backends/webgpu/test/ops/test_q4gsw_requant.py
@@ -0,0 +1,123 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`et_vk.q4gsw_requant` (STE re-quant + int4 pack) export + fp32 code golden.
+
+Writes updated fp32 latent weights back to the 4-bit group-symmetric packed codes
+`et_vk.linear_q4gsw` reads (only the codes move; the per-group scale is frozen).
+Reached by a direct op call in the on-device training graph after the optimizer
+step. The golden is computed in fp32 (not fp64) on purpose: the kernel rounds
+`round(latent / scale)` in IEEE-754 fp32, so a bit-exact contract must round in
+fp32 too -- an fp64 reference would flip codes at half-way ties. CONFIGS reuse
+Llama-3.2-1B linear shapes plus an odd-K shape that exercises the final-nibble
+tail guard.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (n out_features, k in_features, group_size).
+CONFIGS = {
+    "kv_proj": (512, 2048, 64),
+    "q_proj_g32": (2048, 2048, 32),
+    "small_odd_k": (6, 129, 64),  # odd K -> trailing low-nibble-only byte
+}
+
+
+class Q4gswRequantModule(torch.nn.Module):
+    def __init__(self, group_size: int) -> None:
+        super().__init__()
+        self.group_size = group_size
+
+    def forward(self, latent: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+        return torch.ops.et_vk.q4gsw_requant(latent, scales, self.group_size)
+
+
+def _det_inputs(n: int, k: int, gs: int):
+    """Deterministic fp32 latent [N, K] + frozen scales [num_groups, N] (fixed seed).
+
+    Scales are small relative to the latent so `latent / scale` spans well beyond
+    [-8, 7], exercising the clamp on both ends.
+    """
+    num_groups = (k + gs - 1) // gs
+    g = torch.Generator().manual_seed(0)
+    latent = torch.randn(n, k, generator=g, dtype=torch.float32)
+    scales = torch.rand(num_groups, n, generator=g, dtype=torch.float32) * 0.1 + 0.05
+    return latent, scales
+
+
+def _reference_codes(
+    latent: torch.Tensor, scales: torch.Tensor, gs: int
+) -> torch.Tensor:
+    """fp32 truth for the int4 codes: clamp(round(latent / scale), -8, 7), [N, K]."""
+    n, k = latent.shape
+    group_idx = torch.arange(k) // gs  # [K]
+    scale_full = scales.t()[:, group_idx]  # [N, K]: scales[k // gs, n]
+    q = torch.round(latent / scale_full)
+    return torch.clamp(q, -8, 7).to(torch.int64)
+
+
+def _unpack(packed: torch.Tensor, k: int) -> torch.Tensor:
+    """Undo the nibble packing: even k -> low nibble, odd k -> high nibble, code - 8."""
+    n = packed.shape[0]
+    p = packed.to(torch.int64)
+    low = p & 0xF
+    high = (p >> 4) & 0xF
+    codes = torch.zeros(n, k, dtype=torch.int64)
+    n_low = codes[:, 0::2].shape[1]
+    n_high = codes[:, 1::2].shape[1]
+    codes[:, 0::2] = low[:, :n_low] - 8
+    codes[:, 1::2] = high[:, :n_high] - 8
+    return codes
+
+
+def _export(latent: torch.Tensor, scales: torch.Tensor, gs: int):
+    ep = torch.export.export(Q4gswRequantModule(gs).eval(), (latent, scales))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestQ4gswRequant(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (n, k, gs) in CONFIGS.items():
+            with self.subTest(config=name):
+                latent, scales = _det_inputs(n, k, gs)
+                et = _export(latent, scales, gs)
+                self.assertTrue(
+                    _delegated(et),
+                    f"Expected a VulkanBackend delegate (q4gsw_requant {name})",
+                )
+
+    def test_op_matches_fp32_golden(self) -> None:
+        # Op codes vs fp32 quant truth: guards formula+layout, bit-exact.
+        for name, (n, k, gs) in CONFIGS.items():
+            with self.subTest(config=name):
+                latent, scales = _det_inputs(n, k, gs)
+                packed = torch.ops.et_vk.q4gsw_requant(latent, scales, gs)
+                self.assertEqual(packed.dtype, torch.uint8)
+                self.assertEqual(tuple(packed.shape), (n, (k + 1) // 2))
+                got = _unpack(packed, k)
+                golden = _reference_codes(latent, scales, gs)
+                torch.testing.assert_close(got, golden, atol=0, rtol=0)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_quantized_linear_backward.py b/backends/webgpu/test/ops/test_quantized_linear_backward.py
@@ -0,0 +1,157 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Backward of 4-bit quantized linear (`et_vk.linear_q4gsw_backward`) export + fp64 golden.
+
+Mirrors test_quantized_linear.py. The backward computes `d_x = d_out @ dequant(W)` so
+gradients flow through a frozen 4-bit base into a LoRA/DiReFT adapter. CONFIGS reuse the
+real Llama-3.2-1B linear shapes (the backward's d_out is [M, N] and its output d_x is
+[M, K]). The golden is the fp64 dequant-matmul truth; the native test
+(test_webgpu_native.cpp) reconstructs the identical deterministic ramp bit-for-bit.
+"""
+
+import os
+import unittest
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from executorch.backends.vulkan import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+from torchao.quantization.granularity import PerGroup
+from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
+
+
+@dataclass(frozen=True)
+class BwdConfig:
+    name: str
+    m: int  # rows (tokens)
+    k: int  # in_features (== d_x cols)
+    n: int  # out_features (== d_out cols)
+    group_size: int = 32  # K % group_size == 0, K % 8 == 0, N % 8 == 0
+    heavy: bool = False
+
+
+# Mirrored by the C++ kQ4gswBackwardConfigs table (Llama-3.2-1B shapes).
+CONFIGS = [
+    BwdConfig("q_proj", 1, 2048, 2048),  # also o_proj
+    BwdConfig("kv_proj", 1, 2048, 512),
+    BwdConfig("gate_proj", 1, 2048, 8192),
+    BwdConfig("down_proj", 1, 8192, 2048),
+    BwdConfig("q_proj_112", 112, 2048, 2048),  # S=112 multi-row training window
+]
+
+
+def _make_quantized_model(k: int, n: int, group_size: int) -> torch.nn.Module:
+    torch.manual_seed(0)  # load-bearing: fixes the weights the golden uses
+    m = torch.nn.Linear(k, n, bias=False).eval()
+    quantize_(
+        m,
+        IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(group_size)),
+    )
+    return m
+
+
+def _ramp(m_rows: int, cols: int) -> torch.Tensor:
+    """Deterministic fp32 [rows, cols]; the C++ side reconstructs it bit-for-bit.
+
+    v[flat] = ((flat % 17) - 8) / 16 -- exact in fp32 (small modulus, po2 denominator).
+    """
+    flat = np.arange(m_rows * cols, dtype=np.int64)
+    v = ((flat % 17) - 8).astype(np.float32) / np.float32(16.0)
+    return torch.from_numpy(v).reshape(m_rows, cols)
+
+
+def _packed_qweights(m: torch.nn.Module):
+    """The int4 packed weights + per-group scales `et_vk.linear_q4gsw` consumes.
+
+    Recover them the same way the forward op does: `dequant(W)` is [N, K]; the backward op
+    takes the same (weights, weight_scales, group_size) triple the partitioner extracts.
+    """
+    aqt = m.weight  # AffineQuantizedTensor
+    return aqt
+
+
+def _fp64_golden(m: torch.nn.Module, d_out: torch.Tensor) -> np.ndarray:
+    """fp64 truth: d_x = d_out @ dequant(W), dequant(W) is [N, K] -> [M, N]@[N, K] = [M, K]."""
+    wq = m.weight.dequantize()  # [N, K]
+    d_x = d_out.double() @ wq.double()  # [M, K] in fp64
+    return d_x.to(torch.float32).numpy().astype("<f4")
+
+
+class _BackwardModule(torch.nn.Module):
+    """Wraps the linear so autograd through it exercises the registered backward op."""
+
+    def __init__(self, lin: torch.nn.Module) -> None:
+        super().__init__()
+        self.lin = lin
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lin(x)
+
+
+def _export_backward(m: torch.nn.Module, x: torch.Tensor):
+    # Training-style export: forward + backward makes the backward reachable.
+    mod = _BackwardModule(m)
+    ep = torch.export.export(mod, (x,))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+class TestQuantizedLinearBackward(unittest.TestCase):
+    def test_op_matches_fp64_golden(self) -> None:
+        # Op impl (d_out @ dequant(W)) vs fp64 truth: guards backward formula.
+        for cfg in CONFIGS:
+            if cfg.heavy:
+                continue
+            with self.subTest(config=cfg.name):
+                m = _make_quantized_model(cfg.k, cfg.n, cfg.group_size)
+                d_out = _ramp(cfg.m, cfg.n)
+                got = torch.ops.et_vk.linear_q4gsw_backward(
+                    d_out, _packed_qweights(m), None, cfg.group_size
+                )
+                golden = torch.from_numpy(_fp64_golden(m, d_out))
+                torch.testing.assert_close(got, golden, atol=5e-4, rtol=1e-3)
+
+    def test_autograd_backward_matches_golden(self) -> None:
+        # autograd through linear_q4gsw uses the registered backward op.
+        for cfg in CONFIGS:
+            if cfg.heavy:
+                continue
+            with self.subTest(config=cfg.name):
+                m = _make_quantized_model(cfg.k, cfg.n, cfg.group_size)
+                x = _ramp(cfg.m, cfg.k).requires_grad_(True)
+                d_out = _ramp(cfg.m, cfg.n)
+                y = m(x)
+                y.backward(d_out)
+                golden = torch.from_numpy(_fp64_golden(m, d_out))
+                torch.testing.assert_close(x.grad, golden, atol=5e-4, rtol=1e-3)
+
+
+def export_backward_model(cfg: BwdConfig, pte_path: str, golden_path: str) -> None:
+    """Export one config's backward .pte + its fp64 golden (raw LE fp32)."""
+    m = _make_quantized_model(cfg.k, cfg.n, cfg.group_size)
+    x = _ramp(cfg.m, cfg.k)
+    et = _export_backward(m, x)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    _fp64_golden(m, _ramp(cfg.m, cfg.n)).tofile(golden_path)
+    print(f"Exported {pte_path}; golden {golden_path} ({cfg.m * cfg.k} floats)")
+
+
+def export_all_backward_models(out_dir: str, include_heavy: bool = False) -> None:
+    for cfg in CONFIGS:
+        if cfg.heavy and not include_heavy:
+            continue
+        pte = os.path.join(out_dir, f"q4gsw_backward_{cfg.name}.pte")
+        golden = os.path.join(out_dir, f"q4gsw_backward_{cfg.name}.golden.bin")
+        export_backward_model(cfg, pte, golden)
+
+
+if __name__ == "__main__":
+    unittest.main()
