Skip to content

VulkanQuantizer dynamic activation quantization is per tensor, and very lossy on transformer encoders#22431

Description

@msluszniak

馃悰 Describe the bug

VulkanQuantizer's dynamic mode quantizes activations per tensor, which costs about an order of magnitude more accuracy than a per-token scheme on transformer encoders. The code says otherwise:

# backends/vulkan/quantizer/vulkan_quantizer.py, get_symmetric_quantization_config
else:
    # Dynamic quantization: per-token input quantization, no output quantization
    ...
    act_quantization_spec = QuantizationSpec(
        dtype=torch.int8,
        qscheme=torch.per_tensor_affine,   # <- per tensor, not per token
        is_dynamic=True,
        observer_or_fake_quant_ctr=PlaceholderObserver,
    )

The spec matches the kernel: linear_q8ta_q8csw builds QuantizationConfig input_quant_config(8, kPerTensor, {}, false) (QuantizedLinear.cpp), so a single input scale and zero point is what the shader consumes. It is the comment that is wrong, and it sets the wrong expectation for anyone choosing is_dynamic=True.

Measured on two sentence-transformer embedders, mean cosine of the embedding against the fp32 eager model over 8 sentences (all-mpnet-base-v2 / multi-qa-mpnet-base-dot-v1):

config cosine max pairwise-similarity delta
VulkanQuantizer(is_dynamic=True, weight_bits=8) 0.910 / 0.948 0.198 / 0.136
VulkanQuantizer(is_dynamic=True, weight_bits=4) 0.831 / 0.876 0.099 / 0.175
VulkanQuantizer(is_dynamic=False, weight_bits=8) 0.9993 / 0.9993 0.005 / 0.013
torchao Int8DynamicActivationIntxWeightConfig(int8, PerAxis(0)) 0.998 / 0.998 0.011 / 0.021

The last row is the control: dynamic int8 activations are not inherently lossy here. torchao quantizes the activation per row, and lands within 0.002 of fp32 on the same models with the same weight width. The per-tensor scheme moves the pair ranking of the 28 sentence pairs, which is what an embedder is for.

The runtime already has the machinery for per-token activations: linear_dq8ca_q4gsw uses QuantizationConfig input_quant_config(8, kPerChannel, {}, false, true). There is no equivalent for int8 weights, so a linear_dq8ca_q8csw variant would give the quantizer a config worth using on encoder-style models.

Two suggestions, in order of usefulness:

  1. Add a per-token activation variant for int8 weights, so is_dynamic=True has a config that keeps accuracy on transformer graphs.
  2. Meanwhile, say what the current mode does. The linked PR does that part.

Reproduction

from executorch.backends.vulkan.quantizer.vulkan_quantizer import (
    VulkanQuantizer, get_symmetric_quantization_config)
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e

q = VulkanQuantizer().set_global(
    get_symmetric_quantization_config(is_dynamic=True, weight_bits=8))
prepared = prepare_pt2e(torch.export.export(m, inputs, dynamic_shapes=d).module(), q)
prepared(*inputs)
converted = convert_pt2e(prepared)
# compare converted(...) against m(...) with cosine similarity

Needs #22372 first, or _convert_scalars_to_attrs lifts the position-id add as float32 and embedding fails.

Versions

  • executorch==1.4.1, torch==2.14.0.dev20260702, torchao==0.18.0, Python 3.10, macOS arm64
  • Also read on main at 834a4fb020; the comment and the spec are unchanged there.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

Metadata

Metadata

Assignees

No one assigned

    Labels

    module: vulkanIssues related to the Vulkan delegate and code under backends/vulkan/

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions