diff --git a/backends/webgpu/test/ops/test_linear_q4gsw_dw.py b/backends/webgpu/test/ops/test_linear_q4gsw_dw.py new file mode 100644 index 00000000000..7570bfe42cf --- /dev/null +++ b/backends/webgpu/test/ops/test_linear_q4gsw_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_q4gsw_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 Q4gswDwModule(torch.nn.Module): + def forward(self, d_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + return torch.ops.et_vk.linear_q4gsw_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(Q4gswDwModule().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 TestLinearQ4gswDw(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_q4gsw_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_q4gsw_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 new file mode 100644 index 00000000000..a342902ebea --- /dev/null +++ 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 new file mode 100644 index 00000000000..a675be79251 --- /dev/null +++ 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(" 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()