|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""4-bit groupwise-symmetric quantized embedding (`et_vk.embedding_q4gsw`) export |
| 8 | ++ golden for the WebGPU backend. |
| 9 | +
|
| 10 | +Quantizes an nn.Embedding with the Llama EmbeddingQuantHandler recipe (int4 |
| 11 | +groupwise-symmetric, packed) which lowers to `quantized_decomposed.embedding_4bit` |
| 12 | +and fuses under VulkanPartitioner into `et_vk.embedding_q4gsw.default` |
| 13 | +(is_linear_weight=False). Writes a torch-computed golden (the native binary has no |
| 14 | +ATen) via the registered et_vk reference op + the raw int32 indices the native |
| 15 | +test loads and compares. |
| 16 | +
|
| 17 | +Two shapes are exercised: a tiny one and a Llama-3.2-1B-scale one (EMBED=2048, |
| 18 | +GROUP=64) so the per-group scale indexing (32 groups/row) + dequant are validated |
| 19 | +at the real embedding dim, not just a single 64-wide row. |
| 20 | +""" |
| 21 | + |
| 22 | +import unittest |
| 23 | +from collections import namedtuple |
| 24 | + |
| 25 | +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 |
| 26 | + |
| 27 | +import torch |
| 28 | +from executorch.backends.vulkan import VulkanPartitioner |
| 29 | +from executorch.examples.models.llama.source_transformation.quantize import ( |
| 30 | + EmbeddingQuantHandler, |
| 31 | +) |
| 32 | +from executorch.exir import to_edge_transform_and_lower |
| 33 | + |
| 34 | +# vocab rows, embed columns (embed % 32 == 0), group-wise scales, gather indices. |
| 35 | +Shape = namedtuple("Shape", ["name", "vocab", "embed", "group", "indices"]) |
| 36 | +SHAPES = [ |
| 37 | + Shape("small", 64, 64, 32, [1, 5, 63, 0]), |
| 38 | + # Llama-3.2-1B embedding dim + group (small vocab keeps the export light). |
| 39 | + Shape("llama1b", 512, 2048, 64, [1, 5, 511, 0]), |
| 40 | +] |
| 41 | + |
| 42 | + |
| 43 | +class _EmbeddingModel(torch.nn.Module): |
| 44 | + def __init__(self, vocab: int, embed: int) -> None: |
| 45 | + super().__init__() |
| 46 | + self.emb = torch.nn.Embedding(vocab, embed) |
| 47 | + |
| 48 | + def forward(self, idx: torch.Tensor) -> torch.Tensor: |
| 49 | + return self.emb(idx) |
| 50 | + |
| 51 | + |
| 52 | +def _make_quantized_model(shape: Shape) -> torch.nn.Module: |
| 53 | + torch.manual_seed(0) |
| 54 | + return ( |
| 55 | + EmbeddingQuantHandler( |
| 56 | + _EmbeddingModel(shape.vocab, shape.embed).eval(), |
| 57 | + device="cpu", |
| 58 | + bitwidth=4, |
| 59 | + group_size=shape.group, |
| 60 | + packed=True, |
| 61 | + quantize_with_hqq=False, |
| 62 | + ) |
| 63 | + .quantized_model() |
| 64 | + .eval() |
| 65 | + ) |
| 66 | + |
| 67 | + |
| 68 | +def _indices(shape: Shape) -> torch.Tensor: |
| 69 | + return torch.tensor(shape.indices, dtype=torch.long) |
| 70 | + |
| 71 | + |
| 72 | +def _quant_params(qm: torch.nn.Module) -> tuple[torch.Tensor, torch.Tensor, int]: |
| 73 | + sd = qm.state_dict() |
| 74 | + weight = next( |
| 75 | + v for k, v in sd.items() if k.endswith("weight") and v.dtype == torch.uint8 |
| 76 | + ) |
| 77 | + scales = next(v for k, v in sd.items() if k.endswith("scales")) |
| 78 | + if scales.ndim == 1: |
| 79 | + scales = scales.unsqueeze(1) |
| 80 | + embed = weight.shape[1] * 2 |
| 81 | + group_size = embed // scales.shape[1] |
| 82 | + return weight, scales, group_size |
| 83 | + |
| 84 | + |
| 85 | +def _golden(qm: torch.nn.Module, idx: torch.Tensor) -> torch.Tensor: |
| 86 | + # Reference = the registered et_vk dequant+gather op (non-linear branch). |
| 87 | + weight, scales, group_size = _quant_params(qm) |
| 88 | + return torch.ops.et_vk.embedding_q4gsw.default( |
| 89 | + weight, scales, group_size, idx, False |
| 90 | + ) |
| 91 | + |
| 92 | + |
| 93 | +def _export(qm: torch.nn.Module, idx: torch.Tensor): |
| 94 | + ep = torch.export.export(qm, (idx,)) |
| 95 | + return to_edge_transform_and_lower( |
| 96 | + ep, partitioner=[VulkanPartitioner()] |
| 97 | + ).to_executorch() |
| 98 | + |
| 99 | + |
| 100 | +class TestEmbeddingQ4gsw(unittest.TestCase): |
| 101 | + def test_export_delegates(self) -> None: |
| 102 | + for shape in SHAPES: |
| 103 | + with self.subTest(shape=shape.name): |
| 104 | + et = _export(_make_quantized_model(shape), _indices(shape)) |
| 105 | + found = any( |
| 106 | + d.id == "VulkanBackend" |
| 107 | + for plan in et.executorch_program.execution_plan |
| 108 | + for d in plan.delegates |
| 109 | + ) |
| 110 | + self.assertTrue( |
| 111 | + found, "Expected a VulkanBackend delegate (embedding_q4gsw fusion)" |
| 112 | + ) |
| 113 | + |
| 114 | + def test_golden_matches_eager(self) -> None: |
| 115 | + # The torch golden (et_vk reference) must equal torch dequant+gather, so a |
| 116 | + # buggy golden can't fake-pass the native kernel. Run at both shapes so the |
| 117 | + # Llama-scale per-group scale indexing (32 groups/row) is covered. |
| 118 | + for shape in SHAPES: |
| 119 | + with self.subTest(shape=shape.name): |
| 120 | + qm = _make_quantized_model(shape) |
| 121 | + idx = _indices(shape) |
| 122 | + weight, scales, group_size = _quant_params(qm) |
| 123 | + vocab = weight.shape[0] |
| 124 | + embed = weight.shape[1] * 2 |
| 125 | + # fp64 reference dequant, vectorized (no fp32 rounding in oracle). |
| 126 | + w = weight.to(torch.int64) |
| 127 | + nib = torch.empty((vocab, embed), dtype=torch.int64) |
| 128 | + nib[:, 0::2] = (w >> 4) & 0xF # even dim -> high nibble |
| 129 | + nib[:, 1::2] = w & 0xF # odd dim -> low nibble |
| 130 | + scale_exp = scales.to(torch.float64).repeat_interleave( |
| 131 | + group_size, dim=1 |
| 132 | + ) |
| 133 | + deq = (nib - 8).to(torch.float64) * scale_exp |
| 134 | + eager = torch.nn.functional.embedding(idx, deq) |
| 135 | + golden = _golden(qm, idx) |
| 136 | + torch.testing.assert_close(golden.double(), eager, atol=1e-5, rtol=1e-5) |
| 137 | + |
| 138 | + |
| 139 | +def export_embedding_q4gsw_model( |
| 140 | + pte_path: str, golden_path: str, indices_path: str, shape_name: str = "small" |
| 141 | +) -> None: |
| 142 | + """Write the embedding_q4gsw .pte + torch golden (raw LE fp32) + raw LE int32 |
| 143 | + indices (downcast from int64 for the int32-typed GPU buffer). `shape_name` |
| 144 | + selects an entry from SHAPES (default the tiny shape; "llama1b" = EMBED=2048).""" |
| 145 | + shape = next(s for s in SHAPES if s.name == shape_name) |
| 146 | + qm = _make_quantized_model(shape) |
| 147 | + idx = _indices(shape) |
| 148 | + golden = _golden(qm, idx).detach().numpy().astype("<f4") |
| 149 | + et = _export(qm, idx) |
| 150 | + with open(pte_path, "wb") as f: |
| 151 | + f.write(et.buffer) |
| 152 | + golden.tofile(golden_path) |
| 153 | + idx.to(torch.int32).numpy().astype("<i4").tofile(indices_path) |
| 154 | + print( |
| 155 | + f"Exported {pte_path} (shape={shape_name}); golden {golden_path} " |
| 156 | + f"({golden.size} floats); indices {indices_path} ({idx.numel()} int32)" |
| 157 | + ) |
| 158 | + |
| 159 | + |
| 160 | +if __name__ == "__main__": |
| 161 | + unittest.main() |
0 commit comments