From d18198dc1a2bd0922a69066f106e10b034785cb1 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 20 Jul 2026 17:42:06 -0700 Subject: [PATCH] Refactor quantization into a shared extension/llm/export/quant package (#21044) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/21044 Consolidates the quantization/packing logic that was previously duplicated across per-model directories into a single reusable package under extension/llm/export/quant, and introduces a model-free checkpoint loader. Key changes: - New extension/llm/export/quant package exposing a stable API: - convert.py: Convert, to_exportable/to_default, fuse_along_output, identity - quantize.py: quantize_model/quantize_weight/quantize_stream + dequantize - recipe.py: QuantConfig / QuantRecipe / QuantRule - New extension/llm/export/load.py: model-free (iter_checkpoint) and model-aware (load_checkpoint) readers that stream GGUF / safetensors weights, remap names to FQNs, and wrap quantized subclasses in portable Exportable* form. Backend-specific layout (CUDA coalesced int4, MLX gather buffers) is a separate terminal pass. - Extended int4.py / gguf.py to plug into the shared conversion path. - Removed the old per-model examples/models/gemma4_31b/quant/ packers (pack.py, pack_cuda.py, pack_mlx.py) and moved packing into examples/models/gemma4_31b/cuda_packers.py, wiring gemma4_31b and qwen3_5_moe onto the shared package. - Added unit tests for the new package and the gemma4_31b packers/pipelines. Reviewed By: Gasoonjia Differential Revision: D112867629 --- .github/workflows/cuda.yml | 2 +- .github/workflows/mlx.yml | 1 - backends/cuda/coalesced_int4_tensor.py | 44 +- .../shims/tests/gen_plain_mm_test_vectors.py | 14 +- backends/cuda/tests/test_int4_dispatch.py | 49 +- examples/models/gemma4_31b/cuda_packers.py | 160 ++++++ examples/models/gemma4_31b/export.py | 129 +++-- examples/models/gemma4_31b/gguf_loader.py | 202 +------ examples/models/gemma4_31b/inference.py | 2 +- examples/models/gemma4_31b/model.py | 75 +-- examples/models/gemma4_31b/quant/README.md | 54 -- examples/models/gemma4_31b/quant/__init__.py | 11 - examples/models/gemma4_31b/quant/pack.py | 106 ---- examples/models/gemma4_31b/quant/pack_cuda.py | 162 ------ examples/models/gemma4_31b/quant/pack_mlx.py | 103 ---- .../gemma4_31b/quant/tests/test_pack_mlx.py | 173 ------ .../models/gemma4_31b/quantize_and_save.py | 47 +- examples/models/gemma4_31b/tests/BUCK | 18 + .../gemma4_31b/tests/test_cuda_packers.py | 501 ++++++++++++++++++ .../gemma4_31b/tests/test_cuda_pipeline.py | 15 +- .../gemma4_31b/tests/test_mlx_pipeline.py | 46 +- .../models/gemma4_31b/tests/test_pipeline.py | 107 +++- .../models/gemma4_31b/tests/test_recipe.py | 61 +++ examples/models/qwen3_5_moe/export.py | 80 +-- .../models/qwen3_5_moe/quantize_and_save.py | 36 +- .../qwen3_5_moe/test_quantize_roundtrip.py | 11 +- extension/llm/export/BUCK | 64 +++ extension/llm/export/gguf.py | 71 ++- extension/llm/export/int4.py | 51 ++ extension/llm/export/load.py | 315 +++++++++++ extension/llm/export/quant/__init__.py | 20 + extension/llm/export/quant/convert.py | 216 ++++++++ extension/llm/export/quant/quantize.py | 386 ++++++++++++++ extension/llm/export/quant/recipe.py | 62 +++ extension/llm/export/quant/test/BUCK | 52 ++ .../llm/export/quant/test/test_convert.py | 362 +++++++++++++ .../llm/export/quant/test/test_quantize.py | 310 +++++++++++ .../llm/export/quant/test/test_recipe.py | 118 +++++ .../quant/test/test_safetensors_roundtrip.py | 143 +++++ extension/llm/export/test/BUCK | 21 + extension/llm/export/test/test_gguf.py | 54 +- extension/llm/export/test/test_int4.py | 14 + 42 files changed, 3346 insertions(+), 1122 deletions(-) create mode 100644 examples/models/gemma4_31b/cuda_packers.py delete mode 100644 examples/models/gemma4_31b/quant/README.md delete mode 100644 examples/models/gemma4_31b/quant/__init__.py delete mode 100644 examples/models/gemma4_31b/quant/pack.py delete mode 100644 examples/models/gemma4_31b/quant/pack_cuda.py delete mode 100644 examples/models/gemma4_31b/quant/pack_mlx.py delete mode 100644 examples/models/gemma4_31b/quant/tests/test_pack_mlx.py create mode 100644 examples/models/gemma4_31b/tests/BUCK create mode 100644 examples/models/gemma4_31b/tests/test_cuda_packers.py create mode 100644 examples/models/gemma4_31b/tests/test_recipe.py create mode 100644 extension/llm/export/load.py create mode 100644 extension/llm/export/quant/__init__.py create mode 100644 extension/llm/export/quant/convert.py create mode 100644 extension/llm/export/quant/quantize.py create mode 100644 extension/llm/export/quant/recipe.py create mode 100644 extension/llm/export/quant/test/BUCK create mode 100644 extension/llm/export/quant/test/test_convert.py create mode 100644 extension/llm/export/quant/test/test_quantize.py create mode 100644 extension/llm/export/quant/test/test_recipe.py create mode 100644 extension/llm/export/quant/test/test_safetensors_roundtrip.py diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 5178cbcbd0b..03f8556c461 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -211,7 +211,7 @@ jobs: # Run Gemma 4 31B tests (quant unit tests + pipeline integration tests) pip install gguf - python -m pytest examples/models/gemma4_31b/quant/tests/ examples/models/gemma4_31b/tests/ --ignore=examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v -o "addopts=" + python -m pytest examples/models/gemma4_31b/tests/ --ignore=examples/models/gemma4_31b/tests/test_mlx_pipeline.py -v -o "addopts=" unittest-cuda-runtime: name: unittest-cuda-runtime diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 40ebabbcf28..4ebabff91e1 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -85,7 +85,6 @@ jobs: backends/mlx/test/test_serialization_dedup.py \ backends/mlx/test/test_slot_recycling.py \ backends/mlx/test/test_sample.py \ - examples/models/gemma4_31b/quant/tests/test_pack_mlx.py \ examples/models/gemma4_31b/tests/test_mlx_pipeline.py \ -v echo "::endgroup::" diff --git a/backends/cuda/coalesced_int4_tensor.py b/backends/cuda/coalesced_int4_tensor.py index f2c1bbf51b0..09d5e2d066d 100644 --- a/backends/cuda/coalesced_int4_tensor.py +++ b/backends/cuda/coalesced_int4_tensor.py @@ -46,14 +46,13 @@ The coalesced [N, n_groups] layout is exactly what the W4A8 dp4a matvec kernel (``executorch_cuda::int4_plain_mm`` / ``int4_plain_mm.cuh``) reads row-for-row with qdata, so the exported decode graph carries no per-step transpose. The -transpose (and the uint8 re-encoding) is owned by :meth:`from_int4_tensor` so it -is baked into the serialized weight constant once at pack time. +is owned by :meth:`from_exportable_int4_tensor` so +it is baked into the serialized weight constant once at pack time. """ from typing import List, Optional, Tuple import torch -from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor from torchao.utils import TorchAOBaseTensor __all__ = [ @@ -113,10 +112,10 @@ def _unpack_nibble_qdata(qdata: torch.Tensor, N: int, K: int) -> torch.Tensor: class CudaCoalescedInt4Tensor(TorchAOBaseTensor): """INT4 weight, uint8 scale + per-256 fp16 step, uint8 zero + per-256 fp16 step. - ExecuTorch-internal; see the module docstring. Mirrors torchao + ExecuTorch-internal; see the module docstring. Mirrors torchao ``Int4Tensor``'s data/attribute layout (so the common tensor utilities and - serialization work) but owns the [n_groups, N] -> [N, n_groups] transpose and - the uint8 re-encoding via :meth:`from_int4_tensor`. + serialization work) but owns the [n_groups, N] -> [N, n_groups] transpose and + the uint8 re-encoding via :meth:`from_exportable_int4_tensor`. """ tensor_data_names = [ @@ -181,20 +180,25 @@ def _quantization_type(self): return s @classmethod - def from_int4_tensor(cls, t: Int4Tensor) -> "CudaCoalescedInt4Tensor": - """Build a coalesced tensor from a torchao ``Int4Tensor``. - - Owns the transpose AND the uint8 re-encoding: torchao stores - scale/zero_point as (n_groups, N) bf16. The CUDA decode kernel reads the - (N, n_groups) uint8 scale/zero *codes* plus per-256-super-block - (N, K/256) fp16 ``scale_step`` / ``zero_point_step`` (scale = scale_code * - scale_step[:, g//8], zero = zero_code * zero_point_step[:, g//8]). The - transpose + encode here is baked into the serialized weight constant so - the exported decode graph has no per-step transpose/clone. + def from_exportable_int4_tensor( + cls, t: "ExportableInt4Tensor" # noqa: F821 + ) -> "CudaCoalescedInt4Tensor": + """Build a coalesced tensor from an ``ExportableInt4Tensor``. + + Owns the [n_groups, N] -> [N, n_groups] transpose and the uint8 + re-encoding, baked into the serialized weight constant at pack time. + torchao stores scale/zero_point as (n_groups, N) bf16; the CUDA decode + kernel reads (N, n_groups) uint8 scale/zero *codes* plus per-256-super- + block (N, K/256) fp16 ``scale_step`` / ``zero_point_step`` (scale = + scale_code * scale_step[:, g//8], zero = zero_code * zero_point_step[:, + g//8]), so the exported decode graph has no per-step transpose/clone. An + ``ExportableInt4Tensor`` (e.g. a decoded GGUF Q4_K) stores ``group_size`` + as a scalar rather than a ``block_size`` list and carries no activation + quantization, so ``block_size`` is reconstructed as ``[1, group_size]``. """ - scale_codes, scale_step = _encode_uint8_per_super(t.scale, t.block_size[-1]) + scale_codes, scale_step = _encode_uint8_per_super(t.scale, t.group_size) zero_codes, zero_point_step = _encode_uint8_per_super( - t.zero_point, t.block_size[-1] + t.zero_point, t.group_size ) return cls( t.qdata, @@ -202,10 +206,8 @@ def from_int4_tensor(cls, t: Int4Tensor) -> "CudaCoalescedInt4Tensor": scale_step, zero_codes, zero_point_step, - t.block_size, + [1, t.group_size], t.shape, - t.act_pre_scale, - t.activation_dtype, ) def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor: diff --git a/backends/cuda/runtime/shims/tests/gen_plain_mm_test_vectors.py b/backends/cuda/runtime/shims/tests/gen_plain_mm_test_vectors.py index c11eed73293..f64a253828c 100644 --- a/backends/cuda/runtime/shims/tests/gen_plain_mm_test_vectors.py +++ b/backends/cuda/runtime/shims/tests/gen_plain_mm_test_vectors.py @@ -23,8 +23,9 @@ 1. torch.manual_seed(case.seed) on CPU. 2. Draw a random bf16 weight ``[N, K]`` then activation ``[M, K]`` (weight first, then activation: fixed order is part of the seed contract). - 3. quantize_weight(..., bits=4, min_max, asymmetric) -> torchao Int4Tensor. - 4. CudaCoalescedInt4Tensor.from_int4_tensor(...) -> qdata [N, K/2] uint8, + 3. quantize_weight(..., bits=4, min_max, asymmetric) -> torchao Int4Tensor, + wrapped as the canonical ExportableInt4Tensor. + 4. CudaCoalescedInt4Tensor.from_exportable_int4_tensor(...) -> qdata [N, K/2] uint8, scale codes [N, K/gs] uint8, scale_step [N, K/256] fp16, zero_point codes [N, K/gs] uint8, zero_point_step [N, K/256] fp16. 5. expected = F.linear(A, tensor.dequantize(bf16)). @@ -157,8 +158,9 @@ def _i8(t: torch.Tensor) -> List[int]: def build_int4(case: Case) -> Dict[str, tuple]: """Return {array_name: (ctype, [ints])} for one INT4 case (CPU only).""" from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor - from executorch.examples.models.gemma4_31b.quant.quantize import quantize_weight - from executorch.examples.models.gemma4_31b.quant.recipe import QuantConfig + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + from executorch.extension.llm.export.quant.quantize import quantize_weight + from executorch.extension.llm.export.quant.recipe import QuantConfig torch.manual_seed(case.seed) # Weight first, then activation: fixed order is part of the seed contract. @@ -167,7 +169,9 @@ def build_int4(case: Case) -> Dict[str, tuple]: config = QuantConfig(bits=4, group_size=case.gs, symmetric=False, method="min_max") int4 = quantize_weight(w, config) - c = CudaCoalescedInt4Tensor.from_int4_tensor(int4) + c = CudaCoalescedInt4Tensor.from_exportable_int4_tensor( + ExportableInt4Tensor.from_int4_tensor(int4) + ) # bf16 dequant @ F.linear reference (kernel adds activation-quant noise). w_deq = c.dequantize(torch.bfloat16) diff --git a/backends/cuda/tests/test_int4_dispatch.py b/backends/cuda/tests/test_int4_dispatch.py index 8e011720487..b3758683359 100644 --- a/backends/cuda/tests/test_int4_dispatch.py +++ b/backends/cuda/tests/test_int4_dispatch.py @@ -34,12 +34,13 @@ import torch.nn.functional as F from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor from executorch.backends.cuda.quantize_op_dispatch.int4_dispatch import _dequant_matmul -from executorch.examples.models.gemma4_31b.quant.pack_cuda import pack_linear_for_cuda -from executorch.examples.models.gemma4_31b.quant.quantize import ( +from executorch.examples.models.gemma4_31b.cuda_packers import pack_linear_for_cuda +from executorch.extension.llm.export.int4 import ExportableInt4Tensor +from executorch.extension.llm.export.quant.quantize import ( dequantize_weight, quantize_weight, ) -from executorch.examples.models.gemma4_31b.quant.recipe import QuantConfig +from executorch.extension.llm.export.quant.recipe import QuantConfig def _require_cuda(tc: unittest.TestCase) -> None: @@ -48,22 +49,26 @@ def _require_cuda(tc: unittest.TestCase) -> None: def _make_int4_linear(N, K, group_size=128, symmetric=False, bias=False): - """Build an nn.Linear with Int4Tensor weight and return (module, bf16_ref_weight). + """Build an nn.Linear with ExportableInt4Tensor weight + bf16 ref weight. - The bf16 reference is the original unquantized weight, so tests can - measure quantization error against the true value. + Mirrors production: weights are converted to ExportableInt4Tensor (the + canonical portable int4 form) before packing for CUDA. The bf16 reference + is the original unquantized weight, so tests can measure quantization + error against the true value. """ w_bf16 = torch.randn(N, K, dtype=torch.bfloat16) config = QuantConfig( bits=4, group_size=group_size, symmetric=symmetric, method="min_max" ) - int4_w = quantize_weight(w_bf16, config) + exportable_w = ExportableInt4Tensor.from_int4_tensor( + quantize_weight(w_bf16, config) + ) # device="cuda" so the random init draws from the CUDA RNG to match the # same random weight as regular int4 dispatch and fit the same numerical # error tolerance. module = nn.Linear(K, N, bias=bias, dtype=torch.bfloat16, device="cuda") - pack_linear_for_cuda(module, {"weight": int4_w}) + pack_linear_for_cuda(module, {"weight": exportable_w}) module.cuda() return module, w_bf16.cuda() @@ -184,9 +189,11 @@ def _check(self, out, ref, tol=0.15): def test_to_cuda(self): w_bf16 = torch.randn(256, 512, dtype=torch.bfloat16) config = QuantConfig(bits=4, group_size=128, symmetric=False, method="min_max") - int4_w = quantize_weight(w_bf16, config) + exportable_w = ExportableInt4Tensor.from_int4_tensor( + quantize_weight(w_bf16, config) + ) module = nn.Linear(512, 256, bias=False) - pack_linear_for_cuda(module, {"weight": int4_w}) + pack_linear_for_cuda(module, {"weight": exportable_w}) module = module.to("cuda") x = torch.randn(1, 512, dtype=torch.bfloat16, device="cuda") self._check(module(x), F.linear(x, w_bf16.cuda())) @@ -228,6 +235,12 @@ def _make_int4_tensor(N, K, group_size=128, symmetric=False): return quantize_weight(w, config), w +def _make_exportable_int4_tensor(N, K, group_size=128, symmetric=False): + """Build an ``ExportableInt4Tensor`` (canonical portable int4) + bf16 ref.""" + t, w = _make_int4_tensor(N, K, group_size=group_size, symmetric=symmetric) + return ExportableInt4Tensor.from_int4_tensor(t), w + + @contextlib.contextmanager def _record_int4_plain_mm(): """Record calls to the decode custom op without needing a GPU. @@ -278,8 +291,8 @@ def test_stock_int4tensor_does_not_route_to_int4_plain_mm(self): def test_coalesced_tensor_routes_to_int4_plain_mm(self): """CudaCoalescedInt4Tensor with M<=4 routes to the decode custom op.""" - t, _ = _make_int4_tensor(16, 256, group_size=32) - c = CudaCoalescedInt4Tensor.from_int4_tensor(t) + t, _ = _make_exportable_int4_tensor(16, 256, group_size=32) + c = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t) x = torch.randn(1, 256, dtype=torch.bfloat16) # M=1 (decode regime) with _record_int4_plain_mm() as calls: out = F.linear(x, c) @@ -288,8 +301,8 @@ def test_coalesced_tensor_routes_to_int4_plain_mm(self): def test_coalesced_tensor_prefill_uses_dequant(self): """M>4 uses inline dequant (no custom op) and is numerically correct.""" - t, _ = _make_int4_tensor(16, 256, group_size=32) - c = CudaCoalescedInt4Tensor.from_int4_tensor(t) + t, _ = _make_exportable_int4_tensor(16, 256, group_size=32) + c = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t) x = torch.randn(8, 256, dtype=torch.bfloat16) # M=8 > 4 (prefill regime) with _record_int4_plain_mm() as calls: out = F.linear(x, c) @@ -312,10 +325,10 @@ def test_square_shape_not_misrouted(self): F.linear(x, t) self.assertEqual(calls, []) - def test_from_int4_tensor_transpose_correct(self): - """from_int4_tensor owns the (n_groups, N) -> (N, n_groups) transpose.""" - t, _ = _make_int4_tensor(24, 256, group_size=64) - c = CudaCoalescedInt4Tensor.from_int4_tensor(t) + def test_from_exportable_int4_tensor_transpose_correct(self): + """from_exportable_int4_tensor owns the (n_groups, N) -> (N, n_groups) transpose.""" + t, _ = _make_exportable_int4_tensor(24, 256, group_size=64) + c = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t) n_groups = 256 // 64 self.assertEqual(tuple(t.scale.shape), (n_groups, 24)) # torchao layout self.assertEqual(tuple(c.scale.shape), (24, n_groups)) # coalesced layout diff --git a/examples/models/gemma4_31b/cuda_packers.py b/examples/models/gemma4_31b/cuda_packers.py new file mode 100644 index 00000000000..888be4fe8a3 --- /dev/null +++ b/examples/models/gemma4_31b/cuda_packers.py @@ -0,0 +1,160 @@ +# 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. + +"""CUDA weight conversion pass: rewrite portable quantized weights to CUDA layouts. + +``convert_quantized_tensors_for_cuda`` is a terminal pass over a loaded model: +it walks every ``nn.Linear`` / ``nn.Embedding`` and repacks its quantized weight +into the ExecuTorch-internal CUDA layouts read by the decode kernels. It runs +after ``load_checkpoint`` (and any fusion), consuming the canonical portable +forms that loading produced (``ExportableInt4Tensor``, ``ExportableGGUFTensor``, +int8 ``IntxUnpackedToInt8Tensor``). + +Linear layouts: + * ``ExportableInt4Tensor`` -> ``CudaCoalescedInt4Tensor`` + (bakes the scale/zero transpose into the coalesced [N, n_groups] layout). A + Q4_K ``ExportableGGUFTensor`` is decoded to an ``ExportableInt4Tensor`` first, + reaching the same layout from the GGUF blob. + * Q5_K ``ExportableGGUFTensor`` -> ``CudaDp4aPlanarInt5Tensor`` (5-bit ql/qh + split bit-planes, asymmetric). + * Q6_K ``ExportableGGUFTensor`` -> ``CudaDp4aPlanarInt6Tensor`` (6-bit ql/qh + split bit-planes). + * genuine INT8 ``IntxUnpackedToInt8Tensor`` -> left unchanged (int8 dp4a path). + +Embedding layouts: + * GGUF ``ExportableGGUFTensor`` (e.g. the tied token embedding) -> gatherable + int8 ``IntxUnpackedToInt8Tensor``. + * genuine int8 ``IntxUnpackedToInt8Tensor`` -> unchanged. + * int4 (``Exportable``/``Int4Tensor``) -> raises (no int4 embedding op). + +No CUDA is required for packing. ``load_checkpoint`` lives in +``extension/llm/export/load.py``; the model-free conversion helpers in +``extension/llm/export/quant/convert.py``. +""" + +import torch +import torch.nn as nn + +from executorch.extension.llm.export.quant.convert import _is_quantized + +# --------------------------------------------------------------------------- +# Per-module weight conversion + + +def pack_linear_for_cuda(module: nn.Module, weights: dict[str, torch.Tensor]) -> None: + """Convert an ``nn.Linear``'s quantized weight to its CUDA layout, in place. + + Routes by weight type: ``ExportableInt4Tensor`` (or Q4_K + ``ExportableGGUFTensor``) -> coalesced INT4, Q5_K ``ExportableGGUFTensor`` -> + packed INT5, Q6_K ``ExportableGGUFTensor`` -> packed INT6, genuine + ``IntxUnpackedToInt8Tensor`` -> int8 passthrough. + """ + from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor + from executorch.backends.cuda.dp4a_planar_int5_tensor import ( + CudaDp4aPlanarInt5Tensor, + ) + from executorch.backends.cuda.dp4a_planar_int6_tensor import ( + CudaDp4aPlanarInt6Tensor, + ) + from executorch.extension.llm.export.gguf import ExportableGGUFTensor + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + from torchao.quantization import IntxUnpackedToInt8Tensor + + w = weights["weight"] + if isinstance(w, ExportableInt4Tensor): + # Canonical portable int4 (what load_checkpoint produces by default). + # CudaCoalescedInt4Tensor reuses the nibble-packed qdata untouched and + # re-encodes scale/zero into the coalesced [N, n_groups] layout the CUDA + # decode kernel reads (int4_dispatch.py / int4_plain_mm.cuh). The + # transpose is baked into the serialized constant, so the exported decode + # graph carries no per-step transpose/clone. + w = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(w) + elif isinstance(w, ExportableGGUFTensor) and w.ggml_type == "q4_k": + # GGUF Q4_K: decode to an ExportableInt4Tensor (bf16 — fp16 is not a CUDA + # target), then coalesce as above. Same end state as the safetensors int4 + # path; the GGUF blob is just the other source. + w = CudaCoalescedInt4Tensor.from_exportable_int4_tensor( + w.to_exportable_int4_tensor(torch.bfloat16) + ) + elif isinstance(w, ExportableGGUFTensor) and w.ggml_type == "q5_k": + # GGUF Q5_K: genuine 5-bit CudaDp4aPlanarInt5Tensor (ql/qh split + # bit-planes, 0.625 B/elem) for the W5A8 dp4a decode kernel. Asymmetric + # (has dmin), so it carries a zero point like INT4. from_exportable_gguf + # reuses the shared Q5_K decode (gguf.py) and bakes the bit-pack in. + w = CudaDp4aPlanarInt5Tensor.from_exportable_gguf(w) + elif isinstance(w, ExportableGGUFTensor) and w.ggml_type == "q6_k": + # GGUF Q6_K: genuine 6-bit CudaDp4aPlanarInt6Tensor (ql/qh split + # bit-planes, 0.75 B/elem) for the W6A8 dp4a decode kernel. + # from_exportable_gguf reuses the shared Q6_K decode (gguf.py), once. + w = CudaDp4aPlanarInt6Tensor.from_exportable_gguf(w) + elif isinstance(w, IntxUnpackedToInt8Tensor): + # Genuine INT8 weight: left unchanged for the int8 dp4a path. The + # mixed-precision HQQ-INT4 ("sensitive") checkpoint reaches this branch + # for its int8 tensors — edge-layer v_proj/down_proj are quantized to + # INT8 while the rest is INT4 (see GEMMA4_31B_SENSITIVE_RECIPE in + # quantize_and_save.py). Q6_K never reaches here (it arrives as an + # ExportableGGUFTensor, handled above), so int4/int6/int8 routing stays + # unambiguous. + pass + else: + raise ValueError(f"Unsupported weight type: {type(w).__name__}") + module.weight = nn.Parameter(w, requires_grad=False) + + +def pack_embedding_for_cuda( + module: nn.Module, weights: dict[str, torch.Tensor] +) -> None: + """Convert an ``nn.Embedding``'s quantized weight to gatherable int8, in place. + + A GGUF ``ExportableGGUFTensor`` (Q4_K/Q6_K, e.g. the tied token embedding) + is decoded to a gatherable int8 ``IntxUnpackedToInt8Tensor`` -- the packed + int4/int6 matmul layouts can't gather, and the lossless int8 decode halves + the footprint vs bf16. A genuine int8 ``IntxUnpackedToInt8Tensor`` (the + safetensors int8 embedding) is assigned unchanged. INT4 is unsupported (no + int4 embedding op). + """ + from executorch.extension.llm.export.gguf import ExportableGGUFTensor + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor + + w = weights["weight"] + if isinstance(w, ExportableGGUFTensor): + w = w.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16) + elif isinstance(w, (Int4Tensor, ExportableInt4Tensor)): + raise ValueError( + "Only 8-bit embedding quantization is supported on CUDA. " + "INT4 does not implement the embedding op." + ) + module.weight = nn.Parameter(w, requires_grad=False) + + +# --------------------------------------------------------------------------- +# Model pass + + +def convert_quantized_tensors_for_cuda(model: nn.Module) -> None: + """Rewrite every quantized ``nn.Linear`` / ``nn.Embedding`` weight in ``model`` + into its CUDA layout, in place. + + Terminal pass: run after ``load_checkpoint`` and any fusion. Walks the model, + dispatches by module type (the tied-embedding split: ``lm_head`` -> + ``CudaCoalescedInt4Tensor``, ``embed_tokens`` -> gatherable int8), and mutates + each module's ``weight``. Plain (unquantized) weights are left untouched. + """ + for module in model.modules(): + weight = getattr(module, "weight", None) + if weight is None or not _is_quantized(weight): + continue + if isinstance(module, nn.Linear): + pack_linear_for_cuda(module, {"weight": weight}) + elif isinstance(module, nn.Embedding): + pack_embedding_for_cuda(module, {"weight": weight}) + else: + raise ValueError( + f"Cannot convert quantized weight on unsupported module type " + f"{type(module).__name__}: only nn.Linear and nn.Embedding are " + f"handled by the CUDA pass." + ) diff --git a/examples/models/gemma4_31b/export.py b/examples/models/gemma4_31b/export.py index 01f4c3e1957..af1a6a0dfb2 100644 --- a/examples/models/gemma4_31b/export.py +++ b/examples/models/gemma4_31b/export.py @@ -40,27 +40,98 @@ # Load paths +def _build_and_pack( + config: Gemma4_31BConfig, + path: str, + backend: str, + *, + key_map=None, + tie_map=None, +) -> Gemma4_31B: + """Build a meta-device model, stream ``path`` into it, and finalize. + + Shared by ``load_prequantized_model`` (safetensors) and ``load_gguf_model`` + (GGUF). ``load_checkpoint`` picks the reader by ``path`` and asserts every + parameter is materialized; we then fill runtime buffers left on meta (RoPE + tables, KV caches, scalar constants) and switch to eval. Pass ``key_map`` / + ``tie_map`` for formats whose tensor names differ from the model's or that + tie weights (GGUF). + """ + from executorch.extension.llm.export.load import load_checkpoint + + # Validate the backend before any file I/O so an unknown backend fails fast + # rather than erroring later on a missing/large checkpoint. + if backend not in _SUPPORTED_BACKENDS: + raise ValueError( + f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}." + ) + + print("Building model on meta device...") + with torch.device("meta"): + model = Gemma4_31B(config) + + load_checkpoint( + path, + model, + key_map=key_map, + tie_map=tie_map, + dtype=torch.bfloat16, + ) + materialize_runtime_buffers(model, dtype=torch.bfloat16) + _apply_backend_pass(model, backend) + model.eval() + + print(f"Model: {config.num_hidden_layers} layers, hidden={config.hidden_size}") + return model + + def load_prequantized_model( prequantized_dir: str, max_seq_len: int = 4096, backend: str = "cuda", ) -> tuple[Gemma4_31B, Gemma4_31BConfig]: - """Load a quantized checkpoint and pack for the target backend.""" + """Load a quantized safetensors checkpoint and pack for the target backend.""" config = Gemma4_31BConfig.from_hf_config( os.path.join(prequantized_dir, "config.json") ) config.max_seq_len = max_seq_len - print("Building model on meta device...") - with torch.device("meta"): - model = Gemma4_31B(config) - safetensors_path = os.path.join(prequantized_dir, "model.safetensors") print(f"Loading quantized checkpoint from {safetensors_path}...") - _pack_for_backend(model, safetensors_path, backend) - model.eval() + model = _build_and_pack(config, safetensors_path, backend) + return model, config - print(f"Model: {config.num_hidden_layers} layers, hidden={config.hidden_size}") + +def load_gguf_model( + gguf_path: str, + max_seq_len: int = 4096, + backend: str = "cuda", + config: Gemma4_31BConfig | None = None, +) -> tuple[Gemma4_31B, Gemma4_31BConfig]: + """Load a GGUF file and pack for the target backend. + + ``key_map`` (``gguf_to_model_key``) remaps GGUF tensor names to model FQNs. + GGUF ties the token embedding and lm_head into one tensor; ``tie_map`` fans + it out to ``lm_head`` -- untied on CUDA (``clone=True``: ``lm_head`` keeps a + packed matmul weight while ``pack_embedding_for_cuda`` decodes a gatherable + int8 copy) and kept tied on MLX (``clone=False``: both share one quantized + tensor). ``config`` defaults to the full Gemma 4 31B config; pass a smaller + one (e.g. in tests) to load a GGUF for a tiny model. + """ + from executorch.examples.models.gemma4_31b.gguf_loader import gguf_to_model_key + + if config is None: + config = Gemma4_31BConfig(max_seq_len=max_seq_len) + + tie_map = {"embed_tokens.weight": ("lm_head.weight", backend == "cuda")} + print(f"Loading GGUF from {gguf_path}...") + model = _build_and_pack( + config, + gguf_path, + backend, + key_map=gguf_to_model_key, + tie_map=tie_map, + ) return model, config @@ -71,8 +142,9 @@ def load_and_quantize( backend: str = "cuda", ) -> tuple[Gemma4_31B, Gemma4_31BConfig]: """Load bf16 checkpoint, quantize, pack — one shot.""" - from executorch.examples.models.gemma4_31b.quant import pack_model, quantize_model from executorch.examples.models.gemma4_31b.quantize_and_save import _RECIPES + from executorch.extension.llm.export.load import assign_state_dict + from executorch.extension.llm.export.quant import quantize_model recipe = _RECIPES[recipe_name] @@ -89,7 +161,8 @@ def load_and_quantize( print(f"Packing for {backend}...") with torch.device("meta"): model = Gemma4_31B(config) - pack_model(model, state_dict, packers=_get_packers(backend)) + assign_state_dict(model, state_dict) + _apply_backend_pass(model, backend) model.eval() print(f"Model: {config.num_hidden_layers} layers, hidden={config.hidden_size}") @@ -103,35 +176,27 @@ def load_and_quantize( _SUPPORTED_BACKENDS = ("cuda", "mlx") -def _get_packers(backend: str) -> dict: - if backend == "cuda": - from executorch.examples.models.gemma4_31b.quant import DEFAULT_CUDA_PACKERS +def _apply_backend_pass(model: nn.Module, backend: str) -> None: + """Run the backend's terminal layout-conversion pass over a loaded model. - return DEFAULT_CUDA_PACKERS + CUDA repacks quantized weights into the coalesced/planar decode layouts; MLX + (gemma4 is dense — no switch linears) needs no extra pass, the portable forms + export directly. + """ if backend == "mlx": - from executorch.examples.models.gemma4_31b.quant import DEFAULT_MLX_PACKERS + return + if backend == "cuda": + from executorch.examples.models.gemma4_31b.cuda_packers import ( + convert_quantized_tensors_for_cuda, + ) - return DEFAULT_MLX_PACKERS + convert_quantized_tensors_for_cuda(model) + return raise ValueError( f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}." ) -def _pack_for_backend(model: nn.Module, path: str, backend: str) -> None: - if backend == "cuda": - from executorch.examples.models.gemma4_31b.quant import load_and_pack_for_cuda - - load_and_pack_for_cuda(path, model) - elif backend == "mlx": - from executorch.examples.models.gemma4_31b.quant import load_and_pack_for_mlx - - load_and_pack_for_mlx(path, model) - else: - raise ValueError( - f"Unsupported backend: {backend!r}. Supported: {_SUPPORTED_BACKENDS}." - ) - - # --------------------------------------------------------------------------- # Export + lower @@ -522,8 +587,6 @@ def main() -> None: backend=args.backend, ) elif args.gguf: - from executorch.examples.models.gemma4_31b.gguf_loader import load_gguf_model - model, config = load_gguf_model( args.gguf, max_seq_len=args.max_seq_len, backend=args.backend ) diff --git a/examples/models/gemma4_31b/gguf_loader.py b/examples/models/gemma4_31b/gguf_loader.py index 325d76d2306..3d7074c08a6 100644 --- a/examples/models/gemma4_31b/gguf_loader.py +++ b/examples/models/gemma4_31b/gguf_loader.py @@ -4,35 +4,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Load a GGUF file into a Gemma 4 31B model. +"""GGUF tensor-name mapping for Gemma 4 31B. -Streams tensors one at a time via the shared loader in -``extension/llm/export/gguf.py`` (each quantized weight arrives as an -``ExportableGGUFTensor`` wrapping the raw GGUF blob), remaps GGUF names to model -FQNs, handles the tied embed/lm_head, and converts each weight for the target -backend: - -* **MLX**: every quantized weight stays an ``ExportableGGUFTensor`` and is lowered - by the MLX GGUF pattern (Q6_K custom kernels, Q4_K native affine ops) for both - linear and embedding. ``embed_tokens`` and ``lm_head`` stay tied -- they share - the one quantized tensor. -* **CUDA**: Q4_K -> ``Int4Tensor``, Q6_K -> ``CudaDp4aPlanarInt6Tensor`` (a genuine - 6-bit packed weight, lossless, symmetric). ``embed_tokens`` and ``lm_head`` are - untied: ``lm_head`` keeps a packed (int6/int4) matmul weight, while the token - embedding becomes a gatherable ``IntxUnpackedToInt8Tensor`` (int8) -- the truly - packed int4/int6 tensors can't gather. For the Q6_K tied weight the decode is - done once and shared between the two, avoiding a whole-tensor bf16 dequant and - a second decode (see ``_untie_embed_lm_head``). - -Usage: - model, config = load_gguf_model("model.gguf", backend="cuda") - model, config = load_gguf_model("model.gguf", backend="mlx") +Maps GGUF tensor names (``blk.0.attn_q.weight``, ``token_embd.weight``, ...) to +Gemma model FQNs. Used as the ``key_map`` for ``load_checkpoint`` when loading a +GGUF checkpoint -- see ``export.load_gguf_model``, which pairs this with a +``tie_map`` to fan the tied embedding out to ``lm_head``. """ from typing import Optional -import torch - # GGUF pattern → model FQN pattern. ``{}`` is the layer index. _KEY_MAP = { "token_embd.weight": "embed_tokens.weight", @@ -77,176 +58,3 @@ def gguf_to_model_key(gguf_key: str) -> Optional[str]: return model_pat.replace("{}", layer_str) return None - - -def _validate_no_meta(model): - """Ensure all parameters have been loaded.""" - for fqn, p in model.named_parameters(): - if p.device.type == "meta": - raise RuntimeError( - f"Weight '{fqn}' not found in GGUF file " - f"(model/checkpoint version mismatch?)" - ) - for p in model.parameters(): - p.requires_grad_(False) - - -def _convert_weight(model, model_key: str, gtensor, backend: str): - """Convert an ``ExportableGGUFTensor`` to the per-backend module weight.""" - if backend == "mlx": - return gtensor - # CUDA: Q4_K -> torchao Int4Tensor. Q6_K stays the raw ExportableGGUFTensor - # (like MLX) -- the CUDA packer repacks it into CudaDp4aPlanarInt6Tensor via - # CudaDp4aPlanarInt6Tensor.from_exportable_gguf, so the Q6_K block decode is - # owned by gguf.py and reused, not duplicated here. - if gtensor.ggml_type == "q4_k": - return gtensor.to_int4_tensor() - return gtensor - - -def _resolve_tied_lm_head(model, lm_head_weight, packers): - """Assign a tied lm_head (GGUF ties it to the token embedding).""" - from executorch.examples.models.gemma4_31b.quant import pack_one - - lm_head = getattr(model.lm_head, "weight", None) - if lm_head is None or lm_head.device.type != "meta": - return - if lm_head_weight is not None: - pack_one(model, "lm_head.weight", lm_head_weight, packers) - else: - pack_one( - model, "lm_head.weight", model.embed_tokens.weight.data.clone(), packers - ) - - -def _untie_embed_lm_head(model, gtensor, weight, backend): - """Untie the GGUF token-embed / lm_head weight, returning ``(embed, lm_head)``. - - GGUF ties ``embed_tokens`` and ``lm_head`` to one quantized weight. The - returned ``lm_head`` is packed into ``model.lm_head`` after the streaming loop - (``_resolve_tied_lm_head``), or is ``None`` when this function already - assigned it. - - * **MLX**: keep both tied on the raw ``ExportableGGUFTensor``. - * **CUDA** (Q6_K or Q4_K): untie so ``lm_head`` keeps a packed low-bit matmul - weight while the token embedding becomes a gatherable int8 - ``IntxUnpackedToInt8Tensor`` -- the truly packed int4/int6 tensors can't - gather. Instead of dequantizing the whole ~1.4 B-element weight to bf16 - (2 B/elem), decode it once to int8 (1 B/elem; the decode is lossless so the - result is numerically identical), halving the embedding's host + GPU-constant - footprint. The token embedding (Q4_K for the Gemma checkpoint) is the single - biggest weight, so this is the dominant saving vs the bf16 path. ``lm_head``: - - Q6_K -> ``CudaDp4aPlanarInt6Tensor`` from the *same* int8 decode and - assigned here (``pack_linear_for_cuda`` would mis-route an int8 tensor to - the int8 path), so the post-loop resolve is a no-op. - - Q4_K -> kept as the native ``Int4Tensor`` and returned, so - ``_resolve_tied_lm_head`` packs it to ``CudaCoalescedInt4Tensor`` (same - as a regular Q4_K linear). - * **CUDA, other types**: fall back to the bf16 embedding. - """ - if backend == "mlx": - return weight, gtensor - - if gtensor.ggml_type in ("q6_k", "q4_k"): - intx = gtensor.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16) - if gtensor.ggml_type == "q6_k": - import torch.nn as nn - from executorch.backends.cuda.dp4a_planar_int6_tensor import ( - CudaDp4aPlanarInt6Tensor, - ) - - model.lm_head.weight = nn.Parameter( - CudaDp4aPlanarInt6Tensor._from_intx_int8(intx), requires_grad=False - ) - return intx, None - # Q4_K: ``weight`` is the native Int4Tensor; let _resolve_tied_lm_head - # pack it to CudaCoalescedInt4Tensor. Only the embedding switches to int8. - return intx, weight - - from executorch.examples.models.gemma4_31b.quant import dequantize_weight - - return dequantize_weight(weight, torch.bfloat16), weight - - -def load_gguf_model( - gguf_path: str, - max_seq_len: int = 4096, - backend: str = "cuda", - config=None, -) -> tuple: - """Load a GGUF file, remap keys, and convert weights for the target backend. - - Streams tensors one at a time for low peak memory. GGUF ties ``embed_tokens`` - and ``lm_head``: on MLX they stay tied (one shared quantized tensor); on CUDA - they are untied so the embedding can be dequantized for the gather while - ``lm_head`` keeps its quantization. See the module docstring for the - per-backend conversion details. - - ``config`` defaults to the full Gemma 4 31B config; pass a smaller - ``Gemma4_31BConfig`` (e.g. in tests) to load a GGUF for a tiny model. - - Returns ``(model, config)``. - """ - from executorch.examples.models.gemma4_31b.model import ( - Gemma4_31B, - Gemma4_31BConfig, - materialize_runtime_buffers, - ) - from executorch.examples.models.gemma4_31b.quant import pack_one - from executorch.extension.llm.export.gguf import ExportableGGUFTensor, iter_gguf - - if backend == "cuda": - from executorch.examples.models.gemma4_31b.quant import DEFAULT_CUDA_PACKERS - - packers = DEFAULT_CUDA_PACKERS - elif backend == "mlx": - from executorch.examples.models.gemma4_31b.quant import DEFAULT_MLX_PACKERS - - packers = DEFAULT_MLX_PACKERS - else: - raise ValueError(f"Unsupported backend: {backend!r}. Supported: 'cuda', 'mlx'.") - - if config is None: - config = Gemma4_31BConfig(max_seq_len=max_seq_len) - - print("Building model on meta device...") - with torch.device("meta"): - model = Gemma4_31B(config) - - lm_head_weight = None # tied weight resolved into lm_head after the loop - n_processed = 0 - - print(f"Streaming GGUF from {gguf_path}...") - for gguf_name, value in iter_gguf(gguf_path): - model_key = gguf_to_model_key(gguf_name) - if model_key is None: - continue - - if isinstance(value, ExportableGGUFTensor): - weight = _convert_weight(model, model_key, value, backend) - if model_key == "embed_tokens.weight": - weight, lm_head_weight = _untie_embed_lm_head( - model, value, weight, backend - ) - value = weight - elif value.dtype in (torch.float32, torch.float16): - value = value.to(torch.bfloat16) - - pack_one(model, model_key, value, packers) - - n_processed += 1 - if n_processed % 100 == 0: - print(f" Processed {n_processed} tensors...") - - _resolve_tied_lm_head(model, lm_head_weight, packers) - - # Fill RoPE tables / KV caches / scalar constants (left on meta by the - # streaming load), matching load_prequantized_model so the CUDA and eager - # forward paths get bf16 runtime buffers instead of float32 defaults. - materialize_runtime_buffers(model, dtype=torch.bfloat16) - - _validate_no_meta(model) - model.eval() - - print(f"Model: {config.num_hidden_layers} layers, hidden={config.hidden_size}") - return model, config diff --git a/examples/models/gemma4_31b/inference.py b/examples/models/gemma4_31b/inference.py index 121e1deb97e..7cafe92b192 100644 --- a/examples/models/gemma4_31b/inference.py +++ b/examples/models/gemma4_31b/inference.py @@ -217,7 +217,7 @@ def main() -> None: eos_token_ids = {1, 50, 106} if args.gguf: - from executorch.examples.models.gemma4_31b.gguf_loader import load_gguf_model + from executorch.examples.models.gemma4_31b.export import load_gguf_model model, config = load_gguf_model( args.gguf, args.max_seq_len, backend=args.backend diff --git a/examples/models/gemma4_31b/model.py b/examples/models/gemma4_31b/model.py index d953541a244..f9727f96ac8 100644 --- a/examples/models/gemma4_31b/model.py +++ b/examples/models/gemma4_31b/model.py @@ -506,9 +506,14 @@ def from_hf_checkpoint( ) -> tuple["Gemma4_31B", Gemma4_31BConfig]: """Build the model on `meta` and load weights from the HF safetensors checkpoint. - Uses lazy shard-by-shard loading + assign=True so peak memory stays at - roughly one shard's worth of weights. + Streams the shards through the shared ``load_checkpoint`` loader, remapping + HF tensor names via ``_hf_to_model_key`` (which also skips multimodal keys) + and filling the tied ``lm_head`` from ``embed_tokens`` via ``tie_map``. + Runtime buffers (KV caches, RoPE tables) stay on meta -- callers fill them + with ``materialize_runtime_buffers``. """ + from executorch.extension.llm.export.load import load_checkpoint + config = Gemma4_31BConfig.from_hf_config(os.path.join(model_dir, "config.json")) config.max_seq_len = max_seq_len @@ -520,35 +525,11 @@ def from_hf_checkpoint( model = Gemma4_31B(config) print(f"Loading weights from {model_dir}...") - state_dict = _load_and_remap_checkpoint(model_dir, config) - - # Tied embeddings: copy embedding weight into lm_head when missing. - if "lm_head.weight" not in state_dict and "embed_tokens.weight" in state_dict: - state_dict["lm_head.weight"] = state_dict["embed_tokens.weight"] - - missing, unexpected = model.load_state_dict( - state_dict, strict=False, assign=True - ) - - # Runtime buffers (KV caches, RoPE tables, masks) are zero-initialized - # and not in the checkpoint — those are the "expected" missing keys. - runtime_prefixes = ( - ".kv_cache.", - ".inv_freq", - "embed_normalizer", - "logit_softcap", - "cache_positions", - ) - actual_missing = set(missing) - expected = {k for k in actual_missing if any(p in k for p in runtime_prefixes)} - extra = actual_missing - expected - if extra: - print(f" WARNING: missing weight keys: {sorted(extra)[:10]}") - if unexpected: - print(f" WARNING: unexpected keys: {sorted(unexpected)[:10]}") - print( - f" Loaded {len(state_dict)} tensors " - f"({len(expected)} runtime buffers OK)" + load_checkpoint( + model_dir, + model, + key_map=_hf_to_model_key, + tie_map={"embed_tokens.weight": ("lm_head.weight", False)}, ) return model, config @@ -609,38 +590,6 @@ def _hf_to_model_key(hf_key: str) -> Optional[str]: return None -def _load_and_remap_checkpoint(model_dir: str, config: Gemma4_31BConfig) -> dict: - """Stream-load safetensors shards and remap keys to model state_dict keys.""" - from safetensors import safe_open - - index_path = os.path.join(model_dir, "model.safetensors.index.json") - if os.path.exists(index_path): - with open(index_path, "r") as f: - index = json.load(f) - shard_files = sorted(set(index["weight_map"].values())) - elif os.path.exists(os.path.join(model_dir, "model.safetensors")): - shard_files = ["model.safetensors"] - else: - raise FileNotFoundError(f"No safetensors checkpoint in {model_dir}") - - state_dict: dict[str, torch.Tensor] = {} - skipped = 0 - for shard_file in shard_files: - shard_path = os.path.join(model_dir, shard_file) - with safe_open(shard_path, framework="pt", device="cpu") as f: - for ckpt_key in f.keys(): - model_key = _hf_to_model_key(ckpt_key) - if model_key is None: - skipped += 1 - continue - tensor = f.get_tensor(ckpt_key) - # layer_scalar in checkpoint is shape (1,) bf16 — keep as-is. - state_dict[model_key] = tensor - if skipped > 0: - print(f" Skipped {skipped} non-text keys (vision tower, etc.)") - return state_dict - - # --------------------------------------------------------------------------- # Runtime buffer materialization diff --git a/examples/models/gemma4_31b/quant/README.md b/examples/models/gemma4_31b/quant/README.md deleted file mode 100644 index 8906a0faede..00000000000 --- a/examples/models/gemma4_31b/quant/README.md +++ /dev/null @@ -1,54 +0,0 @@ -# quant/ - -Quantization framework: **recipe → quantize → pack**. - -## Files - -| File | Concern | Depends on | -|---|---|---| -| `recipe.py` | **Policy** — what to quantize, what precision, which layers | nothing | -| `quantize.py` | **Computation** — produces torchao subclass tensors | recipe, torchao | -| `pack.py` | **Packing dispatch** — `pack_model` (bulk) and `pack_one` (streaming) | — | -| `pack_cuda.py` | **CUDA packing** — passes Int4Tensor/IntxUnpacked through for CUDA dispatch | pack | -| `pack_mlx.py` | **MLX packing** — converts Int4Tensor → IntxUnpacked, regroups per-axis embeddings | pack | - -GGUF import (unpacking Q4_K/Q6_K blocks) now lives in the shared -`extension/llm/export/gguf.py`. - -## Data flow - -``` -QuantRecipe → quantize_model() → state_dict{Int4Tensor, IntxUnpackedToInt8Tensor, Tensor} → safetensors → state_dict → pack_model() → runtime model -``` - -Quantized weights are stored as torchao tensor subclasses: -- **Int4Tensor** — 4-bit weights (nibble-packed qdata + transposed scale/zero_point) -- **IntxUnpackedToInt8Tensor** — 8-bit weights (int8 qdata + scale + zero_point) - -These are the canonical interchange formats from torchao. Everything left -of `save()` is backend-agnostic. Everything right is backend-specific. - -## Adding a new backend - -Write a `pack_.py` with per-module packers and a default registry: - -```python -def pack_linear_for_metal(module, weights): ... -DEFAULT_METAL_PACKERS = {nn.Linear: pack_linear_for_metal} -``` - -Call `pack_model(model, state_dict, packers=DEFAULT_METAL_PACKERS)`. -No changes to recipe or quantize. - -## On-disk format - -Uses torchao's safetensors integration (`torchao.prototype.safetensors`). -Each tensor subclass is decomposed into its inner tensors -(e.g., `layer._weight_qdata`, `layer._weight_scale`) plus JSON metadata -recording the subclass type and attributes. Plain tensors are stored as-is. -The format is compatible with torchao's `save_pretrained` / `load_pretrained`. - -## TODO - -- `pack_metal.py` — Metal backend packer. -- GGUF quant types (Q5_K, Q8_0): extend `extension/llm/export/gguf.py`. diff --git a/examples/models/gemma4_31b/quant/__init__.py b/examples/models/gemma4_31b/quant/__init__.py deleted file mode 100644 index 7e9ab97a1bb..00000000000 --- a/examples/models/gemma4_31b/quant/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# 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. - -from .pack import ModulePackerFn, pack_model, pack_one # noqa: F401 -from .pack_cuda import DEFAULT_CUDA_PACKERS, load_and_pack_for_cuda # noqa: F401 -from .pack_mlx import DEFAULT_MLX_PACKERS, load_and_pack_for_mlx # noqa: F401 -from .quantize import dequantize_weight, quantize_model, quantize_weight # noqa: F401 -from .recipe import QuantConfig, QuantRecipe, QuantRule # noqa: F401 diff --git a/examples/models/gemma4_31b/quant/pack.py b/examples/models/gemma4_31b/quant/pack.py deleted file mode 100644 index 95abc43546a..00000000000 --- a/examples/models/gemma4_31b/quant/pack.py +++ /dev/null @@ -1,106 +0,0 @@ -# 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. - -"""Backend-agnostic model packing: quantized state dict → runtime model. - -``pack_model`` walks a state dict, groups quantized weights by parent -module, and dispatches to per-module packer functions. Each backend -(``pack_cuda.py``, future ``pack_metal.py``) provides its own packers dict. -""" - -from collections import defaultdict -from typing import Callable - -import torch -import torch.nn as nn - -# Packer signature: receives the module + a dict of its quantized weights -# (keyed by attribute name), modifies module in-place. -ModulePackerFn = Callable[[nn.Module, dict[str, torch.Tensor]], None] - - -def _is_quantized(value: torch.Tensor) -> bool: - """Check if a tensor is a torchao quantized subclass.""" - from torchao.utils import TorchAOBaseTensor - - return isinstance(value, TorchAOBaseTensor) - - -def pack_model( - model: nn.Module, - state_dict: dict[str, torch.Tensor], - packers: dict[type, ModulePackerFn], -) -> None: - """Pack a state dict into ``model`` using the given packers. - - Quantized weights (torchao tensor subclasses) are grouped by parent - module and dispatched to per-module packers. Plain tensors are assigned - directly as parameters or buffers. - """ - # Separate quantized and unquantized - for fqn, value in state_dict.items(): - if not _is_quantized(value): - pack_one(model, fqn, value, packers) - - # Group quantized weights by parent module - module_weights: dict[str, dict[str, torch.Tensor]] = defaultdict(dict) - for fqn, value in state_dict.items(): - if _is_quantized(value): - parts = fqn.rsplit(".", 1) - parent_fqn = parts[0] if len(parts) > 1 else "" - attr = parts[-1] - module_weights[parent_fqn][attr] = value - - for parent_fqn, weights in module_weights.items(): - module = model.get_submodule(parent_fqn) if parent_fqn else model - packer = packers.get(type(module)) - if packer is None: - raise ValueError( - f"No packer registered for {type(module).__name__} at '{parent_fqn}'. " - f"Registered types: {[t.__name__ for t in packers]}." - ) - packer(module, weights) - - for fqn, p in model.named_parameters(): - if p.device.type == "meta": - raise RuntimeError( - f"Weight '{fqn}' not found in checkpoint " - f"(model/checkpoint version mismatch?)" - ) - - for p in model.parameters(): - p.requires_grad_(False) - - -def pack_one( - model: nn.Module, - fqn: str, - value: torch.Tensor, - packers: dict[type, ModulePackerFn], -) -> None: - """Pack a single weight into ``model``. - - Quantized subclass tensors are dispatched to the packer for the parent - module's type. Plain tensors are assigned directly. - """ - parts = fqn.rsplit(".", 1) - parent_fqn = parts[0] if len(parts) > 1 else "" - attr = parts[-1] - parent = model.get_submodule(parent_fqn) if parent_fqn else model - - if _is_quantized(value): - packer = packers.get(type(parent)) - if packer is None: - raise ValueError( - f"No packer registered for {type(parent).__name__} at '{parent_fqn}'. " - f"Registered types: {[t.__name__ for t in packers]}." - ) - packer(parent, {attr: value}) - else: - if isinstance(getattr(parent, attr, None), nn.Parameter): - setattr(parent, attr, nn.Parameter(value, requires_grad=False)) - else: - parent.register_buffer(attr, value) diff --git a/examples/models/gemma4_31b/quant/pack_cuda.py b/examples/models/gemma4_31b/quant/pack_cuda.py deleted file mode 100644 index f3f4cf704bc..00000000000 --- a/examples/models/gemma4_31b/quant/pack_cuda.py +++ /dev/null @@ -1,162 +0,0 @@ -# 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. - -"""CUDA packer: assign quantized weights to model modules. - -Repacks native torchao quantized tensors into the ExecuTorch-internal CUDA -layouts read by the decode kernels: - - * ``Int4Tensor`` -> ``CudaCoalescedInt4Tensor`` (bakes the scale/zero transpose - into the coalesced [N, n_groups] layout). - * Q5_K ``ExportableGGUFTensor`` -> ``CudaDp4aPlanarInt5Tensor`` (the genuine 5-bit - ql/qh split bit-planes, asymmetric; the Q5_K block decode is reused from - gguf.py, not duplicated). - * Q6_K ``ExportableGGUFTensor`` -> ``CudaDp4aPlanarInt6Tensor`` (the genuine 6-bit - ql/qh split bit-planes; the Q6_K block decode is reused from gguf.py, not - duplicated). - -A genuine INT8 ``IntxUnpackedToInt8Tensor`` is left unchanged for the int8 path -(Q6_K no longer arrives as an int8 tensor, so the routing is unambiguous). -The quantize_op_dispatch package (``int4_dispatch`` / ``int6_dispatch`` / -``int8_dispatch``) handles F.linear at runtime. - -No CUDA is required for packing. The backend-agnostic ``pack_model`` -dispatcher lives in ``pack.py``. -""" - -import json - -import torch -import torch.nn as nn - -from .pack import ModulePackerFn, pack_model # noqa: F401 - -# --------------------------------------------------------------------------- -# Per-module packers - - -def pack_linear_for_cuda(module: nn.Module, weights: dict[str, torch.Tensor]) -> None: - """Assign a quantized weight to an ``nn.Linear`` module. - - Routes by weight type: ``Int4Tensor`` -> coalesced INT4, Q5_K - ``ExportableGGUFTensor`` -> packed INT5, Q6_K ``ExportableGGUFTensor`` -> - packed INT6, genuine ``IntxUnpackedToInt8Tensor`` -> int8 passthrough. - """ - from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor - from executorch.backends.cuda.dp4a_planar_int5_tensor import ( - CudaDp4aPlanarInt5Tensor, - ) - from executorch.backends.cuda.dp4a_planar_int6_tensor import ( - CudaDp4aPlanarInt6Tensor, - ) - from executorch.extension.llm.export.gguf import ExportableGGUFTensor - from torchao.quantization import IntxUnpackedToInt8Tensor - from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor - - w = weights["weight"] - if isinstance(w, Int4Tensor): - # Convert to the ExecuTorch-internal CudaCoalescedInt4Tensor, which - # repacks scale/zero from torchao's native [n_groups, N] layout into the - # coalesced [N, n_groups] layout the CUDA decode kernel reads (see - # int4_dispatch.py / int4_plain_mm.cuh). The transpose lives in - # CudaCoalescedInt4Tensor.from_int4_tensor, so it is baked into the - # serialized weight constant and the exported decode graph carries NO - # per-step transpose/clone — AOTInductor (freezing=False) does not - # constant-fold ops on parameters, so the transpose must already live in - # the constant for the coalesced layout to pay off. - w = CudaCoalescedInt4Tensor.from_int4_tensor(w) - elif isinstance(w, ExportableGGUFTensor) and w.ggml_type == "q5_k": - # GGUF Q5_K: repack the native ExportableGGUFTensor into the genuine 5-bit - # CudaDp4aPlanarInt5Tensor (ql/qh split bit-planes, 0.625 B/elem) for the - # W5A8 dp4a decode kernel. Q5_K is asymmetric (has dmin), so this tensor - # carries a zero point like INT4. from_exportable_gguf reuses the shared - # Q5_K decode (gguf.py) then bakes the bit-pack into the weight constant. - w = CudaDp4aPlanarInt5Tensor.from_exportable_gguf(w) - elif isinstance(w, ExportableGGUFTensor) and w.ggml_type == "q6_k": - # GGUF Q6_K: repack the native ExportableGGUFTensor into the genuine 6-bit - # CudaDp4aPlanarInt6Tensor (ql/qh split bit-planes, 0.75 B/elem) for the - # W6A8 dp4a decode kernel. from_exportable_gguf reuses the shared Q6_K - # decode (gguf.py) then bakes the bit-pack into the weight constant, once. - w = CudaDp4aPlanarInt6Tensor.from_exportable_gguf(w) - elif isinstance(w, IntxUnpackedToInt8Tensor): - # Genuine INT8 weight: left unchanged for the int8 dp4a path. The - # mixed-precision HQQ-INT4 ("sensitive") checkpoint reaches this branch - # for its int8 tensors — edge-layer v_proj/down_proj are quantized to - # INT8 while the rest is INT4 (see GEMMA4_31B_SENSITIVE_RECIPE in - # quantize_and_save.py). Q6_K never reaches here (it arrives as an - # ExportableGGUFTensor, handled above), so int4 vs int6 vs int8 routing - # stays unambiguous. - pass - else: - raise ValueError(f"Unsupported weight type: {type(w).__name__}") - module.weight = nn.Parameter(w, requires_grad=False) - - -def pack_embedding_for_cuda( - module: nn.Module, weights: dict[str, torch.Tensor] -) -> None: - """Assign a quantized weight to an ``nn.Embedding`` (INT8 only).""" - from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor - - w = weights["weight"] - if isinstance(w, Int4Tensor): - raise ValueError( - "Only 8-bit embedding quantization is supported on CUDA. " - "INT4 does not implement the embedding op." - ) - module.weight = nn.Parameter(w, requires_grad=False) - - -DEFAULT_CUDA_PACKERS: dict[type, ModulePackerFn] = { - nn.Linear: pack_linear_for_cuda, - nn.Embedding: pack_embedding_for_cuda, -} - - -# --------------------------------------------------------------------------- -# Load + pack (I/O wrapper) - - -def load_and_pack_for_cuda( - path: str, - model: nn.Module, - packers: dict[type, ModulePackerFn] | None = None, -) -> None: - """Load a quantized safetensors file and assign weights to the model.""" - from safetensors import safe_open - from torchao.prototype.safetensors.safetensors_support import ( - unflatten_tensor_state_dict, - ) - - from .pack import pack_one - - _packers = packers or DEFAULT_CUDA_PACKERS - with safe_open(path, framework="pt", device="cpu") as f: - metadata = f.metadata() - all_keys = list(f.keys()) - tensor_names = json.loads(metadata.get("tensor_names", "[]")) - - for name in tensor_names: - parts = name.rsplit(".", 1) - module_fqn = parts[0] if len(parts) > 1 else "" - weight_name = parts[-1] - prefix = ( - f"{module_fqn}._{weight_name}_" if module_fqn else f"_{weight_name}_" - ) - partial = {} - for key in all_keys: - if key.startswith(prefix) or key == name: - partial[key] = f.get_tensor(key) - result, _ = unflatten_tensor_state_dict(partial, metadata) - for fqn, value in result.items(): - pack_one(model, fqn, value, _packers) - - for fqn, p in model.named_parameters(): - if p.device.type == "meta": - raise RuntimeError( - f"Weight '{fqn}' not found in checkpoint " - f"(model/checkpoint version mismatch?)" - ) diff --git a/examples/models/gemma4_31b/quant/pack_mlx.py b/examples/models/gemma4_31b/quant/pack_mlx.py deleted file mode 100644 index 45cea5b2297..00000000000 --- a/examples/models/gemma4_31b/quant/pack_mlx.py +++ /dev/null @@ -1,103 +0,0 @@ -# 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. - -"""MLX packer: convert quantized weights to MLX-compatible format. - -``Int4Tensor`` weights are wrapped as ``ExportableInt4Tensor`` so they export to -``dequantize_int4_tensor -> linear/embedding`` (matched by MLX's Int4 handlers). -``IntxUnpackedToInt8Tensor`` (e.g. int8 / Q6_K) already exports to -``dequantize_affine -> linear`` and is assigned directly. Coarse/per-axis group -sizes are regrouped to an MLX-legal size (16/32/64/128) inside the MLX pattern -handlers at export time (``regroup_affine_scales``), so no pack-time regroup is -needed here. - -The backend-agnostic ``pack_model`` dispatcher lives in ``pack.py``. -""" - -import json - -import torch -import torch.nn as nn - -from .pack import ModulePackerFn, pack_model # noqa: F401 - - -# --------------------------------------------------------------------------- -# Per-module packer - - -def pack_for_mlx(module: nn.Module, weights: dict[str, torch.Tensor]) -> None: - """Pack a quantized weight for MLX. - - ``Int4Tensor`` is wrapped as ``ExportableInt4Tensor`` (exports to - ``dequantize_int4_tensor → linear/embedding``). ``IntxUnpackedToInt8Tensor`` - is assigned directly; coarse/per-axis group sizes are regrouped to an - MLX-legal size in the MLX pattern handlers at export time. - """ - from executorch.extension.llm.export.int4 import ExportableInt4Tensor - from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor - - w = weights["weight"] - if isinstance(w, Int4Tensor): - # Int4 group is MLX-native (32); wrap so it exports to - # dequantize_int4_tensor -> linear/embedding. - w = ExportableInt4Tensor.from_int4_tensor(w) - module.weight = nn.Parameter(w, requires_grad=False) - - -DEFAULT_MLX_PACKERS: dict[type, ModulePackerFn] = { - nn.Linear: pack_for_mlx, - nn.Embedding: pack_for_mlx, -} - - -# --------------------------------------------------------------------------- -# Load + pack (I/O wrapper) - - -def load_and_pack_for_mlx( - path: str, - model: nn.Module, - packers: dict[type, ModulePackerFn] | None = None, -) -> None: - """Load a quantized safetensors file and pack for MLX. - - Streams one weight at a time via torchao's safetensors support. - """ - from safetensors import safe_open - from torchao.prototype.safetensors.safetensors_support import ( - unflatten_tensor_state_dict, - ) - - from .pack import pack_one - - _packers = packers or DEFAULT_MLX_PACKERS - with safe_open(path, framework="pt", device="cpu") as f: - metadata = f.metadata() - all_keys = list(f.keys()) - tensor_names = json.loads(metadata.get("tensor_names", "[]")) - - for name in tensor_names: - parts = name.rsplit(".", 1) - module_fqn = parts[0] if len(parts) > 1 else "" - weight_name = parts[-1] - prefix = ( - f"{module_fqn}._{weight_name}_" if module_fqn else f"_{weight_name}_" - ) - partial = {} - for key in all_keys: - if key.startswith(prefix) or key == name: - partial[key] = f.get_tensor(key) - result, _ = unflatten_tensor_state_dict(partial, metadata) - for fqn, value in result.items(): - pack_one(model, fqn, value, _packers) - - for fqn, p in model.named_parameters(): - if p.device.type == "meta": - raise RuntimeError( - f"Weight '{fqn}' not found in checkpoint " - f"(model/checkpoint version mismatch?)" - ) diff --git a/examples/models/gemma4_31b/quant/tests/test_pack_mlx.py b/examples/models/gemma4_31b/quant/tests/test_pack_mlx.py deleted file mode 100644 index 2861a7ffaec..00000000000 --- a/examples/models/gemma4_31b/quant/tests/test_pack_mlx.py +++ /dev/null @@ -1,173 +0,0 @@ -# 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. - -"""Unit tests for quant/pack_mlx.py. No CUDA or MLX hardware required.""" - -import unittest - -import torch -import torch.nn as nn - -from executorch.examples.models.gemma4_31b.quant.pack import pack_model -from executorch.examples.models.gemma4_31b.quant.pack_mlx import ( - DEFAULT_MLX_PACKERS, - pack_for_mlx, -) -from executorch.examples.models.gemma4_31b.quant.quantize import ( - dequantize_weight, - quantize_weight, -) -from executorch.examples.models.gemma4_31b.quant.recipe import QuantConfig - - -class TestPackLinearForMlx(unittest.TestCase): - def test_int4_wraps_exportable(self): - from executorch.extension.llm.export.int4 import ExportableInt4Tensor - - module = nn.Linear(128, 64, bias=False) - config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") - w = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), config) - pack_for_mlx(module, {"weight": w}) - - self.assertIsInstance(module.weight.data, ExportableInt4Tensor) - self.assertEqual(module.weight.shape, torch.Size([64, 128])) - self.assertFalse(module.weight.requires_grad) - - def test_int8_passes_through(self): - from torchao.quantization import IntxUnpackedToInt8Tensor - - module = nn.Linear(128, 64, bias=False) - config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") - w = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), config) - self.assertIsInstance(w, IntxUnpackedToInt8Tensor) - pack_for_mlx(module, {"weight": w}) - - self.assertIsInstance(module.weight.data, IntxUnpackedToInt8Tensor) - self.assertEqual(module.weight.shape, torch.Size([64, 128])) - - def test_int8_coarse_passes_through(self): - """Linear with a coarse group_size passes through unchanged. - - Regrouping to an MLX-legal group_size now happens in the MLX pattern - handlers at export time, so the packer leaves block_size untouched. - """ - torch.manual_seed(0) - weight = torch.randn(64, 256, dtype=torch.bfloat16) - config = QuantConfig(bits=8, group_size=256, symmetric=True, method="min_max") - w = quantize_weight(weight, config) - before = dequantize_weight(w, torch.float32) - - module = nn.Linear(256, 64, bias=False) - pack_for_mlx(module, {"weight": w}) - - self.assertEqual(module.weight.data.block_size, (1, 256)) - after = dequantize_weight(module.weight.data, torch.float32) - self.assertTrue( - torch.allclose(before, after, atol=1e-5), - f"max diff: {(before - after).abs().max():.6g}", - ) - - -class TestPackLinearGroupSize16(unittest.TestCase): - """Packing group_size=16 weights (GGUF Q6_K) preserves semantics.""" - - def _make_gs16_tensor(self, N=64, K=128): - from torchao.quantization import IntxUnpackedToInt8Tensor - - return IntxUnpackedToInt8Tensor( - qdata=torch.randint(-32, 31, (N, K), dtype=torch.int8), - scale=torch.randn(N, K // 16, dtype=torch.bfloat16), - zero_point=torch.zeros(N, K // 16, dtype=torch.int8), - target_dtype=torch.int8, - block_size=(1, 16), - dtype=torch.bfloat16, - activation_quantization=None, - ) - - def test_dequant_preserves_values(self): - """Packing preserves the dequantized weight values.""" - w = self._make_gs16_tensor(64, 128) - before = dequantize_weight(w, torch.float32) - - module = nn.Linear(128, 64, bias=False) - pack_for_mlx(module, {"weight": w}) - after = dequantize_weight(module.weight.data, torch.float32) - - self.assertTrue( - torch.allclose(before, after, atol=1e-5), - f"max diff: {(before - after).abs().max():.6g}", - ) - - def test_forward_produces_valid_output(self): - """Packed gs=16 weight produces finite output in a linear forward.""" - w = self._make_gs16_tensor(64, 128) - module = nn.Linear(128, 64, bias=False) - pack_for_mlx(module, {"weight": w}) - - x = torch.randn(1, 128, dtype=torch.bfloat16) - out = torch.nn.functional.linear(x, module.weight.data.dequantize()) - self.assertEqual(out.shape, torch.Size([1, 64])) - self.assertFalse(torch.isnan(out).any()) - - -class TestPackEmbeddingForMlx(unittest.TestCase): - def test_compatible_passes_through(self): - module = nn.Embedding(100, 64) - config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") - w = quantize_weight(torch.randn(100, 64, dtype=torch.bfloat16), config) - pack_for_mlx(module, {"weight": w}) - self.assertEqual(module.weight.shape, torch.Size([100, 64])) - - def test_per_axis_passes_through(self): - module = nn.Embedding(50, 256) - config = QuantConfig(bits=8, group_size=256, symmetric=True, method="min_max") - w = quantize_weight(torch.randn(50, 256, dtype=torch.bfloat16), config) - pack_for_mlx(module, {"weight": w}) - self.assertEqual(module.weight.shape, torch.Size([50, 256])) - # Regrouping happens in the MLX handlers at export time, not at pack time. - self.assertEqual(module.weight.data.block_size, (1, 256)) - - def test_int4_wraps_exportable(self): - from executorch.extension.llm.export.int4 import ExportableInt4Tensor - - module = nn.Embedding(100, 64) - config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") - w = quantize_weight(torch.randn(100, 64, dtype=torch.bfloat16), config) - pack_for_mlx(module, {"weight": w}) - self.assertIsInstance(module.weight.data, ExportableInt4Tensor) - self.assertEqual(module.weight.shape, torch.Size([100, 64])) - - -class TestPackModelMlx(unittest.TestCase): - def test_mixed_precision(self): - q4 = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") - q8 = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") - w4 = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), q4) - w8 = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), q8) - - state_dict = { - "q_proj.weight": w4, - "v_proj.weight": w8, - "norm.weight": torch.randn(64, dtype=torch.bfloat16), - } - - with torch.device("meta"): - model = nn.ModuleDict( - { - "q_proj": nn.Linear(128, 64, bias=False), - "v_proj": nn.Linear(128, 64, bias=False), - "norm": nn.LayerNorm(64, bias=False), - } - ) - pack_model(model, state_dict, DEFAULT_MLX_PACKERS) - - self.assertEqual(model.q_proj.weight.shape, torch.Size([64, 128])) - self.assertEqual(model.v_proj.weight.shape, torch.Size([64, 128])) - self.assertEqual(model.norm.weight.shape, torch.Size([64])) - - -if __name__ == "__main__": - unittest.main() diff --git a/examples/models/gemma4_31b/quantize_and_save.py b/examples/models/gemma4_31b/quantize_and_save.py index e654e12f637..b0a68ad691d 100644 --- a/examples/models/gemma4_31b/quantize_and_save.py +++ b/examples/models/gemma4_31b/quantize_and_save.py @@ -6,9 +6,12 @@ """Quantize Gemma 4 31B-IT and save as a quantized checkpoint. -Produces a safetensors file containing torchao tensor subclasses -(``Int4Tensor``, ``IntxUnpackedToInt8Tensor``) that can be loaded and -packed for any backend via ``load_and_pack_for_cuda`` or ``pack_model``. +Streams the HF checkpoint shard-by-shard via ``iter_checkpoint`` and quantizes +each weight with ``quantize_stream`` -- the full model is never instantiated +(model-free). Produces a safetensors file containing torchao tensor subclasses +(``Int4Tensor``, ``IntxUnpackedToInt8Tensor``) that can be loaded and converted +for any backend via ``load_checkpoint`` (from disk) or ``assign_state_dict`` +(in memory). The default recipe runs on CPU. The sensitive recipe requires CUDA for HQQ asymmetric quantization. @@ -24,12 +27,13 @@ import os import shutil -import torch.nn as nn +import torch -from executorch.examples.models.gemma4_31b.model import Gemma4_31B -from executorch.examples.models.gemma4_31b.quant import ( +from executorch.extension.llm.export.load import iter_checkpoint +from executorch.extension.llm.export.quant import ( + identity, QuantConfig, - quantize_model, + quantize_stream, QuantRecipe, QuantRule, ) @@ -107,15 +111,26 @@ def main() -> None: recipe = _RECIPES[args.quant_recipe] - print("Loading checkpoint (lazy, shard-by-shard)...") - model, _ = Gemma4_31B.from_hf_checkpoint(args.model_dir) - - if model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr(): - print("Untying embed_tokens / lm_head...") - model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) - - print(f"Quantizing with recipe '{args.quant_recipe}'...") - state_dict = quantize_model(model, recipe, verbose=True) + print( + f"Streaming + quantizing checkpoint (model-free) with '{args.quant_recipe}'..." + ) + from executorch.examples.models.gemma4_31b.model import _hf_to_model_key + + # Stream raw HF weights, remap to model FQNs, fill the tied lm_head from the + # embedding (clone=True so it quantizes independently), and quantize per + # recipe. Only one logical weight is materialized at a time -- the full model + # is never instantiated. + pairs = iter_checkpoint( + args.model_dir, + key_map=_hf_to_model_key, + convert=identity, + tie_map={"embed_tokens.weight": ("lm_head.weight", True)}, + ) + state_dict = {} + for fqn, value in quantize_stream(pairs, recipe, dtype=torch.bfloat16): + state_dict[fqn] = value + print(f" Quantized: {fqn}", end="\r") + print() os.makedirs(args.output, exist_ok=True) safetensors_path = os.path.join(args.output, "model.safetensors") diff --git a/examples/models/gemma4_31b/tests/BUCK b/examples/models/gemma4_31b/tests/BUCK new file mode 100644 index 00000000000..a4411d89925 --- /dev/null +++ b/examples/models/gemma4_31b/tests/BUCK @@ -0,0 +1,18 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +# 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. + +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target(_kind = runtime.python_test, + name = "test_recipe", + srcs = ["test_recipe.py"], + deps = [ + "//executorch/examples/models/gemma4_31b:quantize_and_save", + ], +) diff --git a/examples/models/gemma4_31b/tests/test_cuda_packers.py b/examples/models/gemma4_31b/tests/test_cuda_packers.py new file mode 100644 index 00000000000..d2d7048f687 --- /dev/null +++ b/examples/models/gemma4_31b/tests/test_cuda_packers.py @@ -0,0 +1,501 @@ +# 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. + +"""Unit tests for cuda_packers.py. + +Tests the public API contract: after packing, modules produce correct +output via F.linear / nn.Embedding at various batch sizes and configs. +""" + +import os +import tempfile +import unittest + +# Register Int4/Int8 F.linear dispatch before any test uses it +import executorch.backends.cuda.quantize_op_dispatch # noqa: F401 +import torch +import torch.nn as nn +from executorch.backends.cuda.dp4a_planar_int6_tensor import ( + _encode_int8_per_super, + CudaDp4aPlanarInt6Tensor, + pack_int6, + unpack_int6, +) +from executorch.examples.models.gemma4_31b.cuda_packers import ( + convert_quantized_tensors_for_cuda, + pack_embedding_for_cuda, + pack_linear_for_cuda, +) +from executorch.extension.llm.export.load import ( + assign_one, + assign_state_dict, + load_checkpoint, +) +from executorch.extension.llm.export.quant.convert import to_exportable +from executorch.extension.llm.export.quant.quantize import quantize_weight +from executorch.extension.llm.export.quant.recipe import QuantConfig +from safetensors.torch import save_file +from torchao.prototype.safetensors.safetensors_support import flatten_tensor_state_dict + + +def _require_cuda(tc: unittest.TestCase) -> None: + if not torch.cuda.is_available(): + tc.skipTest("CUDA required") + + +class TestPackLinearInt4(unittest.TestCase): + """pack_linear_for_cuda with INT4 weights produces correct F.linear output.""" + + def setUp(self): + _require_cuda(self) + torch.manual_seed(0) + self.weight = torch.randn(256, 1024, dtype=torch.bfloat16) + + def _pack(self, symmetric=False, group_size=32): + config = QuantConfig( + bits=4, group_size=group_size, symmetric=symmetric, method="min_max" + ) + q = to_exportable("weight", quantize_weight(self.weight, config)) + module = nn.Linear(1024, 256, bias=False) + pack_linear_for_cuda(module, {"weight": q}) + module.cuda() + return module + + def test_shape_preserved(self): + module = self._pack() + self.assertEqual(module.weight.shape, torch.Size([256, 1024])) + + def test_asymmetric_decode(self): + module = self._pack(symmetric=False) + x = torch.randn(1, 1024, dtype=torch.bfloat16, device="cuda") + ref = torch.nn.functional.linear(x, self.weight.cuda()) + out = module(x) + rel_error = (out.float() - ref.float()).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.15) + + def test_symmetric_decode(self): + module = self._pack(symmetric=True) + x = torch.randn(1, 1024, dtype=torch.bfloat16, device="cuda") + ref = torch.nn.functional.linear(x, self.weight.cuda()) + out = module(x) + rel_error = (out.float() - ref.float()).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.15) + + def test_prefill_batch(self): + module = self._pack(symmetric=False) + x = torch.randn(64, 1024, dtype=torch.bfloat16, device="cuda") + ref = torch.nn.functional.linear(x, self.weight.cuda()) + out = module(x) + rel_error = (out.float() - ref.float()).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.15) + + def test_different_group_sizes(self): + for gs in (32, 64, 128): + with self.subTest(group_size=gs): + module = self._pack(group_size=gs) + x = torch.randn(1, 1024, dtype=torch.bfloat16, device="cuda") + ref = torch.nn.functional.linear(x, self.weight.cuda()) + out = module(x) + rel_error = ( + out.float() - ref.float() + ).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.15) + + +class TestPackLinearInt8(unittest.TestCase): + """pack_linear_for_cuda with INT8 weights produces correct F.linear output.""" + + def setUp(self): + _require_cuda(self) + + def test_matmul_correct(self): + torch.manual_seed(0) + weight = torch.randn(256, 128, dtype=torch.bfloat16) + x = torch.randn(1, 128, dtype=torch.bfloat16) + ref = torch.nn.functional.linear(x.cuda(), weight.cuda()) + + config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") + q = quantize_weight(weight, config) + module = nn.Linear(128, 256, bias=False) + pack_linear_for_cuda(module, {"weight": q}) + module.cuda() + out = module(x.cuda()) + + rel_error = (out.float() - ref.float()).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.02) + + def test_unsupported_type_raises(self): + module = nn.Linear(64, 32, bias=False) + with self.assertRaises(ValueError): + pack_linear_for_cuda(module, {"weight": torch.randn(32, 64)}) + + +class TestPackLinearInt6(unittest.TestCase): + """pack_linear_for_cuda converts a native Q6_K ExportableGGUFTensor (the + gguf_loader output) into a CudaDp4aPlanarInt6Tensor. + + The pack/unpack round-trip is lossless and dequantize() == q * scale (no + CUDA required); the F.linear correctness check is CUDA-only. A genuine INT8 + IntxUnpackedToInt8Tensor is left on the int8 path. + """ + + def setUp(self): + torch.manual_seed(0) + + def _make_int6(self, N, K, gs=16): + q = torch.randint(-32, 32, (N, K), dtype=torch.int8) + scale = (torch.rand(N, K // gs) * 0.1 + 0.01).to(torch.bfloat16) + ql, qh = pack_int6(q) + # The tensor stores int8 scale *codes* + a per-256-super-block [N, K/256] + # fp16 step (scale = code * step[:, g // (256 // gs)]); encode here and + # return the effective scale so the reference dequant matches the kernel. + scale_codes, steps = _encode_int8_per_super(scale.float(), gs) + n_super = steps.shape[1] + gps = (K // gs) // n_super + step_g = steps.to(torch.bfloat16).repeat_interleave(gps, dim=1) + eff_scale = (scale_codes.to(torch.bfloat16) * step_g).to(torch.bfloat16) + t = CudaDp4aPlanarInt6Tensor( + ql, qh, scale_codes, steps, [1, gs], torch.Size([N, K]) + ) + return t, q, eff_scale + + def _make_q6k_gguf(self, N, nb=1): + """Build a synthetic q6_k ExportableGGUFTensor (see test_gguf.py). + + ``K = nb * 256``; fixed non-zero sub-scales + a small super-block scale + keep dequantized magnitudes O(1). + """ + from executorch.extension.llm.export.gguf import ( + _Q6_K_BLOCK_BYTES, + ExportableGGUFTensor, + ) + + g = torch.Generator().manual_seed(0) + blk = torch.randint( + 0, 256, (N * nb, _Q6_K_BLOCK_BYTES), dtype=torch.uint8, generator=g + ) + blk[:, 192:208] = 0x10 # fixed non-zero int8 sub-scales + blk[:, 208:210] = torch.tensor([0.01], dtype=torch.float16).view( + torch.uint8 + ) # super-block scale d + raw = blk.reshape(N, nb * _Q6_K_BLOCK_BYTES) + return ExportableGGUFTensor.from_raw(raw, "q6_k") + + def test_pack_unpack_roundtrip(self): + q = torch.randint(-32, 32, (64, 128), dtype=torch.int8) + ql, qh = pack_int6(q) + self.assertEqual(tuple(ql.shape), (64, 64)) # [N, K/2] + self.assertEqual(tuple(qh.shape), (64, 32)) # [N, K/4] + q_rt = unpack_int6(ql, qh, 64, 128).to(torch.int8) + self.assertTrue(torch.equal(q_rt, q)) + + def test_dequantize_equals_q_scale(self): + t, q, scale = self._make_int6(32, 256, gs=16) + ref = q.to(torch.bfloat16) * scale.to(torch.bfloat16).repeat_interleave( + 16, dim=-1 + ) + self.assertTrue(torch.equal(t.dequantize(), ref)) + + def test_pack_linear_converts_q6k(self): + gt = self._make_q6k_gguf(32, nb=1) # (32, 256) + with torch.device("meta"): + module = nn.Linear(256, 32, bias=False) + pack_linear_for_cuda(module, {"weight": gt}) + self.assertIsInstance(module.weight.data, CudaDp4aPlanarInt6Tensor) + self.assertEqual(module.weight.shape, torch.Size([32, 256])) + + def test_pack_linear_real_int8_passthrough(self): + """A genuine INT8 weight (wide groups, full range) is NOT repacked.""" + from torchao.quantization import IntxUnpackedToInt8Tensor + + q = torch.randint(-128, 128, (32, 128), dtype=torch.int8) + scale = (torch.rand(32, 128 // 32) * 0.1 + 0.01).to(torch.bfloat16) + zero = torch.zeros(32, 128 // 32, dtype=torch.int8) + t = IntxUnpackedToInt8Tensor( + qdata=q, + scale=scale, + zero_point=zero, + target_dtype=torch.int8, + block_size=(1, 32), + dtype=torch.bfloat16, + activation_quantization=None, + ) + with torch.device("meta"): + module = nn.Linear(128, 32, bias=False) + pack_linear_for_cuda(module, {"weight": t}) + self.assertIsInstance(module.weight.data, IntxUnpackedToInt8Tensor) + + def test_matmul_correct(self): + _require_cuda(self) + gt = self._make_q6k_gguf(256, nb=1) # (256, 256) + intx = gt.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16) + q, scale = intx.qdata, intx.scale + module = nn.Linear(256, 256, bias=False) + pack_linear_for_cuda(module, {"weight": gt}) + self.assertIsInstance(module.weight.data, CudaDp4aPlanarInt6Tensor) + module.cuda() + x = torch.randn(1, 256, dtype=torch.bfloat16, device="cuda") + w_ref = ( + q.to(torch.bfloat16) + * scale.to(torch.bfloat16).repeat_interleave(16, dim=-1) + ).cuda() + ref = torch.nn.functional.linear(x, w_ref) + out = module(x) + rel_error = (out.float() - ref.float()).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.02) + + def test_e2e_q6k_export_lower_decode(self): + """Q6_K -> pack -> export + CUDA lower -> run, vs the Q6_K dequant reference. + + Builds a synthetic Q6_K ExportableGGUFTensor, packs it into a + CudaDp4aPlanarInt6Tensor, exports a decode-shaped (M=1) nn.Linear, and + asserts: + * the exported graph captured ``executorch_cuda.int6_plain_mm`` (the + decode custom op chosen for M<=4), + * lowering through the CUDA backend produces an ``executorch_call_delegate``, + * running the exported graph matches the Q6_K dequant reference. + + The lowered .pte is not executed here (that needs the built C-shim + runtime); the eager exported graph already exercises the int6 decode op + through its registered CUDA impl. + """ + _require_cuda(self) + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower + from torch.export import export + + gt = self._make_q6k_gguf(64, nb=1) # (64, 256) + # Reference: the shared Q6_K int8 decode, dequantized (w = q * scale). + intx = gt.to_intx_unpacked_to_int8_tensor(scale_dtype=torch.bfloat16) + w_ref = ( + intx.qdata.to(torch.bfloat16) + * intx.scale.to(torch.bfloat16).repeat_interleave(16, dim=-1) + ).cuda() + + module = nn.Linear(256, 64, bias=False) + pack_linear_for_cuda(module, {"weight": gt}) + self.assertIsInstance(module.weight.data, CudaDp4aPlanarInt6Tensor) + module = module.cuda().eval() + + x = torch.randn(1, 256, dtype=torch.bfloat16, device="cuda") # decode M=1 + with torch.no_grad(): + ep = export(module, (x,), strict=True) + + # The decode (M<=4) path must capture the int6 decode custom op. + targets = [str(n.target) for n in ep.graph.nodes if n.op == "call_function"] + self.assertTrue( + any("int6_plain_mm" in t for t in targets), + f"int6_plain_mm not found in exported graph: {targets}", + ) + + # Run the exported graph and compare against the Q6_K dequant reference. + with torch.no_grad(): + out = ep.module()(x) + ref = torch.nn.functional.linear(x, w_ref) + rel_error = (out.float() - ref.float()).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.02) + + # Lower through the CUDA backend: the int6 weight + decode op must land in + # an executorch_call_delegate. + lowered = to_edge_transform_and_lower( + ep, + partitioner=[ + CudaPartitioner( + [CudaBackend.generate_method_name_compile_spec("forward")] + ) + ], + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + lowered_ep = lowered.exported_program() + self.assertTrue( + any( + n.op == "call_function" and "executorch_call_delegate" in str(n.target) + for n in lowered_ep.graph.nodes + ), + "CUDA lowering produced no delegate call", + ) + + +class TestPackEmbedding(unittest.TestCase): + """pack_embedding_for_cuda with INT8 per-axis weights.""" + + def setUp(self): + _require_cuda(self) + + def test_gather_correct(self): + torch.manual_seed(0) + weight = torch.randn(1000, 64, dtype=torch.bfloat16) + ids = torch.tensor([0, 1, 42, 500, 999]) + ref = weight[ids] + + config = QuantConfig(bits=8, group_size=64, symmetric=True, method="min_max") + q = quantize_weight(weight, config) + module = nn.Embedding(1000, 64) + pack_embedding_for_cuda(module, {"weight": q}) + module.cuda() + out = module(ids.cuda()) + + rel_error = ( + out.cpu().float() - ref.float() + ).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.02) + + def test_rejects_4bit(self): + config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") + q = quantize_weight(torch.randn(100, 64, dtype=torch.bfloat16), config) + module = nn.Embedding(100, 64) + with self.assertRaises(ValueError): + pack_embedding_for_cuda(module, {"weight": q}) + + +class TestPackModel(unittest.TestCase): + """assign_state_dict + convert_quantized_tensors_for_cuda over a whole model.""" + + def setUp(self): + _require_cuda(self) + + def test_mixed_precision(self): + torch.manual_seed(0) + w4 = torch.randn(64, 256, dtype=torch.bfloat16) + w8 = torch.randn(64, 256, dtype=torch.bfloat16) + q4 = to_exportable( + "q_proj.weight", + quantize_weight( + w4, + QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max"), + ), + ) + q8 = quantize_weight( + w8, + QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max"), + ) + with torch.device("meta"): + model = nn.ModuleDict( + { + "q_proj": nn.Linear(256, 64, bias=False), + "v_proj": nn.Linear(256, 64, bias=False), + } + ) + assign_state_dict(model, {"q_proj.weight": q4, "v_proj.weight": q8}) + convert_quantized_tensors_for_cuda(model) + model.cuda() + x = torch.randn(1, 256, dtype=torch.bfloat16, device="cuda") + + ref4 = torch.nn.functional.linear(x, w4.cuda()) + out4 = model.q_proj(x) + self.assertLess( + (out4.float() - ref4.float()).abs().mean().item() + / ref4.float().abs().mean().item(), + 0.15, + ) + + ref8 = torch.nn.functional.linear(x, w8.cuda()) + out8 = model.v_proj(x) + self.assertLess( + (out8.float() - ref8.float()).abs().mean().item() + / ref8.float().abs().mean().item(), + 0.02, + ) + + def test_load_checkpoint_from_disk(self): + torch.manual_seed(0) + weight = torch.randn(64, 256, dtype=torch.bfloat16) + config = QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max") + q = quantize_weight(weight, config) + + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "m.safetensors") + state = { + "proj.weight": q, + "norm.weight": torch.randn(64, dtype=torch.bfloat16), + } + td, md = flatten_tensor_state_dict(state) + save_file(td, path, metadata=md) + + with torch.device("meta"): + model = nn.ModuleDict( + { + "proj": nn.Linear(256, 64, bias=False), + "norm": nn.LayerNorm(64, bias=False), + } + ) + load_checkpoint(path, model) + convert_quantized_tensors_for_cuda(model) + + self.assertEqual(model.proj.weight.shape, torch.Size([64, 256])) + self.assertEqual(model.norm.weight.shape, torch.Size([64])) + + model.proj.cuda() + x = torch.randn(1, 256, dtype=torch.bfloat16, device="cuda") + ref = torch.nn.functional.linear(x, weight.cuda()) + out = model.proj(x) + rel_error = (out.float() - ref.float()).abs().mean() / ref.float().abs().mean() + self.assertLess(rel_error.item(), 0.15) + + +class TestAssignAndPass(unittest.TestCase): + """assign_one / the CUDA pass on CPU (no matmul; no CUDA required).""" + + def test_assign_one_quantized_then_pass(self): + from executorch.backends.cuda.coalesced_int4_tensor import ( + CudaCoalescedInt4Tensor, + ) + + config = QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max") + q = quantize_weight(torch.randn(64, 256, dtype=torch.bfloat16), config) + with torch.device("meta"): + model = nn.ModuleDict({"proj": nn.Linear(256, 64, bias=False)}) + assign_one(model, "proj.weight", to_exportable("proj.weight", q)) + self.assertNotEqual(model.proj.weight.device.type, "meta") + convert_quantized_tensors_for_cuda(model) + self.assertIsInstance(model.proj.weight.data, CudaCoalescedInt4Tensor) + + def test_assign_one_plain_tensor(self): + with torch.device("meta"): + model = nn.ModuleDict({"norm": nn.LayerNorm(64, bias=False)}) + assign_one(model, "norm.weight", torch.randn(64, dtype=torch.bfloat16)) + self.assertEqual(model.norm.weight.dtype, torch.bfloat16) + + +class TestPackErrorPaths(unittest.TestCase): + + def test_unsupported_module_type(self): + class CustomModule(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.randn(32, 64)) + + config = QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max") + q = quantize_weight(torch.randn(32, 64, dtype=torch.bfloat16), config) + + with torch.device("meta"): + model = nn.ModuleDict({"custom": CustomModule()}) + assign_state_dict(model, {"custom.weight": q}) + with self.assertRaises(ValueError) as ctx: + convert_quantized_tensors_for_cuda(model) + self.assertIn("CustomModule", str(ctx.exception)) + + def test_missing_weight_detected(self): + config = QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max") + q = quantize_weight(torch.randn(32, 256, dtype=torch.bfloat16), config) + + with torch.device("meta"): + model = nn.ModuleDict( + { + "a": nn.Linear(256, 32, bias=False), + "b": nn.Linear(256, 32, bias=False), + } + ) + with self.assertRaises(RuntimeError) as ctx: + assign_state_dict(model, {"a.weight": q}) + self.assertIn("b.weight", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/models/gemma4_31b/tests/test_cuda_pipeline.py b/examples/models/gemma4_31b/tests/test_cuda_pipeline.py index c346c1d2f82..41f1ea1e01e 100644 --- a/examples/models/gemma4_31b/tests/test_cuda_pipeline.py +++ b/examples/models/gemma4_31b/tests/test_cuda_pipeline.py @@ -23,18 +23,16 @@ import executorch.backends.cuda.quantize_op_dispatch # noqa: F401 import torch import torch.nn as nn +from executorch.examples.models.gemma4_31b.cuda_packers import ( + convert_quantized_tensors_for_cuda, +) from executorch.examples.models.gemma4_31b.export import ( export_and_lower, + load_gguf_model, load_prequantized_model, ) -from executorch.examples.models.gemma4_31b.gguf_loader import load_gguf_model from executorch.examples.models.gemma4_31b.inference import _move_to_cuda, generate from executorch.examples.models.gemma4_31b.model import Gemma4_31B -from executorch.examples.models.gemma4_31b.quant import ( - DEFAULT_CUDA_PACKERS, - pack_model, - quantize_model, -) from executorch.examples.models.gemma4_31b.tests.test_pipeline import ( build_gguf_checkpoint, build_hf_checkpoint, @@ -44,6 +42,8 @@ save_checkpoint, TINY_CONFIG, ) +from executorch.extension.llm.export.load import assign_state_dict +from executorch.extension.llm.export.quant import quantize_model def _require_cuda(testcase: unittest.TestCase) -> None: @@ -162,7 +162,8 @@ def test_export_from_hf_checkpoint(self): with torch.device("meta"): model = Gemma4_31B(config) - pack_model(model, state_dict, DEFAULT_CUDA_PACKERS) + assign_state_dict(model, state_dict) + convert_quantized_tensors_for_cuda(model) model.eval() export_and_lower(model, config, out_dir) diff --git a/examples/models/gemma4_31b/tests/test_mlx_pipeline.py b/examples/models/gemma4_31b/tests/test_mlx_pipeline.py index 23200a71462..3b86ce24e14 100644 --- a/examples/models/gemma4_31b/tests/test_mlx_pipeline.py +++ b/examples/models/gemma4_31b/tests/test_mlx_pipeline.py @@ -21,14 +21,6 @@ import torch import torch.nn as nn from executorch.examples.models.gemma4_31b.model import Gemma4_31B -from executorch.examples.models.gemma4_31b.quant import ( - DEFAULT_MLX_PACKERS, - pack_model, - QuantConfig, - quantize_model, - QuantRecipe, - QuantRule, -) from executorch.examples.models.gemma4_31b.tests.test_pipeline import ( build_gguf_checkpoint, build_random_tiny_model, @@ -37,6 +29,13 @@ save_checkpoint, TINY_CONFIG, ) +from executorch.extension.llm.export.load import assign_state_dict +from executorch.extension.llm.export.quant import ( + QuantConfig, + quantize_model, + QuantRecipe, + QuantRule, +) _INT4 = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") _INT8 = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") @@ -67,7 +66,7 @@ def test_pack_for_mlx(self): with torch.device("meta"): model = Gemma4_31B(TINY_CONFIG) model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) - pack_model(model, state_dict, DEFAULT_MLX_PACKERS) + assign_state_dict(model, state_dict) for fqn, p in model.named_parameters(): self.assertNotEqual(p.device.type, "meta", f"Weight '{fqn}' still on meta") @@ -81,7 +80,7 @@ def test_forward_after_pack(self): with torch.device("meta"): model = Gemma4_31B(TINY_CONFIG) model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) - pack_model(model, state_dict, DEFAULT_MLX_PACKERS) + assign_state_dict(model, state_dict) model.eval() from executorch.examples.models.gemma4_31b.model import ( @@ -108,7 +107,7 @@ def test_multi_token_forward(self): with torch.device("meta"): model = Gemma4_31B(TINY_CONFIG) model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) - pack_model(model, state_dict, DEFAULT_MLX_PACKERS) + assign_state_dict(model, state_dict) model.eval() from executorch.examples.models.gemma4_31b.model import ( @@ -137,7 +136,7 @@ def test_source_transforms_forward(self): with torch.device("meta"): model = Gemma4_31B(TINY_CONFIG) model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) - pack_model(model, state_dict, DEFAULT_MLX_PACKERS) + assign_state_dict(model, state_dict) model.eval() from executorch.examples.models.gemma4_31b.mlx_source_transformations import ( @@ -192,7 +191,7 @@ def test_source_transforms_use_mlx_ops(self): with torch.device("meta"): model = Gemma4_31B(TINY_CONFIG) model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) - pack_model(model, state_dict, DEFAULT_MLX_PACKERS) + assign_state_dict(model, state_dict) model.eval() mlx_source_transformations(model, dtype=torch.bfloat16) @@ -308,23 +307,24 @@ def test_load_gguf_model_mlx_backend(self): except ModuleNotFoundError: self.skipTest("gguf package not installed") - from executorch.examples.models.gemma4_31b.gguf_loader import load_gguf_model + from executorch.examples.models.gemma4_31b.export import load_gguf_model # Will fail on missing file, but NOT on "Unsupported backend". with self.assertRaisesRegex((FileNotFoundError, OSError, RuntimeError), ".*"): load_gguf_model("/nonexistent.gguf", backend="mlx") def test_mlx_backend_rejects_unknown(self): - from executorch.examples.models.gemma4_31b.gguf_loader import load_gguf_model + from executorch.examples.models.gemma4_31b.export import load_gguf_model with self.assertRaisesRegex(ValueError, "Unsupported backend"): load_gguf_model("/nonexistent.gguf", backend="tpu") def test_gs16_packing_preserves_values(self): """Q6_K-like weight (gs=16) preserves dequantized values after packing.""" - from executorch.examples.models.gemma4_31b.quant.pack_mlx import pack_for_mlx - from executorch.examples.models.gemma4_31b.quant.quantize import ( + from executorch.extension.llm.export.load import assign_one + from executorch.extension.llm.export.quant import ( dequantize_weight, + to_exportable, ) from torchao.quantization import IntxUnpackedToInt8Tensor @@ -340,7 +340,7 @@ def test_gs16_packing_preserves_values(self): before = dequantize_weight(w, torch.float32) module = nn.Linear(128, 64, bias=False) - pack_for_mlx(module, {"weight": w}) + assign_one(module, "weight", to_exportable("weight", w)) after = dequantize_weight(module.weight.data, torch.float32) self.assertTrue( @@ -350,10 +350,8 @@ def test_gs16_packing_preserves_values(self): def test_embedding_packing_preserves_values(self): """MLX embedding packing preserves dequantized weight values.""" - from executorch.examples.models.gemma4_31b.quant.pack_mlx import pack_for_mlx - from executorch.examples.models.gemma4_31b.quant.quantize import ( - dequantize_weight, - ) + from executorch.extension.llm.export.load import assign_one + from executorch.extension.llm.export.quant import dequantize_weight, to_default from torchao.quantization import IntxUnpackedToInt8Tensor w = IntxUnpackedToInt8Tensor( @@ -368,7 +366,7 @@ def test_embedding_packing_preserves_values(self): before = dequantize_weight(w, torch.float32) module = nn.Embedding(256, 128) - pack_for_mlx(module, {"weight": w}) + assign_one(module, "weight", to_default("weight", w)) after = dequantize_weight(module.weight.data, torch.float32) self.assertTrue( @@ -550,7 +548,7 @@ def setUp(self): self.skipTest("gguf package required") def _load(self, tmp): - from executorch.examples.models.gemma4_31b.gguf_loader import load_gguf_model + from executorch.examples.models.gemma4_31b.export import load_gguf_model path = os.path.join(tmp, "tiny.gguf") build_gguf_checkpoint(path) diff --git a/examples/models/gemma4_31b/tests/test_pipeline.py b/examples/models/gemma4_31b/tests/test_pipeline.py index 700e323650f..0a96a8c4438 100644 --- a/examples/models/gemma4_31b/tests/test_pipeline.py +++ b/examples/models/gemma4_31b/tests/test_pipeline.py @@ -27,11 +27,13 @@ Gemma4_31BConfig, RingKVCache, ) -from executorch.examples.models.gemma4_31b.quant import ( +from executorch.extension.llm.export.quant import ( QuantConfig, quantize_model, + quantize_stream, QuantRecipe, QuantRule, + to_default, ) from safetensors import safe_open from safetensors.torch import save_file @@ -143,7 +145,15 @@ def build_random_tiny_model() -> Gemma4_31B: def save_checkpoint(output_dir: str): model = build_random_tiny_model() model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) - state_dict = quantize_model(model, DEFAULT_RECIPE) + # On-disk producer: quantize_stream yields the torchao-native serialization + # form (Int4Tensor) that flatten_tensor_state_dict can write. quantize_model + # would yield the in-memory ExportableInt4Tensor, which torchao can't flatten. + params = ((fqn, p.data) for fqn, p in model.named_parameters()) + state_dict = dict(quantize_stream(params, DEFAULT_RECIPE)) + persistent_keys = set(model.state_dict().keys()) + for fqn, buf in model.named_buffers(): + if fqn in persistent_keys and fqn not in state_dict: + state_dict[fqn] = buf.data os.makedirs(output_dir, exist_ok=True) td, md = flatten_tensor_state_dict(state_dict) save_file(td, os.path.join(output_dir, "model.safetensors"), metadata=md) @@ -263,7 +273,11 @@ def test_roundtrip_preserves_weights(self): model = build_random_tiny_model() model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) - state_dict = quantize_model(model, DEFAULT_RECIPE) + # Save form == quantize_stream (torchao-native Int4Tensor), which is what + # quantize_and_save writes to disk; quantize_model's Exportable form is + # in-memory only and would need torchao registration to flatten. + params = ((fqn, p.data) for fqn, p in model.named_parameters()) + state_dict = dict(quantize_stream(params, DEFAULT_RECIPE, dtype=torch.bfloat16)) with tempfile.TemporaryDirectory() as tmpdir: path = os.path.join(tmpdir, "model.safetensors") @@ -301,6 +315,93 @@ def test_embedding_quantized_as_int8(self): state_dict["embed_tokens.weight"], IntxUnpackedToInt8Tensor ) + def test_model_free_save_matches_quantize_model(self): + """Model-free save (iter_checkpoint + tie_map + quantize_stream) reconstructs + the model-based quantize_model state dict. + + quantize_stream yields the serialization form (torchao-native Int4Tensor); + quantize_model yields the in-memory Exportable form. They match once the + streamed output is passed through the same ``to_default`` convert. + """ + from executorch.examples.models.gemma4_31b.model import _hf_to_model_key + from executorch.extension.llm.export.load import iter_checkpoint + from executorch.extension.llm.export.quant import identity + + # Reference: instantiate the model, untie lm_head, quantize_model. + model = build_random_tiny_model() + model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) + ref = quantize_model(model, DEFAULT_RECIPE) + + # Model-free: stream a tied HF checkpoint (same seed-42 weights), remap + # names, fan lm_head out of the embedding (clone=True), quantize per recipe. + with tempfile.TemporaryDirectory() as tmpdir: + build_hf_checkpoint(tmpdir) + pairs = iter_checkpoint( + tmpdir, + key_map=_hf_to_model_key, + convert=identity, + tie_map={"embed_tokens.weight": ("lm_head.weight", True)}, + ) + streamed = { + fqn: to_default(fqn, v) + for fqn, v in quantize_stream( + pairs, DEFAULT_RECIPE, dtype=torch.bfloat16 + ) + } + + self.assertEqual(set(streamed), set(ref)) + for fqn, r in ref.items(): + got = streamed[fqn] + self.assertIs(type(got), type(r), fqn) + data_names = getattr(r, "tensor_data_names", None) + if data_names: # quantized subclass -> compare payload bit-exactly + for n in data_names: + self.assertTrue(torch.equal(getattr(got, n), getattr(r, n)), fqn) + else: + self.assertTrue(torch.equal(got, r), fqn) + + def test_load_casts_int4_to_fp16(self): + """Loading a bf16-quantized checkpoint with dtype=fp16 re-stamps int4. + + Regression for the convert-then-cast order in iter_checkpoint: a raw + torchao Int4Tensor (whose own .to ignores dtype) is first wrapped as an + ExportableInt4Tensor by convert, then re-stamped to fp16 by maybe_cast. + """ + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + from executorch.extension.llm.export.load import iter_checkpoint + from executorch.extension.llm.export.quant import to_default + from torchao.quantization import IntxUnpackedToInt8Tensor + from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor + + model = build_random_tiny_model() + model.lm_head.weight = nn.Parameter(model.embed_tokens.weight.clone()) + # Save the torchao-native form (quantize_stream) -- the real on-disk form. + params = ((fqn, p.data) for fqn, p in model.named_parameters()) + state_dict = dict(quantize_stream(params, DEFAULT_RECIPE, dtype=torch.bfloat16)) + int4_fqns = {f for f, v in state_dict.items() if isinstance(v, Int4Tensor)} + int8_fqns = { + f for f, v in state_dict.items() if isinstance(v, IntxUnpackedToInt8Tensor) + } + self.assertTrue(int4_fqns and int8_fqns) + + with tempfile.TemporaryDirectory() as tmpdir: + td, md = flatten_tensor_state_dict(state_dict) + save_file(td, os.path.join(tmpdir, "model.safetensors"), metadata=md) + loaded = dict( + iter_checkpoint(tmpdir, convert=to_default, dtype=torch.float16) + ) + + for fqn in int4_fqns: + got = loaded[fqn] + self.assertIsInstance(got, ExportableInt4Tensor, fqn) + self.assertEqual(got.orig_dtype, torch.float16, fqn) + self.assertEqual(got.scale.dtype, torch.float16, fqn) + self.assertEqual(got.dequantize().dtype, torch.float16, fqn) + for fqn in int8_fqns: + got = loaded[fqn] + self.assertIsInstance(got, IntxUnpackedToInt8Tensor, fqn) + self.assertEqual(got.scale.dtype, torch.float16, fqn) + class TestRingKVCache(unittest.TestCase): """Unit tests for the ring-buffer KV cache (CPU, no model needed).""" diff --git a/examples/models/gemma4_31b/tests/test_recipe.py b/examples/models/gemma4_31b/tests/test_recipe.py new file mode 100644 index 00000000000..6456bf53921 --- /dev/null +++ b/examples/models/gemma4_31b/tests/test_recipe.py @@ -0,0 +1,61 @@ +# 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. + +"""Regression tests for the gemma4_31b production quant recipes.""" + +import unittest + + +class TestProductionRecipes(unittest.TestCase): + """Regression tests for the production recipes in quantize_and_save.py.""" + + def test_default_recipe(self): + from executorch.examples.models.gemma4_31b.quantize_and_save import ( + GEMMA4_31B_DEFAULT_RECIPE, + ) + + r = GEMMA4_31B_DEFAULT_RECIPE + self.assertIsNone(r.get_config("layers.0.input_layernorm.weight")) + self.assertIsNone(r.get_config("layers.5.self_attn.q_norm.weight")) + self.assertIsNone(r.get_config("norm.weight")) + embed_cfg = r.get_config("embed_tokens.weight") + self.assertEqual(embed_cfg.bits, 8) + self.assertEqual(embed_cfg.group_size, 5376) + for fqn in ( + "layers.0.self_attn.q_proj.weight", + "layers.0.self_attn.v_proj.weight", + "layers.0.mlp.gate_proj.weight", + "layers.0.mlp.down_proj.weight", + "lm_head.weight", + ): + cfg = r.get_config(fqn) + self.assertEqual(cfg.bits, 4, fqn) + self.assertEqual(cfg.method, "min_max", fqn) + + def test_sensitive_recipe(self): + from executorch.examples.models.gemma4_31b.quantize_and_save import ( + GEMMA4_31B_SENSITIVE_RECIPE, + ) + + r = GEMMA4_31B_SENSITIVE_RECIPE + self.assertIsNone(r.get_config("layers.0.input_layernorm.weight")) + embed_cfg = r.get_config("embed_tokens.weight") + self.assertEqual(embed_cfg.bits, 8) + self.assertEqual(embed_cfg.group_size, 5376) + # Edge v_proj/down_proj → int8 + self.assertEqual(r.get_config("layers.0.self_attn.v_proj.weight").bits, 8) + self.assertEqual(r.get_config("layers.0.mlp.down_proj.weight").bits, 8) + self.assertEqual(r.get_config("layers.58.self_attn.v_proj.weight").bits, 8) + # Middle v_proj/down_proj → int4 + self.assertEqual(r.get_config("layers.30.self_attn.v_proj.weight").bits, 4) + self.assertEqual(r.get_config("layers.30.mlp.down_proj.weight").bits, 4) + # q_proj always int4 + self.assertEqual(r.get_config("layers.0.self_attn.q_proj.weight").bits, 4) + self.assertEqual(r.get_config("layers.30.self_attn.q_proj.weight").bits, 4) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/models/qwen3_5_moe/export.py b/examples/models/qwen3_5_moe/export.py index 385b386204c..f5d4ae209d2 100644 --- a/examples/models/qwen3_5_moe/export.py +++ b/examples/models/qwen3_5_moe/export.py @@ -178,6 +178,7 @@ def load_and_quantize(args): # noqa: C901 if args.prequantized: return load_prequantized_model_mlx( args.prequantized, + args.max_seq_len, use_splitk_decode=use_splitk, ) _prepare_and_quantize_mlx(model, config, args) @@ -300,60 +301,37 @@ def load_prequantized_model(prequantized_dir, max_seq_len=4096, use_splitk_decod return model, config -def load_prequantized_model_mlx(prequantized_dir, use_splitk_decode=True): +def load_prequantized_model_mlx( + prequantized_dir, max_seq_len=4096, use_splitk_decode=True +): """Load an MLX-format prequantized safetensors bundle into a model. - The bundle (from quantize_and_save.py --backend mlx) stores experts already - packed into stacked gather buffers and dense/lm_head weights as - ``ExportableInt4Tensor`` (embedding as ``IntxUnpackedToInt8Tensor``). Loading: + The bundle (from quantize_and_save.py --backend mlx) stores per-expert and + dense weights as torchao ``Int4Tensor`` (embedding as + ``IntxUnpackedToInt8Tensor``) in the modern torchao safetensors format. + Loading: 1. Build the model on meta and re-apply mlx_source_transformations so the module tree, swapped forwards, and MLX KV cache match the bundle. - 2. Convert each SwitchLinear to its packed buffer layout (experts are - already stacked/packed in the bundle) and assign the rest via - load_state_dict. - - max_seq_len is derived from a saved KV cache buffer so the rebuilt cache - geometry and export dynamic-shape bounds always match the bundle. + 2. ``load_checkpoint`` streams the bundle in, converting each ``Int4Tensor`` + to ``ExportableInt4Tensor`` via the default EXPORTABLE_PACKERS and + passing the int8 embedding through unchanged. + 3. ``pack_all_switch_linears`` stacks the per-expert weights into the + mlx::gather_qmm buffers. + + The KV cache is a runtime buffer and is not stored in the bundle; + mlx_source_transformations recreates it (zero-initialized) with geometry + derived from ``max_seq_len``. """ from executorch.backends.mlx.llm.switch import pack_all_switch_linears from executorch.examples.models.qwen3_5_moe.mlx_source_transformations import ( mlx_source_transformations, ) - from executorch.examples.models.qwen3_5_moe.quantize_and_save import ( - load_quantized_state_dict, - ) + from executorch.extension.llm.export.load import load_checkpoint config_path = os.path.join(prequantized_dir, "config.json") safetensors_path = os.path.join(prequantized_dir, "model.safetensors") - print(f"Loading prequantized MLX weights from {safetensors_path}...") - state_dict = load_quantized_state_dict(safetensors_path) - - # Int4Tensor -> ExportableInt4Tensor so dense linears export via - # dequantize_int4_tensor and experts pack via pack_all_switch_linears. - # Coarse/per-axis int8 (e.g. the embedding at group_size=hidden) keeps its - # group_size; the MLX pattern handlers regroup scale/zero_point to an - # MLX-legal size at export time (regroup_affine_scales). - from executorch.extension.llm.export.int4 import ExportableInt4Tensor - from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor - - for key, val in list(state_dict.items()): - if isinstance(val, Int4Tensor): - state_dict[key] = ExportableInt4Tensor.from_int4_tensor(val) - - # Derive max_seq_len from a saved KV cache buffer ([1, H, max_seq_len, D]). - max_seq_len = None - for key, val in state_dict.items(): - if key.endswith(".kv_cache.k_cache"): - max_seq_len = val.shape[2] - break - if max_seq_len is None: - raise RuntimeError( - "Prequantized MLX bundle has no KV cache buffer; cannot infer " - "max_seq_len (is this a CUDA bundle? use --backend cuda)." - ) - config = Qwen35MoEConfig.from_hf_config(config_path) config.max_seq_len = max_seq_len config.use_splitk_decode = use_splitk_decode @@ -371,30 +349,12 @@ def load_prequantized_model_mlx(prequantized_dir, use_splitk_decode=True): fuse_gate_up=False, ) - missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True) - del state_dict - - runtime_prefixes = (".mask", ".inv_freq", ".cache_positions") - expected_missing = {k for k in missing if any(p in k for p in runtime_prefixes)} - weight_missing = set(missing) - expected_missing - if weight_missing: - raise RuntimeError( - f"Prequantized MLX checkpoint is missing {len(weight_missing)} weight " - f"keys (model/checkpoint version mismatch?): {sorted(weight_missing)[:10]}" - ) - if unexpected: - raise RuntimeError( - f"Prequantized MLX checkpoint has {len(unexpected)} unexpected keys " - f"(model/checkpoint version mismatch?): {sorted(unexpected)[:10]}" - ) + print(f"Loading prequantized MLX weights from {safetensors_path}...") + load_checkpoint(safetensors_path, model) # Stack per-expert ExportableInt4Tensors into mlx::gather_qmm buffers. pack_all_switch_linears(model) - # assign=True wraps assigned tensors as Parameter(requires_grad=True), which - # breaks unwrap_tensor_subclass_parameters on int-dtype quantized inners. - for p in model.parameters(): - p.requires_grad_(False) model.eval() print( diff --git a/examples/models/qwen3_5_moe/quantize_and_save.py b/examples/models/qwen3_5_moe/quantize_and_save.py index 3a64fac8251..c6de5fcf706 100644 --- a/examples/models/qwen3_5_moe/quantize_and_save.py +++ b/examples/models/qwen3_5_moe/quantize_and_save.py @@ -293,6 +293,30 @@ def _materialize_meta_buffers(module): ) +def _stream_quantized(module, recipe, prefix, accum, dtype=torch.bfloat16): + """Append (prefixed) quantized params + persistent buffers to ``accum``. + + The lazy dual of ``quantize_model``: parameters stream through + ``quantize_stream`` (recipe applied per weight), while persistent buffers + pass through unchanged (never cast — some, e.g. a position table, are int). + The KV cache is a runtime buffer recreated (zeroed) at load, so it is + skipped rather than serialized. + """ + from executorch.extension.llm.export.quant import quantize_stream + + persistent = set(module.state_dict().keys()) + params = ((n, p.data) for n, p in module.named_parameters()) + seen = set() + for name, val in quantize_stream(params, recipe, dtype): + seen.add(name) + if ".kv_cache." in name: + continue + accum.append((prefix + name, val)) + for name, buf in module.named_buffers(): + if name in persistent and name not in seen and ".kv_cache." not in name: + accum.append((prefix + name, buf.data)) + + def _mlx_quant_recipe(config, group_size): """gemma4-style recipe for Qwen 3.5 MoE (applied to source-transformed model). @@ -301,7 +325,7 @@ def _mlx_quant_recipe(config, group_size): left unquantized. Matched with re.fullmatch. Uses min/max (single-pass) — HQQ's iterative scale search is far too slow across per-expert linears. """ - from executorch.examples.models.gemma4_31b.quant import ( + from executorch.extension.llm.export.quant import ( QuantConfig, QuantRecipe, QuantRule, @@ -325,13 +349,11 @@ def _mlx_quant_recipe(config, group_size): def stream_quantize_and_save_mlx(model_dir, config, args, safetensors_path): """Quantize + save an MLX bundle one decoder layer at a time (low memory). - Uses gemma4's ``quantize_model`` to produce ``Int4Tensor`` / + Uses the shared ``quantize_stream`` to produce ``Int4Tensor`` / ``IntxUnpackedToInt8Tensor`` subclasses. Packing is deferred to load (``pack_all_switch_linears``). Peak memory is ~one bf16 decoder layer. """ import torch.nn as nn - - from executorch.examples.models.gemma4_31b.quant import quantize_model from executorch.examples.models.qwen3_5_moe.mlx_source_transformations import ( mlx_source_transformations, ) @@ -372,8 +394,7 @@ def stream_quantize_and_save_mlx(model_dir, config, args, safetensors_path): sort_experts=True, fuse_gate_up=False, ) - for name, val in quantize_model(block, recipe).items(): - accum.append((prefix + name, val)) + _stream_quantized(block, recipe, prefix, accum) del block print( f" Streamed layer {i + 1}/{config.num_hidden_layers}", end="\r", flush=True @@ -409,8 +430,7 @@ def __init__(self): # Matches _swap_rms_norm: precompute (1 + weight) used by F.rms_norm. top.norm._rms_weight = nn.Parameter(1.0 + top.norm.weight.data) - for name, val in quantize_model(top, recipe).items(): - accum.append((name, val)) + _stream_quantized(top, recipe, "", accum) print("Writing quantized weights...") return save_quantized_tensors(accum, safetensors_path) diff --git a/examples/models/qwen3_5_moe/test_quantize_roundtrip.py b/examples/models/qwen3_5_moe/test_quantize_roundtrip.py index d4ff64bdb6f..430f31b3c4d 100644 --- a/examples/models/qwen3_5_moe/test_quantize_roundtrip.py +++ b/examples/models/qwen3_5_moe/test_quantize_roundtrip.py @@ -218,9 +218,9 @@ class TestSerializerFormatDispatch(unittest.TestCase): @staticmethod def _cpu_subclass(): - from executorch.examples.models.gemma4_31b.quant import ( + from executorch.extension.llm.export.quant import ( QuantConfig, - quantize_model, + quantize_stream, QuantRecipe, QuantRule, ) @@ -237,7 +237,12 @@ def _cpu_subclass(): ) ] ) - return "0.weight", next(iter(quantize_model(model, recipe).values())) + # Serializer round-trip: use the serialization form (torchao-native + # Int4Tensor from quantize_stream), which is what quantize_and_save + # writes. quantize_model yields an ExportableInt4Tensor (in-memory form) + # that torchao's flatten cannot serialize. + params = ((n, p.data) for n, p in model.named_parameters()) + return next(iter(quantize_stream(params, recipe))) def _assert_subclass_equal(self, a, b): names_a, attrs_a = a.__tensor_flatten__() diff --git a/extension/llm/export/BUCK b/extension/llm/export/BUCK index 4d17118d123..74780f949a1 100644 --- a/extension/llm/export/BUCK +++ b/extension/llm/export/BUCK @@ -59,6 +59,24 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "int4", + srcs = [ + "int4.py", + ], + _is_external_target = True, + base_module = "executorch.extension.llm.export", + visibility = [ + "//executorch/backends/...", + "//executorch/examples/...", + "//executorch/extension/llm/...", + ], + deps = [ + "//caffe2:torch", + "//pytorch/ao:torchao", + ], +) + fbcode_target(_kind = runtime.python_library, name = "gguf", srcs = [ @@ -72,6 +90,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/extension/llm/...", ], deps = [ + ":int4", "//caffe2:torch", "//pytorch/ao:torchao", "fbsource//third-party/pypi/gguf:gguf", @@ -79,6 +98,51 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "quant", + srcs = [ + "quant/__init__.py", + "quant/convert.py", + "quant/quantize.py", + "quant/recipe.py", + ], + _is_external_target = True, + base_module = "executorch.extension.llm.export", + visibility = [ + "//executorch/backends/...", + "//executorch/examples/...", + "//executorch/extension/llm/...", + ], + deps = [ + ":gguf", + ":int4", + "//caffe2:torch", + "//executorch/backends/cuda:dp4a_planar_int6_tensor", + "//pytorch/ao:torchao", + ], +) + +fbcode_target(_kind = runtime.python_library, + name = "load", + srcs = [ + "load.py", + ], + _is_external_target = True, + base_module = "executorch.extension.llm.export", + visibility = [ + "//executorch/backends/...", + "//executorch/examples/...", + "//executorch/extension/llm/...", + ], + deps = [ + ":gguf", + ":quant", + "//caffe2:torch", + "//pytorch/ao:torchao", + "fbsource//third-party/pypi/safetensors:safetensors", + ], +) + fbcode_target(_kind = runtime.python_binary, name = "export_llm", srcs = [ diff --git a/extension/llm/export/gguf.py b/extension/llm/export/gguf.py index 0f5565764c9..d7d8ae0d96c 100644 --- a/extension/llm/export/gguf.py +++ b/extension/llm/export/gguf.py @@ -10,13 +10,14 @@ ``ExportableGGUFTensor`` wraps the *raw* GGUF block bytes for one tensor and defers all unpacking, serving as the canonical GGUF loading representation: -* ``load_gguf(path)`` -> ``{name -> ExportableGGUFTensor | Tensor}`` (quantized - tensors become subclasses; F32/F16 stay plain). No unpacking at load. +* ``iter_gguf(path)`` -> streams ``(name, ExportableGGUFTensor | Tensor)`` pairs + (quantized tensors become subclasses; F32/F16 stay plain). No unpacking at load. * As a weight, it dequantizes via the ``torchao::dequantize_gguf`` custom op (gguf-package eager body) then a plain ``linear`` / ``embedding`` -- a backend can pattern-match ``dequantize_gguf`` -> linear/embedding to fuse. -* ``.to_int4_tensor()`` / ``.to_intx_unpacked_to_int8_tensor()`` convert into - torchao subclasses (``Int4Tensor`` / ``IntxUnpackedToInt8Tensor``) instead. +* ``.to_exportable_int4_tensor()`` / ``.to_intx_unpacked_to_int8_tensor()`` convert + into ExecuTorch/torchao subclasses (``ExportableInt4Tensor`` / + ``IntxUnpackedToInt8Tensor``) instead. The quant type is a string (``"q4_k"`` / ``"q5_k"`` / ``"q6_k"``); the ``gguf`` package's integer ``GGMLQuantizationType`` ids are an internal lookup detail. @@ -28,7 +29,7 @@ from __future__ import annotations -from typing import Dict, Iterator, Optional, Tuple, Union +from typing import Iterator, Optional, Tuple, Union import numpy as np import torch @@ -307,9 +308,9 @@ class ExportableGGUFTensor(TorchAOBaseTensor): Stores the exact GGUF ``block_q*_K`` byte layout (no repacking) plus the quant type string (``"q4_k"`` / ``"q5_k"`` / ``"q6_k"``). ``aten.linear`` / ``aten.embedding`` dequantize via the ``torchao::dequantize_gguf`` op (then a - plain linear/embedding); :meth:`to_int4_tensor` / - :meth:`to_intx_unpacked_to_int8_tensor` convert to torchao subclasses - instead. + plain linear/embedding); :meth:`to_exportable_int4_tensor` / + :meth:`to_intx_unpacked_to_int8_tensor` convert to ExecuTorch/torchao + subclasses instead. """ tensor_data_names = ["raw"] @@ -352,35 +353,57 @@ def from_raw( """Build from a ``(N, row_bytes)`` uint8 GGUF block blob.""" return cls(raw.contiguous(), ggml_type, orig_dtype) + def to(self, *args, **kwargs) -> "ExportableGGUFTensor": + """Move device and/or set the output dtype *without* dequantizing. + + Mirrors ``IntxUnpackedToInt8Tensor.to``: the raw GGUF blocks are only + moved across devices; ``dtype`` becomes the dtype :meth:`dequantize` + will produce. Casting via ``.to`` therefore preserves the quantized + payload (unlike ``aten._to_copy``, which dequantizes for export). + """ + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs.pop("device") + dtype = kwargs.pop("dtype") + assert dtype.is_floating_point, f"expected a floating dtype; got {dtype}" + return ExportableGGUFTensor(self.raw.to(device), self.ggml_type, dtype) + def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> Tensor: """Dequantize to a plain float tensor using the ``gguf`` package.""" return torch.ops.torchao.dequantize_gguf( self.raw, self.ggml_type, output_dtype or self.orig_dtype ) - def to_int4_tensor(self) -> Tensor: - """Convert a Q4_K tensor to a torchao ``Int4Tensor``.""" - from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor + def to_exportable_int4_tensor( + self, output_dtype: Optional[torch.dtype] = None + ) -> Tensor: + """Convert a Q4_K tensor to an ``ExportableInt4Tensor``. + + Threads ``output_dtype`` (default ``orig_dtype``) into the scale/zero_point + and the dequantized output dtype, so callers targeting fp16 (e.g. MLX) get + fp16 params while CUDA can pass bf16. The packed int nibbles are unchanged. + """ + from executorch.extension.llm.export.int4 import ExportableInt4Tensor if self.ggml_type != "q4_k": raise NotImplementedError( - f"to_int4_tensor only supports q4_k; got {self.ggml_type!r}" + f"to_exportable_int4_tensor only supports q4_k; got {self.ggml_type!r}" ) + dtype = output_dtype or self.orig_dtype N, K = int(self.shape[0]), int(self.shape[1]) q, eff_scale, eff_min = _q4_k_fields(self.raw, N, K) zero = torch.where( eff_scale != 0, eff_min / eff_scale, torch.zeros_like(eff_min) ) - # Nibble-pack for Int4Tensor: even index -> low nibble, odd -> high. + # Nibble-pack: even index -> low nibble, odd -> high. packed = q[:, ::2] | (q[:, 1::2] << 4) - return Int4Tensor( - qdata=packed, - # Int4Tensor scale/zero layout is (K // gs, N) -- transposed. - scale=eff_scale.to(torch.bfloat16).t().contiguous(), - zero_point=zero.to(torch.bfloat16).t().contiguous(), - block_size=[1, Q4_K_GROUP_SIZE], - shape=torch.Size([N, K]), + return ExportableInt4Tensor( + packed, + # ExportableInt4Tensor scale/zero layout is (K // gs, N) -- transposed. + eff_scale.to(dtype).t().contiguous(), + zero.to(dtype).t().contiguous(), + Q4_K_GROUP_SIZE, + dtype, ) def to_intx_unpacked_to_int8_tensor( @@ -552,11 +575,3 @@ def iter_gguf( yield tensor.name, flat.view(torch.bfloat16).reshape(shape).clone() else: raise ValueError(f"Unsupported GGUF quant type: {tensor.tensor_type}") - - -def load_gguf(path: str) -> Dict[str, Union[ExportableGGUFTensor, Tensor]]: - """Load a GGUF file into ``{name -> ExportableGGUFTensor | Tensor}``. - - Holds all tensors at once; use :func:`iter_gguf` for low peak memory. - """ - return dict(iter_gguf(path)) diff --git a/extension/llm/export/int4.py b/extension/llm/export/int4.py index 59251ae1875..89b2515b2e9 100644 --- a/extension/llm/export/int4.py +++ b/extension/llm/export/int4.py @@ -97,6 +97,37 @@ def from_int4_tensor(cls, w: Tensor) -> "ExportableInt4Tensor": w.dtype, ) + @classmethod + def from_intx_unpacked_to_int8_tensor(cls, w: Tensor) -> "ExportableInt4Tensor": + """Build from a 4-bit weight-only ``IntxUnpackedToInt8Tensor``. + + Raises ``ValueError`` if ``w`` is not 4-bit (``target_dtype`` is not + ``torch.int4``) or carries activation quantization — neither can be + represented as an ``ExportableInt4Tensor``. + + Nibble-packs the unpacked int8 ``qdata`` (values in [-8, 7]) into uint8 + and offsets ``qdata``/``zero_point`` to the unsigned [0, 15] convention; + the ``+8`` offset cancels in ``scale * (q - z)``, so the dequantized + values are unchanged. ``scale``/``zero_point`` are transposed from the + IntxUnpacked ``(N, K//gs)`` layout to the Int4Tensor ``(K//gs, N)``. + """ + if w.target_dtype != torch.int4: + raise ValueError( + "from_intx_unpacked_to_int8_tensor requires a 4-bit tensor " + f"(target_dtype=torch.int4), got target_dtype={w.target_dtype}" + ) + if w.activation_quantization is not None: + raise ValueError( + "from_intx_unpacked_to_int8_tensor requires a weight-only tensor; " + "activation quantization cannot be represented as an " + "ExportableInt4Tensor" + ) + q = (w.qdata.to(torch.int16) + 8).to(torch.uint8) # [-8, 7] -> [0, 15] + packed = q[..., ::2] | (q[..., 1::2] << 4) # (N, K) -> (N, K//2) + scale = w.scale.t().contiguous() + zero_point = (w.zero_point.to(w.scale.dtype) + 8).t().contiguous() + return cls(packed, scale, zero_point, int(w.block_size[-1]), w.dtype) + def dequantize(self, output_dtype=None): return torch.ops.torchao.dequantize_int4_tensor( self.qdata, @@ -106,6 +137,26 @@ def dequantize(self, output_dtype=None): output_dtype=output_dtype or self.orig_dtype, ) + def to(self, *args, **kwargs) -> "ExportableInt4Tensor": + """Move device and/or set the output dtype *without* dequantizing. + + Mirrors ``IntxUnpackedToInt8Tensor.to``: the packed int4 ``qdata`` only + moves across devices, while ``scale``/``zero_point`` and the output + dtype follow ``dtype``. Use :meth:`dequantize` to materialize a dense + tensor. + """ + kwargs = self._get_to_kwargs(*args, **kwargs) + device = kwargs.pop("device") + dtype = kwargs.pop("dtype") + assert dtype.is_floating_point, f"expected a floating dtype; got {dtype}" + return ExportableInt4Tensor( + self.qdata.to(device), + self.scale.to(device=device, dtype=dtype), + self.zero_point.to(device=device, dtype=dtype), + self.group_size, + dtype, + ) + __torch_function__ = torch._C._disabled_torch_function_impl diff --git a/extension/llm/export/load.py b/extension/llm/export/load.py new file mode 100644 index 00000000000..77c8f8ee74b --- /dev/null +++ b/extension/llm/export/load.py @@ -0,0 +1,315 @@ +# 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. + +"""Load a checkpoint (quantized or unquantized) into a meta-device model. + +``iter_checkpoint`` is the model-free reader: it picks a reader by ``path`` +(``.gguf`` -> raw GGUF blobs; a ``.safetensors`` file or a directory of them -> +safetensors, quantized or plain), remaps each tensor name to a model FQN via +``key_map`` (returning ``None`` skips it), then applies ``convert`` -- a +``(fqn, tensor) -> tensor`` step (default ``to_default``) that wraps quantized +subclasses in their portable ``Exportable*`` form. It yields ``(fqn, tensor)`` +pairs and never touches a model, so ``dict(iter_checkpoint(path, convert=identity))`` +is a plain state dict. + +``load_checkpoint`` is the model-aware wrapper: it streams ``iter_checkpoint`` +into ``model`` via ``assign_one`` (param-vs-buffer, no packer dispatch), fans a +weight out to a second FQN via ``tie_map`` (e.g. a tied lm_head, keyed by the +*remapped* FQN and applied only when the destination is still on meta), and +asserts every *parameter* is materialized once the stream is exhausted (runtime +buffers may still be on meta -- callers fill those afterward, e.g. via +``materialize_runtime_buffers``). Backend-specific layout conversion (CUDA +coalesced int4, MLX gather buffers) is a separate terminal pass over the loaded +model, not part of loading. +""" + +import contextlib +import json +import os +from typing import Callable, Iterable, Optional + +import torch +import torch.nn as nn +from executorch.extension.llm.export.quant.convert import ( + Convert, + maybe_cast, + to_default, +) + +# Remap a checkpoint tensor name to a model FQN, or ``None`` to skip it. +KeyMap = Callable[[str], Optional[str]] + +# Fan a source tensor out to a second module: remapped source FQN -> +# (destination FQN, clone). ``clone=True`` packs an independent copy of the +# tensor into the destination (untie); ``clone=False`` packs the same tensor +# object, so a pass-through packer leaves both modules sharing one storage (tie). +# The fan-out is applied only when the destination is still on meta, so a real +# weight already present in the checkpoint always wins. +TieMap = dict[str, tuple[str, bool]] + +Pair = tuple[str, torch.Tensor] + +# Custom raw reader: ``path -> (name, tensor)`` pairs (before key_map/convert/ +# dtype). When given to ``iter_checkpoint`` / ``load_checkpoint`` it overrides the +# built-in ``.gguf`` / ``.safetensors`` reader selection, so callers can inject a +# format-specific reader (e.g. guac's MLX affine reader) without this module +# needing to know about it. +RawIter = Callable[[str], Iterable[Pair]] + + +def _split_parent(model: nn.Module, fqn: str) -> tuple[nn.Module, str]: + """Resolve ``fqn`` to ``(parent_module, attr_name)``.""" + parts = fqn.rsplit(".", 1) + parent_fqn = parts[0] if len(parts) > 1 else "" + attr = parts[-1] + parent = model.get_submodule(parent_fqn) if parent_fqn else model + return parent, attr + + +def assert_all_materialized(model: nn.Module) -> None: + """Raise if any parameter is still on the meta device after loading.""" + for fqn, p in model.named_parameters(): + if p.device.type == "meta": + raise RuntimeError( + f"Weight '{fqn}' not found in checkpoint " + f"(model/checkpoint version mismatch?)" + ) + + +def assign_one(model: nn.Module, fqn: str, value: torch.Tensor) -> None: + """Assign one weight into ``model`` as a parameter or buffer, in place. + + Chooses parameter vs buffer by inspecting the existing attribute on the + parent module (a slot that is currently an ``nn.Parameter`` -- typically a + meta tensor -- becomes a non-grad ``nn.Parameter``; anything else is + registered as a buffer). No conversion or packer dispatch: ``value`` is + assigned as given (already in its final portable/quantized form). + """ + parent, attr = _split_parent(model, fqn) + if isinstance(getattr(parent, attr, None), nn.Parameter): + setattr(parent, attr, nn.Parameter(value, requires_grad=False)) + else: + parent.register_buffer(attr, value) + + +def _resolve_shard_paths(path: str) -> list[str]: + """Return the safetensors shard file paths for ``path``. + + A file path yields itself; a directory is resolved via + ``model.safetensors.index.json`` (sharded) or a single ``model.safetensors``. + """ + if not os.path.isdir(path): + return [path] + + index_path = os.path.join(path, "model.safetensors.index.json") + if os.path.exists(index_path): + with open(index_path, "r") as f: + index = json.load(f) + return [ + os.path.join(path, shard) + for shard in sorted(set(index["weight_map"].values())) + ] + single = os.path.join(path, "model.safetensors") + if os.path.exists(single): + return [single] + raise FileNotFoundError(f"No safetensors checkpoint in {path}") + + +def _iter_safetensors(path: str) -> Iterable[Pair]: + """Yield ``(name, tensor)`` pairs from a safetensors checkpoint. + + ``path`` may be a single ``.safetensors`` file or a directory (sharded via + ``model.safetensors.index.json``, or a single ``model.safetensors``). All + shards are opened together (mmap, so no tensor data is read up front) and a + ``key -> shard`` map is built, so a logical weight whose sub-tensors are + split across shards can still be reassembled. + + A torchao-quantized checkpoint (identified by a ``tensor_names`` metadata + key) has each logical weight's subclass sub-tensors gathered -- from whatever + shard holds them -- and reassembled via ``unflatten_tensor_state_dict`` into + a single quantized tensor. A plain checkpoint yields its tensors as-is. + Either way only one logical weight is materialized at a time. Casting to a + target dtype is the caller's job (``iter_checkpoint``). + """ + from safetensors import safe_open + + with contextlib.ExitStack() as stack: + handles = [ + stack.enter_context(safe_open(p, framework="pt", device="cpu")) + for p in _resolve_shard_paths(path) + ] + key_to_handle = {} + metadata: dict = {} + tensor_names: list = [] + for handle in handles: + shard_meta = handle.metadata() or {} + metadata.update(shard_meta) + tensor_names.extend( + name + for name in json.loads(shard_meta.get("tensor_names", "[]")) + if name not in tensor_names + ) + for key in handle.keys(): + key_to_handle[key] = handle + + if not tensor_names: + for key, handle in key_to_handle.items(): + yield key, handle.get_tensor(key) + return + + from torchao.prototype.safetensors.safetensors_support import ( + unflatten_tensor_state_dict, + ) + + for name in tensor_names: + parts = name.rsplit(".", 1) + module_fqn = parts[0] if len(parts) > 1 else "" + weight_name = parts[-1] + prefix = ( + f"{module_fqn}._{weight_name}_" if module_fqn else f"_{weight_name}_" + ) + partial = { + key: handle.get_tensor(key) + for key, handle in key_to_handle.items() + if key == name or key.startswith(prefix) + } + result, _ = unflatten_tensor_state_dict(partial, metadata) + for fqn, value in result.items(): + yield fqn, value + + +def _iter_gguf(path: str) -> Iterable[Pair]: + """Yield ``(gguf_name, tensor)`` pairs from a GGUF file. + + Quantized weights arrive as ``ExportableGGUFTensor`` (the raw GGUF blob). + Names are the raw GGUF tensor names -- pass a ``key_map`` to + ``iter_checkpoint`` to remap them to model FQNs. Casting to a target dtype is + the caller's job (``iter_checkpoint``). + """ + from executorch.extension.llm.export.gguf import iter_gguf + + yield from iter_gguf(path) + + +def iter_checkpoint( + path: str, + *, + key_map: Optional[KeyMap] = None, + convert: Convert = to_default, + tie_map: Optional[TieMap] = None, + dtype: Optional[torch.dtype] = None, + raw_iter: Optional[RawIter] = None, +) -> Iterable[Pair]: + """Yield ``(fqn, tensor)`` pairs from a checkpoint, model-free. + + ``raw_iter`` (when given) supplies the raw ``(name, tensor)`` reader and + overrides the selection below; otherwise the reader is picked by ``path``: + + * ``*.gguf`` -- raw GGUF blobs; quantized weights become + ``ExportableGGUFTensor``. + * ``*.safetensors`` -- a single safetensors file. + * a directory -- a safetensors checkpoint, sharded via + ``model.safetensors.index.json`` or a single + ``model.safetensors``. + + For each weight, ``key_map`` remaps the stream name to a model FQN first + (identity when ``None``; return ``None`` to skip), then ``convert`` maps the + value (default ``to_default``: ``Int4Tensor -> ExportableInt4Tensor``, others + unchanged), then ``dtype`` casts the *converted* weight (plain tensors + normally; quantized subclasses keep their payload and only re-stamp their + dequantized output dtype). Casting runs after ``convert`` so a raw torchao + ``Int4Tensor`` -- whose own ``.to`` ignores dtype -- is first wrapped as an + ``ExportableInt4Tensor`` (which ``maybe_cast`` *can* re-stamp, e.g. to fp16 + for MLX). ``convert`` therefore sees the *remapped* FQN and the *uncast* + tensor. Only one logical weight is materialized at a time. + + ``tie_map`` fans a source weight out to a second FQN (e.g. a tied lm_head, + keyed by the *remapped* source FQN). The tied copy is emitted after the stream + is exhausted and only when the destination did not appear on its own, so an + untied checkpoint always wins; ``clone=True`` emits an independent copy. + """ + if raw_iter is not None: + raw = raw_iter(path) + elif path.endswith(".gguf"): + raw = _iter_gguf(path) + elif path.endswith(".safetensors") or os.path.isdir(path): + raw = _iter_safetensors(path) + else: + raise ValueError( + f"Cannot load checkpoint from '{path}': expected a .gguf file, a " + ".safetensors file, or a directory with a safetensors checkpoint." + ) + seen: set[str] = set() + owed: dict[str, Pair] = {} # dst FQN -> (source value, clone) for tie fan-out + for name, value in raw: + fqn = key_map(name) if key_map is not None else name + if fqn is None: + continue + out = maybe_cast(convert(fqn, value), dtype) + seen.add(fqn) + yield fqn, out + if tie_map is not None and fqn in tie_map: + dst, clone = tie_map[fqn] + owed[dst] = (out, clone) + for dst, (value, clone) in owed.items(): + if dst not in seen: + yield dst, value.clone() if clone else value + + +def load_checkpoint( + path: str, + model: nn.Module, + *, + key_map: Optional[KeyMap] = None, + convert: Convert = to_default, + tie_map: Optional[TieMap] = None, + dtype: Optional[torch.dtype] = None, + raw_iter: Optional[RawIter] = None, +) -> None: + """Load a checkpoint and assign it into ``model`` in place. + + Streams ``iter_checkpoint(path, key_map=key_map, convert=convert, + tie_map=tie_map, dtype=dtype)`` (see it for reader selection and the + ``key_map`` / ``convert`` / ``tie_map`` / ``dtype`` semantics) and assigns each + ``(fqn, tensor)`` via ``assign_one`` (param-vs-buffer, no packer dispatch -- + the tensor is already in its final portable form). + + Asserts every parameter is materialized afterward; runtime buffers left on + meta are the caller's responsibility (e.g. ``materialize_runtime_buffers``). + Backend-specific layout conversion (CUDA coalesced int4, MLX gather buffers) + is a separate terminal pass over the loaded model. + """ + for fqn, value in iter_checkpoint( + path, + key_map=key_map, + convert=convert, + tie_map=tie_map, + dtype=dtype, + raw_iter=raw_iter, + ): + assign_one(model, fqn, value) + assert_all_materialized(model) + + +def assign_state_dict( + model: nn.Module, + state_dict: dict[str, torch.Tensor], + convert: Convert = to_default, +) -> None: + """Convert + assign an in-memory state dict into ``model`` in place. + + The in-memory dual of :func:`load_checkpoint`: applies ``convert`` (default + ``to_default``) to each weight and assigns it via :func:`assign_one`. Use + this when weights are already materialized in memory (e.g. the output of + ``quantize_model``) rather than streamed from disk. Backend layout conversion + (if any) is a separate terminal pass over the assigned model. Asserts every + parameter is materialized afterward. + """ + for fqn, value in state_dict.items(): + assign_one(model, fqn, convert(fqn, value)) + assert_all_materialized(model) + for p in model.parameters(): + p.requires_grad_(False) diff --git a/extension/llm/export/quant/__init__.py b/extension/llm/export/quant/__init__.py new file mode 100644 index 00000000000..21644b5de82 --- /dev/null +++ b/extension/llm/export/quant/__init__.py @@ -0,0 +1,20 @@ +# 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. + +from .convert import ( # noqa: F401 + Convert, + fuse_along_output, + identity, + to_default, + to_exportable, +) +from .quantize import ( # noqa: F401 + dequantize_weight, + quantize_model, + quantize_stream, + quantize_weight, +) +from .recipe import QuantConfig, QuantRecipe, QuantRule # noqa: F401 diff --git a/extension/llm/export/quant/convert.py b/extension/llm/export/quant/convert.py new file mode 100644 index 00000000000..139fcfdc656 --- /dev/null +++ b/extension/llm/export/quant/convert.py @@ -0,0 +1,216 @@ +# 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. + +"""Model-free weight representation conversion. + +Two ``convert`` helpers map a raw torchao quantized weight to its portable +``Exportable*`` form. They are ``(fqn, tensor) -> tensor`` so the same function +drives ``iter_checkpoint`` (streaming from disk) and ``assign_state_dict`` (an +in-memory quantized state dict), both in ``extension/llm/export/load.py``: + + * ``to_exportable`` wraps a torchao ``Int4Tensor`` as + ``ExportableInt4Tensor`` -- so it exports to the portable + ``dequantize_int4_tensor`` op instead of torchao's hardware-specific ``mslk`` + kernels -- and passes every other tensor through unchanged. + * ``to_default`` (the default ``convert``) additionally repacks a 4-bit + ``IntxUnpackedToInt8Tensor`` (stored one value per int8 byte) into the + nibble-packed ``ExportableInt4Tensor`` to halve its footprint. + * ``identity`` yields the raw tensor/subclass untouched. + +These only convert a weight's *representation*; assigning it into a model +(``assign_one`` / ``assign_state_dict``) and backend-specific layout conversion +(CUDA coalesced int4, MLX gather buffers) live elsewhere. + +:func:`fuse_along_output` is the multi-weight companion: it concatenates several +weights (e.g. ``gate_proj|up_proj`` or ``q|k|v``) along the output-channel dim, +subclass-aware, so fusion happens once on the portable form before backend +packing. It handles the same tensor types as :func:`maybe_cast`. +""" + +from collections.abc import Sequence +from typing import Callable, Optional + +import torch + +# A weight conversion step: (fqn, tensor) -> tensor. Runs after key remapping, +# so it may key off the model FQN; the built-ins below ignore the name. +Convert = Callable[[str, torch.Tensor], torch.Tensor] + + +def maybe_cast(value: torch.Tensor, dtype: Optional[torch.dtype]) -> torch.Tensor: + """Cast a weight to ``dtype`` only when it is safe and meaningful. + + * ``dtype`` is None / already matches / non-floating -> returned unchanged. + * a plain (unquantized) tensor -> cast normally. + * a quantized subclass whose ``.to`` preserves the payload + (``ExportableGGUFTensor``, ``ExportableInt4Tensor``, + ``IntxUnpackedToInt8Tensor`` -- only ``scale``/``zero_point`` and the + dequantized output dtype change) -> cast via that ``.to``. + * any other subclass -> returned unchanged, so the cast can never silently + dequantize it (via an ``aten._to_copy`` override) or no-op it (torchao's + own ``Int4Tensor.to`` ignores dtype). + + Shared by ``iter_checkpoint`` (which casts the *converted* weight, so a raw + ``Int4Tensor`` is wrapped as an ``ExportableInt4Tensor`` first and can then be + re-stamped, e.g. to fp16 for MLX) and ``quantize_stream`` (which casts each + weight after applying the recipe), so both cast identically. + """ + if dtype is None or value.dtype == dtype or not value.dtype.is_floating_point: + return value + if type(value) is torch.Tensor: + return value.to(dtype) + + from executorch.extension.llm.export.gguf import ExportableGGUFTensor + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + from torchao.quantization import IntxUnpackedToInt8Tensor + + castable = (ExportableGGUFTensor, ExportableInt4Tensor, IntxUnpackedToInt8Tensor) + return value.to(dtype) if isinstance(value, castable) else value + + +def _is_quantized(value: torch.Tensor) -> bool: + """Check if a tensor is a torchao quantized subclass.""" + from torchao.utils import TorchAOBaseTensor + + return isinstance(value, TorchAOBaseTensor) + + +def fuse_along_output(tensors: Sequence[torch.Tensor]) -> torch.Tensor: + """Concatenate weights along the output-channel (``N``) dim, subclass-aware. + + Fuses e.g. ``gate_proj|up_proj`` or ``q|k|v`` into one weight *before* + backend packing, so both the MLX and CUDA packers consume a single fused + weight (no backend-specific concat needed). ``N`` is dim 0 of the logical + ``(N, K)`` weight. Handles the same types as :func:`maybe_cast`: + + * a plain (unquantized) tensor -> ``torch.cat(dim=0)``. + * ``ExportableInt4Tensor`` -> ``qdata`` on dim 0, ``scale``/``zero_point`` on + dim 1 (they are stored transposed as ``(K // group_size, N)``). + * ``IntxUnpackedToInt8Tensor`` / ``ExportableGGUFTensor`` -> every packed + field on dim 0 (all are ``N``-major). + + Exact: ``scale``/``zero_point`` are per-(group, output-channel) and each + packed row is an independent group / GGUF super-block, so concatenating + output channels leaves every channel's metadata untouched. All inputs must + be the same type and share their quant params + inner (``K``) shape; + ``torch.cat`` enforces the shapes and a mismatched attribute raises + ``ValueError``. A single tensor is returned unchanged; an empty input or an + unsupported subclass raises. + """ + tensors = list(tensors) + if not tensors: + raise ValueError("fuse_along_output requires at least one tensor") + first = tensors[0] + if len(tensors) == 1: + return first + + # Fusability: all inputs must be the same kind. Plain weights must all be + # plain; quantized weights must be the exact same subclass (with matching + # attributes, checked in _fuse_subclass). torch.cat then enforces the inner + # (K) shape. + if not _is_quantized(first): + if any(_is_quantized(t) for t in tensors): + raise TypeError( + "fuse_along_output cannot mix plain and quantized tensors; got " + f"{[type(t).__name__ for t in tensors]}" + ) + return torch.cat(tensors, dim=0) + if any(type(t) is not type(first) for t in tensors): + raise TypeError( + "fuse_along_output requires all tensors to be the same type; got " + f"{[type(t).__name__ for t in tensors]}" + ) + + from executorch.extension.llm.export.gguf import ExportableGGUFTensor + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + from torchao.quantization import IntxUnpackedToInt8Tensor + + # Per-data-field concat axis (the output-channel dim) by subclass layout. + axis_by_field = { + ExportableInt4Tensor: {"qdata": 0, "scale": 1, "zero_point": 1}, + IntxUnpackedToInt8Tensor: {"qdata": 0, "scale": 0, "zero_point": 0}, + ExportableGGUFTensor: {"raw": 0}, + }.get(type(first)) + if axis_by_field is None: + raise TypeError( + f"fuse_along_output cannot fuse tensors of type {type(first).__name__}" + ) + return _fuse_subclass(tensors, axis_by_field) + + +def _fuse_subclass( + tensors: list[torch.Tensor], axis_by_field: dict[str, int] +) -> torch.Tensor: + """Rebuild a TorchAO subclass from its data fields concatenated per axis. + + Relies on the ``(*tensor_data_names, *tensor_attribute_names)`` constructor + convention shared by these subclasses; attributes must match across inputs. + """ + first = tensors[0] + attrs = [getattr(first, name) for name in first.tensor_attribute_names] + for t in tensors[1:]: + for name, ref in zip(first.tensor_attribute_names, attrs): + if getattr(t, name) != ref: + raise ValueError( + "fuse_along_output: cannot fuse tensors with different " + f"'{name}' ({getattr(t, name)!r} != {ref!r})" + ) + fused_data = [ + torch.cat([getattr(t, name) for t in tensors], dim=axis_by_field[name]) + for name in first.tensor_data_names + ] + return type(first)(*fused_data, *attrs) + + +def _to_exportable_int4(w: torch.Tensor, *, repack_intx: bool) -> torch.Tensor: + """Wrap 4-bit torchao weights as ``ExportableInt4Tensor``. + + Always converts a torchao ``Int4Tensor``. When ``repack_intx`` is set, also + repacks a 4-bit ``IntxUnpackedToInt8Tensor`` into the nibble-packed int4 + representation (raising if it carries activation quantization, which int4 + export can't represent). Any other tensor is returned unchanged. + """ + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + from torchao.quantization import IntxUnpackedToInt8Tensor + from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor + + if isinstance(w, Int4Tensor): + return ExportableInt4Tensor.from_int4_tensor(w) + if ( + repack_intx + and isinstance(w, IntxUnpackedToInt8Tensor) + and w.target_dtype == torch.int4 + ): + return ExportableInt4Tensor.from_intx_unpacked_to_int8_tensor(w) + return w + + +def to_exportable(fqn: str, w: torch.Tensor) -> torch.Tensor: + """Default convert: ``Int4Tensor -> ExportableInt4Tensor``; others unchanged. + + ``Int4Tensor`` is wrapped so it exports to the portable + ``dequantize_int4_tensor -> linear/embedding`` op (torchao's own + ``Int4Tensor`` would otherwise lower to hardware-specific ``mslk`` kernels). + All other subclasses -- including 4-bit ``IntxUnpackedToInt8Tensor`` -- already + export to portable ``dequantize_affine`` and are passed through. + """ + return _to_exportable_int4(w, repack_intx=False) + + +def to_default(fqn: str, w: torch.Tensor) -> torch.Tensor: + """Like :func:`to_exportable`, but also repacks a 4-bit + ``IntxUnpackedToInt8Tensor`` into the nibble-packed ``ExportableInt4Tensor`` + (halving its footprint vs the int8 container). + + Opt-in (e.g. the MLX gemma4 export), since intx already exports portably via + ``dequantize_affine`` and int4 is purely a memory optimization here. + """ + return _to_exportable_int4(w, repack_intx=True) + + +def identity(fqn: str, w: torch.Tensor) -> torch.Tensor: + """No-op convert: yield the raw tensor/subclass unchanged.""" + return w diff --git a/extension/llm/export/quant/quantize.py b/extension/llm/export/quant/quantize.py new file mode 100644 index 00000000000..747b38bba2b --- /dev/null +++ b/extension/llm/export/quant/quantize.py @@ -0,0 +1,386 @@ +# 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. + +"""Quantize weights to torchao tensor subclasses. + +``quantize_weight`` quantizes a single tensor given a ``QuantConfig``, +returning an ``Int4Tensor`` (4-bit) or ``IntxUnpackedToInt8Tensor`` (5-, 6-, +or 8-bit). + +``quantize_model`` walks a model's parameters, applies a ``QuantRecipe``, +and returns a single state dict containing both quantized subclass tensors +and unquantized plain tensors. + +``quantize_stream`` and ``quantize_model`` differ by design: ``quantize_stream`` +is the *serialization* producer (yields torchao-native ``Int4Tensor``, which +round-trips through torchao's safetensors allowlist and is what +``quantize_and_save`` writes to disk), while ``quantize_model`` is the +*in-memory model* producer -- the dual of :func:`load_checkpoint`, applying +``convert`` (default ``to_default``: ``Int4Tensor -> ExportableInt4Tensor``) +then :func:`maybe_cast` so the state dict is castable (e.g. fp16 for MLX) and +in the export-canonical form the backend packers consume. +""" + +from collections.abc import Iterable, Iterator + +import torch +import torch.nn as nn + +from .convert import Convert, maybe_cast, to_default +from .recipe import QuantConfig, QuantRecipe + +# --------------------------------------------------------------------------- +# Per-weight quantization + + +def _quantize_affine_min_max( + weight: torch.Tensor, + group_size: int, + qmin: int, + qmax: int, + symmetric: bool, + zero_point_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Affine min/max quantization. Returns (int_data, scale, zero_point). + + ``int_data`` is int8 in [qmin, qmax]; ``scale``/``zero_point`` keep the + ``choose_qparams_affine`` block-reduced shape. + """ + from torchao.quantization.quant_primitives import ( + choose_qparams_affine, + MappingType, + quantize_affine, + ) + + mapping = MappingType.SYMMETRIC if symmetric else MappingType.ASYMMETRIC + block_size = tuple([1] * (weight.ndim - 1) + [group_size]) + weight = weight.float() + scale, zero_point = choose_qparams_affine( + weight, + mapping, + block_size, + target_dtype=torch.int8, + quant_min=qmin, + quant_max=qmax, + scale_dtype=torch.bfloat16, + zero_point_dtype=zero_point_dtype, + ) + int_data = quantize_affine( + weight, + block_size, + scale, + zero_point, + output_dtype=torch.int8, + quant_min=qmin, + quant_max=qmax, + ) + return int_data, scale, zero_point + + +def _quantize_scale_only_hqq( + weight_2d: torch.Tensor, + group_size: int, + qmin: int, + qmax: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Scale-only HQQ on a 2D weight. Returns (int_data int8, scale bf16).""" + from torchao.quantization.quant_primitives import ( + _choose_qparams_and_quantize_scale_only_hqq, + ) + + qdata, scale = _choose_qparams_and_quantize_scale_only_hqq( + weight_2d, [1, group_size], qmin, qmax + ) + return qdata.to(torch.int8), scale.to(torch.bfloat16) + + +def _quantize_min_max( + weight: torch.Tensor, + config: QuantConfig, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Standard min/max 4-bit quantization. Returns (int_data, scale, zero_point).""" + qmin, qmax = (-8, 7) if config.symmetric else (0, 15) + return _quantize_affine_min_max( + weight, + config.group_size, + qmin, + qmax, + config.symmetric, + zero_point_dtype=torch.bfloat16, + ) + + +def _quantize_hqq_asymmetric( + weight: torch.Tensor, + config: QuantConfig, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Full HQQ (asymmetric, optimizes scale + zero). Requires CUDA.""" + from torchao.quantization.quant_primitives import ( + _choose_qparams_and_quantize_affine_hqq, + ) + + device = weight.device + if device.type != "cuda": + device = torch.device("cuda") + + W_q, scale, zero, _shape = _choose_qparams_and_quantize_affine_hqq( + weight, + nbits=config.bits, + group_size=config.group_size, + axis=1, + compute_dtype=torch.bfloat16, + device=str(device), + raw_output=True, + ) + + int_data = W_q.to(torch.int8) + scale = scale.to(torch.bfloat16).reshape(*weight.shape[:-1], -1) + zero = zero.to(torch.bfloat16).reshape(*weight.shape[:-1], -1) + + return int_data, scale, zero + + +def _quantize_hqq_symmetric( + weight: torch.Tensor, + config: QuantConfig, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Scale-only HQQ (symmetric 4-bit, optimizes scale only). Runs on CPU or CUDA.""" + orig_shape = weight.shape + weight_2d = weight.reshape(-1, weight.shape[-1]) if weight.ndim > 2 else weight + int_data, scale = _quantize_scale_only_hqq(weight_2d, config.group_size, -8, 7) + int_data = int_data.reshape(orig_shape) + scale = scale.reshape(*orig_shape[:-1], -1) + zero_point = torch.zeros_like(scale) + return int_data, scale, zero_point + + +def _to_int4_tensor( + int_data: torch.Tensor, + scale: torch.Tensor, + zero_point: torch.Tensor, + config: QuantConfig, +) -> torch.Tensor: + """Wrap quantized 4-bit data into an Int4Tensor.""" + from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor + + # Normalize 4-bit signed [-8, 7] to unsigned [0, 15] for storage. + if config.symmetric: + int_data = int_data + 8 + zero_point = torch.full_like(scale, 8.0) + + # Int4Tensor stores qdata as nibble-packed uint8 (N, K//2) + q = int_data.to(torch.uint8) + packed = q[..., ::2] | (q[..., 1::2] << 4) + + # Int4Tensor stores scale/zero as (K//gs, N) — transposed from our (N, K//gs) + return Int4Tensor( + qdata=packed, + scale=scale.t().contiguous(), + zero_point=zero_point.t().contiguous(), + block_size=[1, config.group_size], + shape=torch.Size(int_data.shape), + ) + + +def _to_intx_tensor( + weight: torch.Tensor, + config: QuantConfig, +) -> torch.Tensor: + """Quantize to 5-, 6-, or 8-bit and wrap in IntxUnpackedToInt8Tensor. + + Quantizes in float32 for numerical precision, then constructs the + subclass directly. We avoid ``from_hp`` because it quantizes in the + input dtype (bf16), which loses precision for small-magnitude weights. + + Sub-byte data (e.g. 6-bit) is still stored unpacked in an int8 container; + ``target_dtype`` records the true bit width for the export/runtime path. + """ + from torchao.quantization import IntxUnpackedToInt8Tensor + + qmin = -(1 << (config.bits - 1)) + qmax = (1 << (config.bits - 1)) - 1 + target_dtype = getattr(torch, f"int{config.bits}") + + if config.method == "hqq": + if not config.symmetric: + raise ValueError( + "intx HQQ only supports symmetric quantization " + "(HQQ_SCALE_ONLY). Use method='min_max' for asymmetric intx." + ) + w2d = weight.float().reshape(-1, weight.shape[-1]) + qdata, scale = _quantize_scale_only_hqq(w2d, config.group_size, qmin, qmax) + qdata = qdata.reshape(weight.shape) + scale = scale.reshape(weight.shape[0], -1) + zero_point = torch.zeros_like(scale, dtype=torch.int8) + else: + qdata, scale, zero_point = _quantize_affine_min_max( + weight, + config.group_size, + qmin, + qmax, + config.symmetric, + zero_point_dtype=torch.int8, + ) + N, n_groups = weight.shape[0], weight.shape[-1] // config.group_size + scale = scale.reshape(N, n_groups) + zero_point = zero_point.reshape(N, n_groups) + + return IntxUnpackedToInt8Tensor( + qdata=qdata, + scale=scale, + zero_point=zero_point, + target_dtype=target_dtype, + block_size=(1, config.group_size), + dtype=torch.bfloat16, + activation_quantization=None, + ) + + +def quantize_weight(weight: torch.Tensor, config: QuantConfig) -> torch.Tensor: + """Quantize ``weight`` to a torchao tensor subclass. + + Returns ``Int4Tensor`` for 4-bit or ``IntxUnpackedToInt8Tensor`` for 5-, + 6-, and 8-bit. + """ + if config.bits in (5, 6, 8): + return _to_intx_tensor(weight, config) + + if config.bits != 4: + raise ValueError(f"Unsupported bits={config.bits}") + + if config.method == "min_max": + int_data, scale, zero_point = _quantize_min_max(weight, config) + elif config.method == "hqq": + if config.symmetric: + int_data, scale, zero_point = _quantize_hqq_symmetric(weight, config) + else: + int_data, scale, zero_point = _quantize_hqq_asymmetric(weight, config) + else: + raise ValueError( + f"Unknown quantization method: {config.method!r}. " + f"Supported: 'min_max', 'hqq'." + ) + + return _to_int4_tensor(int_data, scale, zero_point, config) + + +def dequantize_weight( + weight: torch.Tensor, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Dequantize a torchao quantized tensor back to a dense tensor.""" + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + from torchao.quantization import IntxUnpackedToInt8Tensor + from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor + + if isinstance(weight, Int4Tensor): + return ExportableInt4Tensor.from_int4_tensor(weight).dequantize(dtype) + + if isinstance(weight, ExportableInt4Tensor): + return weight.dequantize(dtype) + + if isinstance(weight, IntxUnpackedToInt8Tensor): + return weight.to(dtype).dequantize() + + # ExportableGGUFTensor (native GGUF Q4_K/Q6_K) carries its own gguf-package + # dequant. The tied CUDA token embedding keeps the raw GGUF tensor and is + # dequantized to bf16 here for the gather. Imported lazily to avoid a hard + # extension/llm dependency. + from executorch.extension.llm.export.gguf import ExportableGGUFTensor + + if isinstance(weight, ExportableGGUFTensor): + return weight.dequantize(dtype) + + # CudaDp4aPlanarInt6Tensor (GGUF Q6_K on CUDA) carries its own dequant + # (symmetric, ql/qh split bit-planes). Imported lazily to avoid a hard + # backends/cuda dependency. + from executorch.backends.cuda.dp4a_planar_int6_tensor import ( + CudaDp4aPlanarInt6Tensor, + ) + + if isinstance(weight, CudaDp4aPlanarInt6Tensor): + return weight.dequantize(dtype) + + raise TypeError(f"Cannot dequantize {type(weight).__name__}") + + +# --------------------------------------------------------------------------- +# Per-weight streaming quantization + + +def quantize_stream( + pairs: Iterable[tuple[str, torch.Tensor]], + recipe: QuantRecipe, + dtype: torch.dtype = torch.bfloat16, +) -> Iterator[tuple[str, torch.Tensor]]: + """Lazy, per-weight dual of :func:`quantize_model`. + + For each ``(fqn, tensor)`` pair, apply the recipe -- quantize to a torchao + subclass when ``recipe.get_config(fqn)`` matches, else keep the tensor -- then + :func:`maybe_cast` the result to ``dtype`` (a no-op for non-float, already + matching, or non-castable values, so a routed int buffer is left intact and a + quantized subclass only has its dequantized output dtype re-stamped). + Model-free -- the caller owns iteration order and decides what to feed in + (typically a module's parameters). + + This is the *serialization* form: int4 weights stay as torchao-native + ``Int4Tensor`` (a no-op under ``maybe_cast``, so int4 is always bf16 here), + which round-trips through torchao's safetensors allowlist and is what + ``quantize_and_save`` writes to disk. The ``dtype`` re-stamp for int4 happens + later at the convert boundary (:func:`load_checkpoint` or + :func:`quantize_model`, which wrap it as an ``ExportableInt4Tensor`` first). + """ + for fqn, tensor in pairs: + config = recipe.get_config(fqn) + value = tensor if config is None else quantize_weight(tensor, config) + yield fqn, maybe_cast(value, dtype) + + +# --------------------------------------------------------------------------- +# Per-model quantization + + +def quantize_model( + model: nn.Module, + recipe: QuantRecipe, + *, + convert: Convert = to_default, + dtype: torch.dtype = torch.bfloat16, + verbose: bool = False, +) -> dict[str, torch.Tensor]: + """Walk model parameters + persistent buffers, apply recipe. + + Returns a single state dict for *in-memory* model use (the dual of + :func:`load_checkpoint`). Parameters run through :func:`quantize_stream`, + then each value is ``convert``-ed (default ``to_default``: + ``Int4Tensor -> ExportableInt4Tensor``) and :func:`maybe_cast` to ``dtype`` + -- convert-then-cast, so a raw int4 weight is wrapped as an + ``ExportableInt4Tensor`` (which ``maybe_cast`` *can* re-stamp, e.g. to fp16 + for MLX) and comes out in the export-canonical form the backend packers + consume. Contrast with :func:`quantize_stream`, whose output stays + torchao-native for serialization. + + Non-persistent buffers (KV cache, RoPE tables) are excluded; persistent + buffers are passed through unchanged (never converted or cast -- they may be + non-float). + """ + state: dict[str, torch.Tensor] = {} + persistent_keys = set(model.state_dict().keys()) + + n_params = sum(1 for _ in model.named_parameters()) + params = ((fqn, param.data) for fqn, param in model.named_parameters()) + for i, (fqn, qt) in enumerate(quantize_stream(params, recipe, dtype)): + state[fqn] = maybe_cast(convert(fqn, qt), dtype) + if verbose: + print(f" Quantized {i + 1}/{n_params}: {fqn}", end="\r") + if verbose: + print() + + for fqn, buf in model.named_buffers(): + if fqn in persistent_keys and fqn not in state: + state[fqn] = buf.data + + return state diff --git a/extension/llm/export/quant/recipe.py b/extension/llm/export/quant/recipe.py new file mode 100644 index 00000000000..4c08faa3ed2 --- /dev/null +++ b/extension/llm/export/quant/recipe.py @@ -0,0 +1,62 @@ +# 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. + +"""Quantization recipe: declares what to quantize and how. + +A ``QuantRecipe`` is an ordered list of ``QuantRule`` objects matched against +weight FQNs. First match wins. The recipe says nothing about packing format, +tensor subclass, or target backend. +""" + +import re +from dataclasses import dataclass, field +from typing import Optional + + +@dataclass(frozen=True) +class QuantConfig: + """Per-weight quantization parameters (quantization-time only). + + Not stored in the serialized checkpoint — torchao tensor subclasses + carry their own metadata. This is purely for driving ``quantize_weight``. + """ + + bits: int # 4, 6, or 8 + group_size: int + symmetric: bool # True = no zero point + method: str # "min_max" | "hqq" + + +@dataclass +class QuantRule: + """A single recipe rule: regex pattern + config + optional layer filter.""" + + pattern: str # regex matched against weight FQN + config: Optional[QuantConfig] # None = skip (leave unquantized) + layers: Optional[set[int]] = field(default=None, repr=False) # None = all layers + + +@dataclass +class QuantRecipe: + """Ordered list of rules. First match wins.""" + + rules: list[QuantRule] + + def get_config(self, fqn: str) -> Optional[QuantConfig]: + """Return the ``QuantConfig`` for a weight FQN, or ``None`` to skip.""" + layer_idx = self._extract_layer_idx(fqn) + for rule in self.rules: + if rule.layers is not None: + if layer_idx is None or layer_idx not in rule.layers: + continue + if re.fullmatch(rule.pattern, fqn): + return rule.config + return None + + @staticmethod + def _extract_layer_idx(fqn: str) -> Optional[int]: + m = re.search(r"layers\.(\d+)\.", fqn) + return int(m.group(1)) if m else None diff --git a/extension/llm/export/quant/test/BUCK b/extension/llm/export/quant/test/BUCK new file mode 100644 index 00000000000..8b5b0581a9a --- /dev/null +++ b/extension/llm/export/quant/test/BUCK @@ -0,0 +1,52 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +# 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. + +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target(_kind = runtime.python_test, + name = "test_convert", + srcs = ["test_convert.py"], + deps = [ + "//caffe2:torch", + "//executorch/extension/llm/export:int4", + "//executorch/extension/llm/export:load", + "//executorch/extension/llm/export:quant", + ], +) + +fbcode_target(_kind = runtime.python_test, + name = "test_quantize", + srcs = ["test_quantize.py"], + deps = [ + "//caffe2:torch", + "//executorch/extension/llm/export:int4", + "//executorch/extension/llm/export:quant", + "//pytorch/ao:torchao", + "fbsource//third-party/pypi/parameterized:parameterized", + ], +) + +fbcode_target(_kind = runtime.python_test, + name = "test_recipe", + srcs = ["test_recipe.py"], + deps = [ + "//executorch/extension/llm/export:quant", + "fbsource//third-party/pypi/parameterized:parameterized", + ], +) + +fbcode_target(_kind = runtime.python_test, + name = "test_safetensors_roundtrip", + srcs = ["test_safetensors_roundtrip.py"], + deps = [ + "//caffe2:torch", + "//pytorch/ao:torchao", + "fbsource//third-party/pypi/safetensors:safetensors", + ], +) diff --git a/extension/llm/export/quant/test/test_convert.py b/extension/llm/export/quant/test/test_convert.py new file mode 100644 index 00000000000..f392e57dbdb --- /dev/null +++ b/extension/llm/export/quant/test/test_convert.py @@ -0,0 +1,362 @@ +# 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. + +"""Unit tests for the weight-conversion helpers (``convert.py``) and the +model-assignment helpers (``assign_one`` / ``assign_state_dict`` in ``load.py``). +No hardware required. +""" + +import unittest + +import torch +import torch.nn as nn + +from executorch.extension.llm.export.load import assign_one, assign_state_dict +from executorch.extension.llm.export.quant.convert import ( + fuse_along_output, + to_default, + to_exportable, +) +from executorch.extension.llm.export.quant.quantize import ( + dequantize_weight, + quantize_weight, +) +from executorch.extension.llm.export.quant.recipe import QuantConfig + + +class TestConvertLinear(unittest.TestCase): + def test_int4_wraps_exportable(self): + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + + module = nn.Linear(128, 64, bias=False) + config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") + w = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), config) + assign_one(module, "weight", to_exportable("weight", w)) + + self.assertIsInstance(module.weight.data, ExportableInt4Tensor) + self.assertEqual(module.weight.shape, torch.Size([64, 128])) + self.assertFalse(module.weight.requires_grad) + + def test_int8_passes_through(self): + from torchao.quantization import IntxUnpackedToInt8Tensor + + module = nn.Linear(128, 64, bias=False) + config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") + w = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), config) + self.assertIsInstance(w, IntxUnpackedToInt8Tensor) + assign_one(module, "weight", to_exportable("weight", w)) + + self.assertIsInstance(module.weight.data, IntxUnpackedToInt8Tensor) + self.assertEqual(module.weight.shape, torch.Size([64, 128])) + + def test_int8_coarse_passes_through(self): + """A coarse group_size passes through unchanged. + + Regrouping to a legal group_size happens in the backend pattern + handlers at export time, so conversion leaves block_size untouched. + """ + torch.manual_seed(0) + weight = torch.randn(64, 256, dtype=torch.bfloat16) + config = QuantConfig(bits=8, group_size=256, symmetric=True, method="min_max") + w = quantize_weight(weight, config) + before = dequantize_weight(w, torch.float32) + + module = nn.Linear(256, 64, bias=False) + assign_one(module, "weight", to_exportable("weight", w)) + + self.assertEqual(module.weight.data.block_size, (1, 256)) + after = dequantize_weight(module.weight.data, torch.float32) + self.assertTrue( + torch.allclose(before, after, atol=1e-5), + f"max diff: {(before - after).abs().max():.6g}", + ) + + +class TestConvertLinearGroupSize16(unittest.TestCase): + """Converting group_size=16 weights (GGUF Q6_K) preserves semantics.""" + + def _make_gs16_tensor(self, N=64, K=128): + from torchao.quantization import IntxUnpackedToInt8Tensor + + return IntxUnpackedToInt8Tensor( + qdata=torch.randint(-32, 31, (N, K), dtype=torch.int8), + scale=torch.randn(N, K // 16, dtype=torch.bfloat16), + zero_point=torch.zeros(N, K // 16, dtype=torch.int8), + target_dtype=torch.int8, + block_size=(1, 16), + dtype=torch.bfloat16, + activation_quantization=None, + ) + + def test_dequant_preserves_values(self): + """Conversion preserves the dequantized weight values.""" + w = self._make_gs16_tensor(64, 128) + before = dequantize_weight(w, torch.float32) + + module = nn.Linear(128, 64, bias=False) + assign_one(module, "weight", to_exportable("weight", w)) + after = dequantize_weight(module.weight.data, torch.float32) + + self.assertTrue( + torch.allclose(before, after, atol=1e-5), + f"max diff: {(before - after).abs().max():.6g}", + ) + + def test_forward_produces_valid_output(self): + """Converted gs=16 weight produces finite output in a linear forward.""" + w = self._make_gs16_tensor(64, 128) + module = nn.Linear(128, 64, bias=False) + assign_one(module, "weight", to_exportable("weight", w)) + + x = torch.randn(1, 128, dtype=torch.bfloat16) + out = torch.nn.functional.linear(x, module.weight.data.dequantize()) + self.assertEqual(out.shape, torch.Size([1, 64])) + self.assertFalse(torch.isnan(out).any()) + + +class TestConvertEmbedding(unittest.TestCase): + def test_compatible_passes_through(self): + module = nn.Embedding(100, 64) + config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") + w = quantize_weight(torch.randn(100, 64, dtype=torch.bfloat16), config) + assign_one(module, "weight", to_exportable("weight", w)) + self.assertEqual(module.weight.shape, torch.Size([100, 64])) + + def test_per_axis_passes_through(self): + module = nn.Embedding(50, 256) + config = QuantConfig(bits=8, group_size=256, symmetric=True, method="min_max") + w = quantize_weight(torch.randn(50, 256, dtype=torch.bfloat16), config) + assign_one(module, "weight", to_exportable("weight", w)) + self.assertEqual(module.weight.shape, torch.Size([50, 256])) + # Regrouping happens in the backend handlers at export time, not here. + self.assertEqual(module.weight.data.block_size, (1, 256)) + + def test_int4_wraps_exportable(self): + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + + module = nn.Embedding(100, 64) + config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") + w = quantize_weight(torch.randn(100, 64, dtype=torch.bfloat16), config) + assign_one(module, "weight", to_exportable("weight", w)) + self.assertIsInstance(module.weight.data, ExportableInt4Tensor) + self.assertEqual(module.weight.shape, torch.Size([100, 64])) + + +class TestConvertIntx4SpaceVsExportable(unittest.TestCase): + """to_exportable leaves 4-bit intx as-is; to_default repacks to int4.""" + + def _make_intx4(self, N=64, K=128, gs=32): + from torchao.quantization import IntxUnpackedToInt8Tensor + + return IntxUnpackedToInt8Tensor( + qdata=torch.randint(-8, 7, (N, K), dtype=torch.int8), + scale=torch.randn(N, K // gs, dtype=torch.bfloat16), + zero_point=torch.zeros(N, K // gs, dtype=torch.int8), + target_dtype=torch.int4, + block_size=(1, gs), + dtype=torch.bfloat16, + activation_quantization=None, + ) + + def test_exportable_leaves_intx4_unchanged(self): + from torchao.quantization import IntxUnpackedToInt8Tensor + + w = self._make_intx4() + module = nn.Linear(128, 64, bias=False) + assign_one(module, "weight", to_exportable("weight", w)) + self.assertIsInstance(module.weight.data, IntxUnpackedToInt8Tensor) + + def test_space_repacks_intx4(self): + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + + w = self._make_intx4() + module = nn.Linear(128, 64, bias=False) + assign_one(module, "weight", to_default("weight", w)) + self.assertIsInstance(module.weight.data, ExportableInt4Tensor) + self.assertEqual(module.weight.shape, torch.Size([64, 128])) + + def test_space_convert_via_assign_state_dict(self): + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + + w = self._make_intx4() + with torch.device("meta"): + model = nn.ModuleDict({"q_proj": nn.Linear(128, 64, bias=False)}) + assign_state_dict(model, {"q_proj.weight": w}, convert=to_default) + self.assertIsInstance(model.q_proj.weight.data, ExportableInt4Tensor) + + +class TestAssignStateDict(unittest.TestCase): + def test_mixed_precision(self): + q4 = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") + q8 = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") + w4 = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), q4) + w8 = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), q8) + + state_dict = { + "q_proj.weight": w4, + "v_proj.weight": w8, + "norm.weight": torch.randn(64, dtype=torch.bfloat16), + } + + with torch.device("meta"): + model = nn.ModuleDict( + { + "q_proj": nn.Linear(128, 64, bias=False), + "v_proj": nn.Linear(128, 64, bias=False), + "norm": nn.LayerNorm(64, bias=False), + } + ) + assign_state_dict(model, state_dict) + + self.assertEqual(model.q_proj.weight.shape, torch.Size([64, 128])) + self.assertEqual(model.v_proj.weight.shape, torch.Size([64, 128])) + self.assertEqual(model.norm.weight.shape, torch.Size([64])) + + def test_default_convert(self): + """assign_state_dict with no convert uses to_exportable.""" + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + + q4 = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") + w4 = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), q4) + state_dict = {"q_proj.weight": w4} + + with torch.device("meta"): + model = nn.ModuleDict({"q_proj": nn.Linear(128, 64, bias=False)}) + assign_state_dict(model, state_dict) + + self.assertIsInstance(model.q_proj.weight.data, ExportableInt4Tensor) + self.assertEqual(model.q_proj.weight.shape, torch.Size([64, 128])) + + def test_missing_weight_detected(self): + q4 = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") + w4 = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), q4) + + with torch.device("meta"): + model = nn.ModuleDict( + { + "a": nn.Linear(128, 64, bias=False), + "b": nn.Linear(128, 64, bias=False), + } + ) + with self.assertRaises(RuntimeError) as ctx: + assign_state_dict(model, {"a.weight": w4}) + self.assertIn("b.weight", str(ctx.exception)) + + +class TestFuseAlongOutput(unittest.TestCase): + """fuse_along_output concatenates weights along the output-channel dim.""" + + @staticmethod + def _int4(rows, cols=128, group_size=32): + config = QuantConfig( + bits=4, group_size=group_size, symmetric=False, method="min_max" + ) + w = torch.randn(rows, cols, dtype=torch.bfloat16) + return to_default("w", quantize_weight(w, config)) + + @staticmethod + def _intx8(rows, cols=128, group_size=32): + config = QuantConfig( + bits=8, group_size=group_size, symmetric=False, method="min_max" + ) + return quantize_weight(torch.randn(rows, cols, dtype=torch.bfloat16), config) + + @staticmethod + def _gguf(rows): + from executorch.extension.llm.export.gguf import ( + _Q4_K_BLOCK_BYTES, + ExportableGGUFTensor, + ) + + raw = torch.randint(0, 256, (rows, _Q4_K_BLOCK_BYTES), dtype=torch.uint8) + return ExportableGGUFTensor(raw, "q4_k", torch.bfloat16) + + def test_plain_tensor(self): + a = torch.randn(8, 16) + b = torch.randn(24, 16) + fused = fuse_along_output([a, b]) + self.assertEqual(fused.shape, torch.Size([32, 16])) + self.assertTrue(torch.equal(fused, torch.cat([a, b], dim=0))) + + def test_single_tensor_returned_unchanged(self): + a = self._int4(8) + self.assertIs(fuse_along_output([a]), a) + + def test_empty_raises(self): + with self.assertRaises(ValueError): + fuse_along_output([]) + + def test_int4_fuses_exactly(self): + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + + a = self._int4(8) + b = self._int4(24) + fused = fuse_along_output([a, b]) + + self.assertIsInstance(fused, ExportableInt4Tensor) + self.assertEqual(fused.shape, torch.Size([32, 128])) + self.assertEqual(fused.group_size, a.group_size) + self.assertEqual(fused.orig_dtype, a.orig_dtype) + # qdata is N-major (dim 0); scale/zero_point are transposed (N on dim 1). + self.assertTrue(torch.equal(fused.qdata, torch.cat([a.qdata, b.qdata], dim=0))) + self.assertTrue(torch.equal(fused.scale, torch.cat([a.scale, b.scale], dim=1))) + self.assertTrue( + torch.equal( + fused.zero_point, torch.cat([a.zero_point, b.zero_point], dim=1) + ) + ) + ref = torch.cat([dequantize_weight(a), dequantize_weight(b)], dim=0) + self.assertTrue(torch.allclose(dequantize_weight(fused), ref)) + + def test_intx_int8_fuses_exactly(self): + from torchao.quantization import IntxUnpackedToInt8Tensor + + a = self._intx8(8) + b = self._intx8(24) + fused = fuse_along_output([a, b]) + + self.assertIsInstance(fused, IntxUnpackedToInt8Tensor) + self.assertEqual(fused.shape, torch.Size([32, 128])) + self.assertEqual(list(fused.block_size), list(a.block_size)) + self.assertEqual(fused.target_dtype, a.target_dtype) + # Every packed field is N-major (dim 0) for IntxUnpackedToInt8Tensor. + self.assertTrue(torch.equal(fused.qdata, torch.cat([a.qdata, b.qdata], dim=0))) + self.assertTrue(torch.equal(fused.scale, torch.cat([a.scale, b.scale], dim=0))) + ref = torch.cat([dequantize_weight(a), dequantize_weight(b)], dim=0) + self.assertTrue(torch.allclose(dequantize_weight(fused), ref)) + + def test_gguf_fuses_exactly(self): + from executorch.extension.llm.export.gguf import ExportableGGUFTensor + + a = self._gguf(2) + b = self._gguf(3) + fused = fuse_along_output([a, b]) + + self.assertIsInstance(fused, ExportableGGUFTensor) + self.assertEqual(fused.shape, torch.Size([5, 256])) + self.assertEqual(fused.ggml_type, "q4_k") + self.assertEqual(fused.orig_dtype, torch.bfloat16) + # Each row is an independent super-block, so exact fusion == cat of raw + # bytes (random bytes decode to NaN scales, so don't compare dequant). + self.assertTrue(torch.equal(fused.raw, torch.cat([a.raw, b.raw], dim=0))) + + def test_type_mismatch_raises(self): + with self.assertRaises(TypeError): + fuse_along_output([self._int4(8), self._intx8(8)]) + + def test_plain_quantized_mix_raises(self): + with self.assertRaises(TypeError): + fuse_along_output([torch.randn(8, 128), self._int4(8)]) + + def test_mismatched_quant_param_raises(self): + a = self._int4(8, group_size=32) + b = self._int4(8, group_size=64) + with self.assertRaises(ValueError): + fuse_along_output([a, b]) + + +if __name__ == "__main__": + unittest.main() diff --git a/extension/llm/export/quant/test/test_quantize.py b/extension/llm/export/quant/test/test_quantize.py new file mode 100644 index 00000000000..4688f93a898 --- /dev/null +++ b/extension/llm/export/quant/test/test_quantize.py @@ -0,0 +1,310 @@ +# 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. + +"""Unit tests for quant/quantize.py.""" + +import unittest + +import torch +import torch.nn as nn + +from executorch.extension.llm.export.int4 import ExportableInt4Tensor +from executorch.extension.llm.export.quant.convert import to_default +from executorch.extension.llm.export.quant.quantize import ( + dequantize_weight, + quantize_model, + quantize_stream, + quantize_weight, +) +from executorch.extension.llm.export.quant.recipe import ( + QuantConfig, + QuantRecipe, + QuantRule, +) +from parameterized import parameterized +from torchao.quantization import IntxUnpackedToInt8Tensor +from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor + + +class TestQuantizeWeight(unittest.TestCase): + @parameterized.expand( + [ + ("4bit_asym", 4, 32, False), + ("4bit_sym", 4, 32, True), + ("4bit_gs64", 4, 64, False), + ("8bit_sym", 8, 32, True), + ] + ) + def test_output_type(self, _name, bits, gs, sym): + config = QuantConfig(bits=bits, group_size=gs, symmetric=sym, method="min_max") + result = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), config) + if bits == 4: + self.assertIsInstance(result, Int4Tensor) + self.assertEqual(result.shape, torch.Size([64, 128])) + else: + self.assertIsInstance(result, IntxUnpackedToInt8Tensor) + self.assertEqual(result.shape, torch.Size([64, 128])) + + def test_quantize_dequantize_roundtrip(self): + torch.manual_seed(0) + weight = torch.randn(64, 128, dtype=torch.bfloat16) + config = QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max") + q = quantize_weight(weight, config) + dequant = dequantize_weight(q, dtype=torch.bfloat16) + rel_error = ( + dequant.float() - weight.float() + ).abs().mean() / weight.float().abs().mean() + self.assertLess(rel_error.item(), 0.15) + + def test_dequantize_output_dtype(self): + config = QuantConfig(bits=4, group_size=32, symmetric=False, method="min_max") + q = quantize_weight(torch.randn(32, 64, dtype=torch.bfloat16), config) + self.assertEqual(dequantize_weight(q, torch.float32).dtype, torch.float32) + self.assertEqual(dequantize_weight(q, torch.bfloat16).dtype, torch.bfloat16) + + def test_dequantize_symmetric_4bit(self): + torch.manual_seed(1) + weight = torch.randn(32, 64, dtype=torch.bfloat16) + config = QuantConfig(bits=4, group_size=32, symmetric=True, method="min_max") + q = quantize_weight(weight, config) + dequant = dequantize_weight(q) + self.assertEqual(dequant.shape, (32, 64)) + rel_error = ( + dequant.float() - weight.float() + ).abs().mean() / weight.float().abs().mean() + self.assertLess(rel_error.item(), 0.15) + + def test_dequantize_int8(self): + torch.manual_seed(2) + weight = torch.randn(32, 64, dtype=torch.bfloat16) + config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") + q = quantize_weight(weight, config) + dequant = dequantize_weight(q, dtype=torch.bfloat16) + rel_error = ( + dequant.float() - weight.float() + ).abs().mean() / weight.float().abs().mean() + self.assertLess(rel_error.item(), 0.02) + + def test_int8_small_weights_bf16_precision(self): + """INT8 quantization of small bf16 weights must use full int8 range. + + Regression: IntxUnpackedToInt8Tensor.from_hp quantizes in bf16, + which collapses per-group scales to a single value for weights + with abs_mean ~0.01 (small-magnitude weights). Our _to_intx_tensor + casts to float32 first to avoid this. + """ + torch.manual_seed(42) + weight = torch.randn(64, 128, dtype=torch.bfloat16) * 0.01 + config = QuantConfig(bits=8, group_size=32, symmetric=True, method="min_max") + q = quantize_weight(weight, config) + dequant = dequantize_weight(q, dtype=torch.bfloat16) + rel_error = ( + dequant.float() - weight.float() + ).abs().mean() / weight.float().abs().mean() + self.assertLess(rel_error.item(), 0.02) + + def test_dequantize_int8_asymmetric(self): + torch.manual_seed(3) + weight = torch.randn(32, 64, dtype=torch.bfloat16) + config = QuantConfig(bits=8, group_size=32, symmetric=False, method="min_max") + q = quantize_weight(weight, config) + dequant = dequantize_weight(q, dtype=torch.bfloat16) + rel_error = ( + dequant.float() - weight.float() + ).abs().mean() / weight.float().abs().mean() + self.assertLess(rel_error.item(), 0.02) + + def test_int8_per_axis(self): + """Per-axis (group_size == K) used for embeddings.""" + weight = torch.randn(256, 64, dtype=torch.bfloat16) + config = QuantConfig(bits=8, group_size=64, symmetric=True, method="min_max") + q = quantize_weight(weight, config) + self.assertIsInstance(q, IntxUnpackedToInt8Tensor) + dequant = dequantize_weight(q, dtype=torch.bfloat16) + rel_error = ( + dequant.float() - weight.float() + ).abs().mean() / weight.float().abs().mean() + self.assertLess(rel_error.item(), 0.01) + + @parameterized.expand( + [ + ("unknown_method", QuantConfig(4, 32, False, "bogus"), "bogus"), + ("unsupported_bits", QuantConfig(3, 32, False, "min_max"), None), + ("hqq_8bit_asym", QuantConfig(8, 32, False, "hqq"), "symmetric"), + ] + ) + def test_invalid_config_raises(self, _name, config, expected_substr): + with self.assertRaises(ValueError) as ctx: + quantize_weight(torch.randn(32, 64), config) + if expected_substr: + self.assertIn(expected_substr, str(ctx.exception)) + + +class TestQuantizeWeightHQQ(unittest.TestCase): + def setUp(self): + if not torch.cuda.is_available(): + self.skipTest("CUDA required for HQQ") + + def test_quantize_dequantize_roundtrip(self): + torch.manual_seed(0) + weight = torch.randn(64, 128, dtype=torch.bfloat16, device="cuda") + config = QuantConfig(bits=4, group_size=32, symmetric=False, method="hqq") + q = quantize_weight(weight, config) + dequant = dequantize_weight(q, dtype=torch.bfloat16).cpu() + rel_error = ( + dequant.float() - weight.cpu().float() + ).abs().mean() / weight.cpu().float().abs().mean() + self.assertLess(rel_error.item(), 0.15) + + def test_symmetric_scale_only(self): + config = QuantConfig(bits=4, group_size=32, symmetric=True, method="hqq") + q = quantize_weight(torch.randn(64, 128, dtype=torch.bfloat16), config) + self.assertIsInstance(q, Int4Tensor) + + def test_int8_hqq_roundtrip(self): + torch.manual_seed(0) + weight = torch.randn(64, 128, dtype=torch.bfloat16) + config = QuantConfig(bits=8, group_size=32, symmetric=True, method="hqq") + q = quantize_weight(weight, config) + self.assertIsInstance(q, IntxUnpackedToInt8Tensor) + dequant = dequantize_weight(q, dtype=torch.bfloat16) + rel_error = ( + dequant.float() - weight.float() + ).abs().mean() / weight.float().abs().mean() + self.assertLess(rel_error.item(), 0.02) + + +class TestQuantizeModel(unittest.TestCase): + def test_applies_recipe(self): + model = nn.ModuleDict( + { + "embed": nn.Embedding(32, 16), + "proj": nn.Linear(16, 32, bias=False), + "norm": nn.LayerNorm(32), + } + ) + model.to(dtype=torch.bfloat16) + for p in model.parameters(): + p.data.normal_(0, 0.02) + + recipe = QuantRecipe( + rules=[ + QuantRule(r"embed\.weight", None), + QuantRule(r"norm\.weight", None), + QuantRule(r".*\.weight", QuantConfig(4, 16, False, "min_max")), + ] + ) + state = quantize_model(model, recipe) + + # quantize_model returns the in-memory export form: int4 is wrapped as + # ExportableInt4Tensor (not raw torchao Int4Tensor). + self.assertIsInstance(state["proj.weight"], ExportableInt4Tensor) + self.assertIs(type(state["embed.weight"]), torch.Tensor) + self.assertIs(type(state["norm.weight"]), torch.Tensor) + + def test_persistent_buffers_included(self): + model = nn.Module() + model.weight = nn.Parameter(torch.randn(16, 32, dtype=torch.bfloat16)) + model.register_buffer("scalar", torch.ones(1)) + model.register_buffer("temp", torch.zeros(4), persistent=False) + + recipe = QuantRecipe(rules=[QuantRule(r".*", None)]) + state = quantize_model(model, recipe) + + self.assertIn("scalar", state) + self.assertNotIn("temp", state) + + def test_unquantized_cast_to_dtype(self): + model = nn.ModuleDict({"proj": nn.Linear(16, 8, bias=False)}) + model.proj.weight.data = torch.randn(8, 16, dtype=torch.float32) + + recipe = QuantRecipe(rules=[QuantRule(r".*", None)]) + state = quantize_model(model, recipe, dtype=torch.float16) + + self.assertEqual(state["proj.weight"].dtype, torch.float16) + + def test_empty_model(self): + state = quantize_model(nn.Module(), QuantRecipe(rules=[])) + self.assertEqual(len(state), 0) + + +class TestQuantizeStream(unittest.TestCase): + """quantize_stream is the lazy, per-weight dual of quantize_model.""" + + def _model_and_recipe(self): + model = nn.ModuleDict( + { + "embed": nn.Embedding(32, 16), + "proj": nn.Linear(16, 32, bias=False), + "norm": nn.LayerNorm(32), + } + ) + model.to(dtype=torch.bfloat16) + for p in model.parameters(): + p.data.normal_(0, 0.02) + recipe = QuantRecipe( + rules=[ + QuantRule(r"embed\.weight", None), + QuantRule(r"norm\.weight", None), + QuantRule(r".*\.weight", QuantConfig(4, 16, False, "min_max")), + ] + ) + return model, recipe + + def test_is_lazy(self): + """Returns an iterator; nothing is quantized until consumed.""" + _model, recipe = self._model_and_recipe() + gen = quantize_stream(iter([]), recipe) + self.assertTrue(hasattr(gen, "__next__")) + with self.assertRaises(StopIteration): + next(gen) + + def test_matches_quantize_model_params(self): + """Streamed params reconstruct quantize_model's parameter entries. + + quantize_stream yields the serialization form (torchao-native + ``Int4Tensor``); quantize_model yields the in-memory form + (``ExportableInt4Tensor``). They match once the stream output is passed + through the same ``to_default`` convert quantize_model applies. + """ + model, recipe = self._model_and_recipe() + ref = quantize_model(model, recipe) + params = ((n, p.data) for n, p in model.named_parameters()) + streamed = {n: to_default(n, v) for n, v in quantize_stream(params, recipe)} + + self.assertEqual(set(streamed), {n for n, _ in model.named_parameters()}) + for name, val in streamed.items(): + r = ref[name] + self.assertIs(type(val), type(r)) + data_names = getattr(val, "tensor_data_names", None) + if data_names: # quantized subclass: compare payload bit-exactly + for n in data_names: + self.assertTrue(torch.equal(getattr(val, n), getattr(r, n))) + else: + self.assertTrue(torch.equal(val, r)) + + def test_unmatched_cast_to_dtype(self): + recipe = QuantRecipe(rules=[QuantRule(r".*", None)]) + out = dict( + quantize_stream( + [("x", torch.randn(4, 4, dtype=torch.float32))], + recipe, + dtype=torch.float16, + ) + ) + self.assertEqual(out["x"].dtype, torch.float16) + + def test_non_float_passthrough_uncast(self): + """A routed non-float tensor (e.g. an int position table) is not cast.""" + recipe = QuantRecipe(rules=[QuantRule(r".*", None)]) + positions = torch.arange(8, dtype=torch.long) + out = dict(quantize_stream([("pos", positions)], recipe, dtype=torch.bfloat16)) + self.assertEqual(out["pos"].dtype, torch.long) + self.assertTrue(torch.equal(out["pos"], positions)) + + +if __name__ == "__main__": + unittest.main() diff --git a/extension/llm/export/quant/test/test_recipe.py b/extension/llm/export/quant/test/test_recipe.py new file mode 100644 index 00000000000..5fe6f3be394 --- /dev/null +++ b/extension/llm/export/quant/test/test_recipe.py @@ -0,0 +1,118 @@ +# 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. + +"""Unit tests for quant/recipe.py. CPU only — no CUDA, no model, no torchao.""" + +import unittest + +from executorch.extension.llm.export.quant.recipe import ( + QuantConfig, + QuantRecipe, + QuantRule, +) +from parameterized import parameterized + +_Q4 = QuantConfig(4, 32, True, "min_max") +_Q8 = QuantConfig(8, 32, True, "min_max") + + +class TestQuantRecipeGetConfig(unittest.TestCase): + """Tests for ``QuantRecipe.get_config`` — the core matching logic.""" + + @parameterized.expand( + [ + ( + "first_match_wins", + [QuantRule(r".*v_proj\.weight", _Q8), QuantRule(r".*\.weight", _Q4)], + "layers.0.self_attn.v_proj.weight", + 8, + ), + ( + "fallthrough_to_catchall", + [QuantRule(r".*v_proj\.weight", _Q8), QuantRule(r".*\.weight", _Q4)], + "layers.0.self_attn.q_proj.weight", + 4, + ), + ( + "none_rule_skips", + [ + QuantRule(r"embed_tokens\.weight", None), + QuantRule(r".*\.weight", _Q4), + ], + "embed_tokens.weight", + None, + ), + ( + "unmatched_returns_none", + [QuantRule(r"foo", _Q4)], + "bar.weight", + None, + ), + ( + "empty_recipe", + [], + "anything", + None, + ), + ( + "fullmatch_not_partial", + [QuantRule(r"foo", _Q4)], + "foo.bar", + None, + ), + ( + "fullmatch_exact", + [QuantRule(r"foo", _Q4)], + "foo", + 4, + ), + ] + ) + def test_get_config(self, _name, rules, fqn, expected_bits): + recipe = QuantRecipe(rules=rules) + config = recipe.get_config(fqn) + if expected_bits is None: + self.assertIsNone(config) + else: + self.assertEqual(config.bits, expected_bits) + + +class TestQuantRecipeLayerFilter(unittest.TestCase): + """Tests for the ``layers`` field on ``QuantRule``.""" + + def test_layer_filter(self): + edge = set(range(5)) | set(range(55, 60)) + recipe = QuantRecipe( + rules=[ + QuantRule(r".*norm\.weight", None), + QuantRule(r".*\.(v_proj|down_proj)\.weight", _Q8, layers=edge), + QuantRule(r".*\.weight", _Q4), + ] + ) + # Edge v_proj → 8-bit + self.assertEqual(recipe.get_config("layers.0.self_attn.v_proj.weight").bits, 8) + self.assertEqual(recipe.get_config("layers.58.self_attn.v_proj.weight").bits, 8) + # Middle v_proj → falls through → 4-bit + self.assertEqual(recipe.get_config("layers.30.self_attn.v_proj.weight").bits, 4) + # q_proj always 4-bit + self.assertEqual(recipe.get_config("layers.0.self_attn.q_proj.weight").bits, 4) + # Non-layer FQN skips layer-filtered rule, hits catch-all + self.assertEqual(recipe.get_config("lm_head.weight").bits, 4) + + def test_layer_filter_with_none_config(self): + """Skip rule scoped to specific layers.""" + recipe = QuantRecipe( + rules=[ + QuantRule(r".*\.weight", None, layers={0}), + QuantRule(r".*\.weight", _Q4), + ] + ) + self.assertIsNone(recipe.get_config("layers.0.mlp.gate_proj.weight")) + self.assertEqual(recipe.get_config("layers.1.mlp.gate_proj.weight").bits, 4) + + +if __name__ == "__main__": + unittest.main() diff --git a/extension/llm/export/quant/test/test_safetensors_roundtrip.py b/extension/llm/export/quant/test/test_safetensors_roundtrip.py new file mode 100644 index 00000000000..1f4ba044e7b --- /dev/null +++ b/extension/llm/export/quant/test/test_safetensors_roundtrip.py @@ -0,0 +1,143 @@ +# 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. + +"""Smoke tests: torchao subclasses survive safetensors roundtrip.""" + +import os +import tempfile +import unittest + +import torch + +from safetensors import safe_open +from safetensors.torch import save_file +from torchao.prototype.safetensors.safetensors_support import ( + flatten_tensor_state_dict, + unflatten_tensor_state_dict, +) + + +def save(state_dict, path): + tensors_data, metadata = flatten_tensor_state_dict(state_dict) + save_file(tensors_data, path, metadata=metadata) + + +def load(path): + with safe_open(path, framework="pt", device="cpu") as f: + metadata = f.metadata() + tensors = {k: f.get_tensor(k) for k in f.keys()} + result, _ = unflatten_tensor_state_dict(tensors, metadata) + return result + + +from torchao.quantization import IntxUnpackedToInt8Tensor +from torchao.quantization.quantize_.workflows.int4.int4_tensor import Int4Tensor + + +def _make_int4(shape, group_size=32): + """Build a random Int4Tensor.""" + N, K = shape + packed = torch.randint(0, 255, (N, K // 2), dtype=torch.uint8) + scale = torch.randn(K // group_size, N, dtype=torch.bfloat16) + zp = torch.zeros(K // group_size, N, dtype=torch.bfloat16) + return Int4Tensor( + qdata=packed, + scale=scale, + zero_point=zp, + block_size=[1, group_size], + shape=torch.Size([N, K]), + ) + + +def _make_int8(shape, group_size=32): + """Build a random IntxUnpackedToInt8Tensor.""" + N, K = shape + return IntxUnpackedToInt8Tensor( + qdata=torch.randint(-128, 127, (N, K), dtype=torch.int8), + scale=torch.randn(N, K // group_size, dtype=torch.bfloat16), + zero_point=torch.zeros(N, K // group_size, dtype=torch.int8), + target_dtype=torch.int8, + block_size=(1, group_size), + dtype=torch.bfloat16, + activation_quantization=None, + ) + + +class TestSaveLoad(unittest.TestCase): + def test_int4_roundtrip(self): + """Int4Tensor survives save/load.""" + t = _make_int4((64, 128)) + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "m.safetensors") + save({"layer.weight": t}, path) + loaded = load(path) + + self.assertIn("layer.weight", loaded) + self.assertIsInstance(loaded["layer.weight"], Int4Tensor) + self.assertTrue(torch.equal(t.qdata, loaded["layer.weight"].qdata)) + self.assertTrue(torch.equal(t.scale, loaded["layer.weight"].scale)) + + def test_int8_roundtrip(self): + """IntxUnpackedToInt8Tensor survives save/load.""" + t = _make_int8((64, 128)) + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "m.safetensors") + save({"layer.weight": t}, path) + loaded = load(path) + + self.assertIn("layer.weight", loaded) + self.assertIsInstance(loaded["layer.weight"], IntxUnpackedToInt8Tensor) + self.assertTrue(torch.equal(t.qdata, loaded["layer.weight"].qdata)) + + def test_mixed_state_dict(self): + """Mixed Int4 + Int8 + plain tensor roundtrip.""" + state = { + "linear.weight": _make_int4((64, 128)), + "embed.weight": _make_int8((100, 64)), + "norm.weight": torch.randn(64, dtype=torch.bfloat16), + } + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "m.safetensors") + save(state, path) + loaded = load(path) + + self.assertEqual(set(state.keys()), set(loaded.keys())) + self.assertIsInstance(loaded["linear.weight"], Int4Tensor) + self.assertIsInstance(loaded["embed.weight"], IntxUnpackedToInt8Tensor) + self.assertIsInstance(loaded["norm.weight"], torch.Tensor) + self.assertTrue(torch.equal(state["norm.weight"], loaded["norm.weight"])) + + def test_plain_tensor_only(self): + """State dict with only plain tensors roundtrips.""" + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "m.safetensors") + save({"model.norm.weight": torch.randn(64, dtype=torch.bfloat16)}, path) + loaded = load(path) + self.assertIn("model.norm.weight", loaded) + + def test_3d_int4(self): + """3D Int4Tensor (batched/expert weights) roundtrips.""" + # 3D: (num_experts, N, K//2) packed + N, K, gs = 32, 64, 32 + packed = torch.randint(0, 255, (4, N, K // 2), dtype=torch.uint8) + scale = torch.randn(4, K // gs, N, dtype=torch.bfloat16) + zp = torch.zeros(4, K // gs, N, dtype=torch.bfloat16) + t = Int4Tensor( + qdata=packed, + scale=scale, + zero_point=zp, + block_size=[1, 1, gs], + shape=torch.Size([4, N, K]), + ) + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "m.safetensors") + save({"experts.w1": t}, path) + loaded = load(path) + self.assertTrue(torch.equal(t.qdata, loaded["experts.w1"].qdata)) + + +if __name__ == "__main__": + unittest.main() diff --git a/extension/llm/export/test/BUCK b/extension/llm/export/test/BUCK index 5537a8c5f29..9a40024abbe 100644 --- a/extension/llm/export/test/BUCK +++ b/extension/llm/export/test/BUCK @@ -17,3 +17,24 @@ fbcode_target(_kind = runtime.python_test, "//caffe2:torch", ], ) + +fbcode_target(_kind = runtime.python_test, + name = "test_int4", + srcs = ["test_int4.py"], + deps = [ + "//caffe2:torch", + "//executorch/extension/llm/export:int4", + "//pytorch/ao:torchao", + ], +) + +fbcode_target(_kind = runtime.python_test, + name = "test_gguf", + srcs = ["test_gguf.py"], + deps = [ + "//caffe2:torch", + "//executorch/extension/llm/export:gguf", + "fbsource//third-party/pypi/gguf:gguf", + "fbsource//third-party/pypi/numpy:numpy", + ], +) diff --git a/extension/llm/export/test/test_gguf.py b/extension/llm/export/test/test_gguf.py index f35216a5e20..8a5de780892 100644 --- a/extension/llm/export/test/test_gguf.py +++ b/extension/llm/export/test/test_gguf.py @@ -12,7 +12,7 @@ * ``ExportableGGUFTensor.dequantize`` (and the ``torchao::dequantize_gguf`` op, whose eager body uses ``gguf``) reproduces ``gguf.dequantize``; -* our hand-written ``to_int4_tensor`` / ``to_intx_unpacked_to_int8_tensor`` +* our hand-written ``to_exportable_int4_tensor`` / ``to_intx_unpacked_to_int8_tensor`` unpack matches ``gguf.dequantize`` (within bf16 storage tolerance); * using the subclass as a weight dispatches linear/embedding to the fused ops. @@ -90,23 +90,6 @@ def _gguf_ref(raw: torch.Tensor, qtype) -> torch.Tensor: return torch.from_numpy(np.asarray(gguf.dequantize(raw.numpy(), qtype))).float() -def _int4_to_float(w) -> torch.Tensor: - """Dequantize an ``Int4Tensor`` from its stored fields. - - ``Int4Tensor`` has no working ``dequantize()`` on CPU (``aten.dequantize`` is - unimplemented and the linear path needs fbgemm), so reconstruct directly - from its public fields (this still exercises our nibble-packing). - """ - N, K = int(w.shape[0]), int(w.shape[1]) - gs = w.block_size[1] - q = torch.empty(N, K, dtype=torch.float32) - q[:, ::2] = (w.qdata & 0x0F).float() - q[:, 1::2] = (w.qdata >> 4).float() - scale = w.scale.t().float().repeat_interleave(gs, dim=1) - zero = w.zero_point.t().float().repeat_interleave(gs, dim=1) - return scale * (q - zero) - - @unittest.skipUnless(_HAS_GGUF, "gguf package not installed") class TestExportableGGUFTensor(unittest.TestCase): def test_dequantize_matches_gguf(self): @@ -123,6 +106,19 @@ def test_dequantize_matches_gguf(self): # .dequantize() routes through gguf, so it should match exactly. self.assertTrue(torch.equal(mine, ref), f"{qtype}") + def test_to_preserves_quantization(self): + raw = _make_q4k_raw(N=3, nb=2) + t = ExportableGGUFTensor.from_raw(raw, "q4_k", torch.bfloat16) + cast = t.to(torch.float16) + # .to() sets the output dtype without dequantizing: same subclass, the + # raw blocks are untouched, and dequantize() now yields the new dtype. + self.assertIsInstance(cast, ExportableGGUFTensor) + self.assertEqual(cast.dtype, torch.float16) + self.assertTrue(torch.equal(cast.raw, t.raw)) + self.assertEqual(cast.dequantize().dtype, torch.float16) + with self.assertRaises(AssertionError): + t.to(torch.int8) + def test_to_intx_unpacked_matches_reference(self): # Reference is the gguf-package dequant (ExportableGGUFTensor.dequantize); # the Intx tensor's dequantize exercises our unpacking (Q4_K/Q5_K/Q6_K). @@ -249,23 +245,33 @@ def test_iter_gguf_preserves_native_float_dtype(self): self.assertEqual(got["f16"].dtype, torch.float16) self.assertEqual(got["f32"].dtype, torch.float32) - def test_to_int4_tensor_matches_reference(self): + def test_to_exportable_int4_tensor_matches_reference(self): + from executorch.extension.llm.export.int4 import ExportableInt4Tensor + raw = _make_q4k_raw(N=3, nb=2) t = ExportableGGUFTensor.from_raw(raw, "q4_k") - w = t.to_int4_tensor() + w = t.to_exportable_int4_tensor() + self.assertIsInstance(w, ExportableInt4Tensor) self.assertEqual(tuple(w.shape), (3, 512)) - self.assertEqual(list(w.block_size), [1, Q4_K_GROUP_SIZE]) - # Int4Tensor has no CPU dequantize(); reconstruct from its packed fields - # (this still exercises our nibble-packing) against the gguf reference. + self.assertEqual(w.group_size, Q4_K_GROUP_SIZE) self.assertTrue( torch.allclose( - _int4_to_float(w), + w.dequantize(torch.float32), t.dequantize(torch.float32), rtol=1e-2, atol=5e-2, ) ) + def test_to_exportable_int4_tensor_threads_dtype(self): + raw = _make_q4k_raw(N=3, nb=2) + t = ExportableGGUFTensor.from_raw(raw, "q4_k") + w = t.to_exportable_int4_tensor(torch.float16) + self.assertEqual(w.scale.dtype, torch.float16) + self.assertEqual(w.zero_point.dtype, torch.float16) + self.assertEqual(w.orig_dtype, torch.float16) + self.assertEqual(w.dequantize().dtype, torch.float16) + def test_dequantize_gguf_op_matches_reference(self): for ggml_type, make in ( ("q4_k", _make_q4k_raw), diff --git a/extension/llm/export/test/test_int4.py b/extension/llm/export/test/test_int4.py index 9414248d59a..5a71f146c7d 100644 --- a/extension/llm/export/test/test_int4.py +++ b/extension/llm/export/test/test_int4.py @@ -76,6 +76,20 @@ def test_subclass_embedding_dispatches_to_dequant(self): ref = torch.nn.functional.embedding(idx, t.dequantize(torch.bfloat16)) self.assertTrue(torch.equal(out, ref)) + def test_to_preserves_quantization(self): + it, _, _, _ = _make_int4_tensor(N=8, K=64, gs=32) + t = ExportableInt4Tensor.from_int4_tensor(it) + cast = t.to(torch.float16) + # .to() sets the output dtype without dequantizing: same subclass, the + # packed int4 qdata is untouched, scale/output dtype follow the cast. + self.assertIsInstance(cast, ExportableInt4Tensor) + self.assertEqual(cast.dtype, torch.float16) + self.assertTrue(torch.equal(cast.qdata, t.qdata)) + self.assertEqual(cast.scale.dtype, torch.float16) + self.assertEqual(cast.dequantize().dtype, torch.float16) + with self.assertRaises(AssertionError): + t.to(torch.int8) + class TestExportableInt4TensorExport(unittest.TestCase): """Exporting a module whose weight is an ``ExportableInt4Tensor`` should lower