diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index 9c5fee2e099..2f3990278bb 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -52,6 +52,7 @@ from .fixed_linear_keep_dim import FixedLinearKeepDim from .fold_qdq import FoldQDQ from .fuse_consecutive_cast import FuseConsecutiveCast +from .fuse_consecutive_reshape import FuseConsecutiveReshape from .fuse_consecutive_transpose import FuseConsecutiveTranspose from .i64_to_i32 import I64toI32 from .insert_cast_for_fp_act_quantized_weight import InsertCastForFpActQuantizedWeight @@ -121,6 +122,7 @@ FixedLinearKeepDim, FoldQDQ, FuseConsecutiveCast, + FuseConsecutiveReshape, FuseConsecutiveTranspose, I64toI32, InsertCastForFpActQuantizedWeight, diff --git a/backends/qualcomm/_passes/backends/gpu/qnn_gpu_pass_manager.py b/backends/qualcomm/_passes/backends/gpu/qnn_gpu_pass_manager.py index cf66934da9c..413c9712d69 100644 --- a/backends/qualcomm/_passes/backends/gpu/qnn_gpu_pass_manager.py +++ b/backends/qualcomm/_passes/backends/gpu/qnn_gpu_pass_manager.py @@ -32,7 +32,7 @@ def get_passes_dependency_for_capture_program(cls): return deps @classmethod - def get_annotation_passes(cls): + def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): # The annotation pipeline is skipped for the GPU backend, as it does not # support quantized data types. Return an empty list to indicate a no-op. return [] diff --git a/backends/qualcomm/_passes/backends/htp/qnn_htp_pass_manager.py b/backends/qualcomm/_passes/backends/htp/qnn_htp_pass_manager.py index 2873b413625..0715b746d4b 100644 --- a/backends/qualcomm/_passes/backends/htp/qnn_htp_pass_manager.py +++ b/backends/qualcomm/_passes/backends/htp/qnn_htp_pass_manager.py @@ -36,9 +36,13 @@ def get_passes_dependency_for_capture_program(cls): return deps @classmethod - def get_annotation_passes(cls): + def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): passes = [DecomposeReciprocal, RecomposeHadamard] - passes.extend(super().get_annotation_passes()) + passes.extend( + super().get_annotation_passes( + convert_linear_to_conv2d=convert_linear_to_conv2d, + ) + ) return passes @classmethod diff --git a/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py b/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py index a94766ab335..69001a7483f 100644 --- a/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py +++ b/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py @@ -64,9 +64,13 @@ def _validate_edge_passes(self) -> None: ), "Please ensure LpaiPartitionFallbackSupport is the last edge pass before ResolveDebugHandle." @classmethod - def get_annotation_passes(cls): + def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): passes = [DecomposeHardsigmoid, DecomposeReciprocal] - passes.extend(super().get_annotation_passes()) + passes.extend( + super().get_annotation_passes( + convert_linear_to_conv2d=convert_linear_to_conv2d, + ) + ) return passes @classmethod diff --git a/backends/qualcomm/_passes/convert_linear_to_conv2d.py b/backends/qualcomm/_passes/convert_linear_to_conv2d.py index 03f73736647..c2937a6fa60 100644 --- a/backends/qualcomm/_passes/convert_linear_to_conv2d.py +++ b/backends/qualcomm/_passes/convert_linear_to_conv2d.py @@ -4,11 +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. +from typing import Optional + import torch from executorch.backends.qualcomm._passes.utils import copy_meta from executorch.backends.qualcomm.builders.node_visitor import dq_ops -from executorch.backends.qualcomm.builders.utils import get_parameter +from executorch.backends.qualcomm.builders.utils import ( + get_attr_from_target, + get_parameter, +) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass, PassResult @@ -25,12 +30,47 @@ def _pad_list_to_4(lst): class ConvertLinearToConv2d(ExportPass): """ - Replace aten.linear.default with equivalent 1x1 conv2d using call_function nodes. + Replace linear with an equivalent 1x1 convolution using call_function nodes. """ - def __init__(self, edge_program: torch.export.ExportedProgram): + def __init__( + self, edge_program: Optional[torch.export.ExportedProgram] = None + ) -> None: super().__init__() self.edge_program = edge_program + self.aten_linear_target = torch.ops.aten.linear.default + self.edge_linear_target = exir_ops.edge.aten.linear.default + self._op_table = { + # (is_edge) -> {role: target} + True: { + "view": exir_ops.edge.aten.view_copy.default, + "permute": exir_ops.edge.aten.permute_copy.default, + "conv": exir_ops.edge.aten.convolution.default, + "linear": exir_ops.edge.aten.linear.default, + }, + False: { + "view": torch.ops.aten.reshape.default, + "permute": torch.ops.aten.permute.default, + "conv": torch.ops.aten.conv2d.default, + "linear": torch.ops.aten.linear.default, + }, + } + + def _rewrite_module_stack(self, node, suffix: str): + """ + Append a suffix to nn_module_stack so the quant recipe can differentiate + the converted node. Without this, the conv inherits the linear's + nn_module_stack verbatim, which makes regex targeting in the recipe + ambiguous (e.g. ``model.lm_head`` vs ``model.lm_head.conv``). + """ + stack = node.meta.get("nn_module_stack") + if not stack: + return + stack = dict(stack) + last_key = next(reversed(stack)) + path, qualname = stack[last_key] + stack[last_key] = (path + suffix, qualname) + node.meta["nn_module_stack"] = stack def _register_tensor( self, @@ -79,47 +119,26 @@ def _reshape_weight_for_all_users( self, graph_module: torch.fx.GraphModule, weight_placeholder_node: torch.fx.Node, + ops: dict, ) -> torch.fx.Node: - weight_val = get_parameter(weight_placeholder_node, self.edge_program) + weight_val = ( + get_attr_from_target(graph_module, weight_placeholder_node.target) + if weight_placeholder_node.op == "get_attr" + else get_parameter(weight_placeholder_node, self.edge_program) + ) assert weight_val is not None, "Cannot get the weight in linear node." weight_val = weight_val.reshape(*weight_val.shape, 1, 1).contiguous().detach() get_attr_node = self._register_tensor( graph_module, weight_placeholder_node, weight_val ) - if list(weight_placeholder_node.users)[0].target in dq_ops: - # Scenarios where multiple linear nodes share the same weights, such as the embedding and lm_head in LLM. - for dq_node in list(weight_placeholder_node.users): - if ( - list(dq_node.users)[0].target - is not exir_ops.edge.aten.linear.default - ): - # Add a safety check to prevent replacing weights of non-linear nodes with updated weights. - continue - # For shared weights, a dequantize node is inserted per user after quantization. - # Reuse the dequantize node after weight updates by replacing its input with the corresponding `get_attr` node. - dq_node.replace_input_with(weight_placeholder_node, get_attr_node) - - fake_mode = detect_fake_mode(get_attr_node.meta["val"]) - converter = fake_mode.fake_tensor_converter - dq_node.meta["val"] = converter.from_real_tensor(fake_mode, weight_val) - - # Update block size for per-block quant - if dq_node.target is exir_ops.edge.torchao.dequantize_affine.default: - new_args = list(dq_node.args) - # pad block size - new_args[1] = _pad_list_to_4(list(new_args[1])) - dq_node.args = tuple(new_args) - - return dq_node - else: - for user in list(weight_placeholder_node.users): - if user.target is not exir_ops.edge.aten.linear.default: - # Add a safety check to prevent replacing weights of non-linear nodes with updated weights. - continue - user.replace_input_with(weight_placeholder_node, get_attr_node) + for user in list(weight_placeholder_node.users): + if user.target is not ops["linear"]: + # Add a safety check to prevent replacing weights of non-linear nodes with updated weights. + continue + user.replace_input_with(weight_placeholder_node, get_attr_node) - return get_attr_node + return get_attr_node def call(self, graph_module: GraphModule): graph = graph_module.graph @@ -128,14 +147,16 @@ def call(self, graph_module: GraphModule): preprocessed_linear_weights_set = set() for node in graph.nodes: - if node.target is exir_ops.edge.aten.linear.default: + if node.target in {self.aten_linear_target, self.edge_linear_target}: + is_edge = node.target == self.edge_linear_target + ops = self._op_table[is_edge] input_node = node.args[0] - weight_placeholder_node = ( - # QDQ graph - node.args[1].args[0] - if node.args[1].target in dq_ops - # FP graph - else node.args[1] + weight_placeholder_node = node.args[1] + assert weight_placeholder_node.target not in dq_ops, ( + "ConvertLinearToConv2d does not handle quantized weights. " + "A quantized flow must convert during the annotation " + "pipeline, before q/dq nodes exist. Please set " + "quantizer.set_convert_linear_to_conv2d(True)." ) bias_arg = node.args[2] if len(node.args) > 2 else None @@ -158,18 +179,16 @@ def call(self, graph_module: GraphModule): cur_meta_val = cur_meta_val.reshape(shape) reshape_node = self._create_node( graph_module, - exir_ops.edge.aten.view_copy.default, + ops["view"], (input_node, shape), cur_meta_val, ) - # This pass is scheduled after the `FoldQDQ` pass. After copying the metadata, - # the quantization attributes are also propagated to the target node. order = (0, 3, 1, 2) if rank == 4 else (0, 2, 3, 1) cur_meta_val = cur_meta_val.permute(order) permute_node = self._create_node( graph_module, - exir_ops.edge.aten.permute_copy.default, + ops["permute"], (reshape_node, order) if rank <= 3 else (input_node, order), cur_meta_val, ) @@ -177,14 +196,10 @@ def call(self, graph_module: GraphModule): # Step 2: reshape weight if weight_placeholder_node.name not in preprocessed_linear_weights_set: weight_arg = self._reshape_weight_for_all_users( - graph_module, weight_placeholder_node + graph_module, weight_placeholder_node, ops ) # Add the name of the preprocessed weights to the list. - preprocessed_linear_weights_set.add( - weight_arg.args[0].name - if weight_arg.target in dq_ops - else weight_arg.name - ) + preprocessed_linear_weights_set.add(weight_arg.name) else: # Skip preprocessing as the weights have already been processed due to shared weights. weight_arg = node.args[1] @@ -198,14 +213,47 @@ def call(self, graph_module: GraphModule): output_padding = [0, 0] groups = 1 """ + The two dialects take different argument lists: + Spec for `aten.convolution` (https://docs.pytorch.org/docs/stable/torch.compiler_ir.html) convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups) -> Tensor + + aten::conv2d(Tensor input, Tensor weight, Tensor? bias=None, SymInt[2] stride=[1, 1], + SymInt[2] padding=[0, 0], SymInt[2] dilation=[1, 1], SymInt groups=1) -> Tensor + + conv2d has no transposed / output_padding, so passing the 9-arg + convolution list to it raises "expected at most 7 argument(s)". """ conv_args = ( - permute_node, - weight_arg, - bias_arg, + ( + permute_node, + weight_arg, + bias_arg, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + ) + if is_edge + else ( + permute_node, + weight_arg, + bias_arg, + stride, + padding, + dilation, + groups, + ) + ) + # Shape inference always goes through the 9-arg convolution; + # conv2d is just the narrower spelling of the same op. + cur_meta_val = exir_ops.edge.aten.convolution.default( + cur_meta_val, + weight_meta_val, + bias_meta_val, stride, padding, dilation, @@ -213,15 +261,9 @@ def call(self, graph_module: GraphModule): output_padding, groups, ) - cur_meta_val = exir_ops.edge.aten.convolution.default( - cur_meta_val, - weight_meta_val, - bias_meta_val, - *conv_args[3:], - ) conv_node = self._create_node( graph_module, - exir_ops.edge.aten.convolution.default, + ops["conv"], conv_args, cur_meta_val, meta_source_node=node, @@ -235,7 +277,7 @@ def call(self, graph_module: GraphModule): cur_meta_val = cur_meta_val.permute(order) permute_node = self._create_node( graph_module, - exir_ops.edge.aten.permute_copy.default, + ops["permute"], (conv_node, order), cur_meta_val, ) @@ -244,7 +286,7 @@ def call(self, graph_module: GraphModule): cur_meta_val = cur_meta_val.reshape(target_shape) reshape_node = self._create_node( graph_module, - exir_ops.edge.aten.view_copy.default, + ops["view"], (permute_node, target_shape), cur_meta_val, ) @@ -253,6 +295,7 @@ def call(self, graph_module: GraphModule): node.replace_all_uses_with(permute_node) graph.erase_node(node) + self._rewrite_module_stack(conv_node, ".conv") dead_code_elimination_pass(graph_module) return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/fuse_consecutive_reshape.py b/backends/qualcomm/_passes/fuse_consecutive_reshape.py new file mode 100644 index 00000000000..f474a9191b9 --- /dev/null +++ b/backends/qualcomm/_passes/fuse_consecutive_reshape.py @@ -0,0 +1,53 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# 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. + + +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_REQUANTIZE +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.passes import dead_code_elimination_pass + + +class FuseConsecutiveReshape(ExportPass): + """ + Collapse chains of consecutive view_copy into a single view_copy. + This condition is not added in remove_redundancy.py since + quantized(QDQ) models are hard to handle, and remove_redundancy runs before FoldQDQ. + Some reshape shows up at edge dialect, so annotation phase of remove_redundancy + also won't work. + This pass is designed to be conservative and will not fold if requantize meta is + found during traversal. The only exception is when the requantize meta is on the + last view_copy of the chain, since that node is the one being rewritten rather + than a node being folded away. + """ + + def __init__(self): + super().__init__() + self.view = exir_ops.edge.aten.view_copy.default + + def _fuse(self, graph_module: torch.fx.GraphModule): + for node in graph_module.graph.nodes: + if node.target != self.view: + continue + # Walk back to the first non-view source node; every intermediate view + # is redundant because this node restates the full target shape. + # Need to ensure this optimization doesn't break requantize logic. + src = original_src = node.args[0] + while ( + isinstance(src, torch.fx.Node) + and src.target == self.view + and QCOM_REQUANTIZE not in src.meta + ): + src = src.args[0] + + if src is not original_src and QCOM_REQUANTIZE not in src.meta: + node.args = (src, *node.args[1:]) + + def call(self, graph_module: torch.fx.GraphModule): + self._fuse(graph_module) + dead_code_elimination_pass(graph_module) + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index cd42c024147..e7e98bbcdae 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -55,6 +55,7 @@ FixedLinearKeepDim, FoldQDQ, FuseConsecutiveCast, + FuseConsecutiveReshape, FuseConsecutiveTranspose, I64toI32, InsertCastForFpActQuantizedWeight, @@ -148,6 +149,7 @@ def get_default_pass_activations(cls): (ExpandBroadcastTensorShape, True), (FixedLinearKeepDim, True), (FoldQDQ, True), + (FuseConsecutiveReshape, True), (I64toI32, True), (InsertCastForFpActQuantizedWeight, True), (LayoutTransform, True), @@ -161,9 +163,9 @@ def get_default_pass_activations(cls): ] @classmethod - def get_annotation_passes(cls): + def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): """Return annotation pipeline pass classes. Override in subclasses to add backend-specific passes.""" - return [ + passes = [ RemoveRedundancy, RecomposePixelUnshuffle, RecomposeRmsNorm, @@ -200,6 +202,11 @@ def get_annotation_passes(cls): InsertReshapeForReduceOps, ] + if convert_linear_to_conv2d: + passes.append(ConvertLinearToConv2d) + + return passes + @classmethod def get_export_passes(cls): """Return export pipeline pass classes. Override in subclasses to add backend-specific passes.""" @@ -292,6 +299,7 @@ def get_passes_dependency_for_capture_program(cls): DecomposeAny: [RemoveRedundancy], DecomposeAtan2: [RemoveRedundancy], DecomposeColIm: [FoldQDQ], + FuseConsecutiveReshape: [FoldQDQ], DecomposePDist: [RemoveRedundancy], DecomposeDiagonal: [RemoveRedundancy], DecomposeDivMode: [RemoveRedundancy], @@ -423,9 +431,12 @@ def _instantiate_passes(self, pass_classes, **available_kwargs): def transform_for_annotation_pipeline( self, graph_module: GraphModule, + convert_linear_to_conv2d: bool = False, ): self._instantiate_passes( - self.get_annotation_passes(), + self.get_annotation_passes( + convert_linear_to_conv2d=convert_linear_to_conv2d, + ), quantization_capture=True, ) return self._transform(graph_module) diff --git a/backends/qualcomm/builders/utils.py b/backends/qualcomm/builders/utils.py index 745e3324eb0..936a0be5b69 100755 --- a/backends/qualcomm/builders/utils.py +++ b/backends/qualcomm/builders/utils.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import Dict, Optional +from typing import Any, Dict, Optional import torch from torch._export.utils import ( @@ -17,6 +17,25 @@ ) +def get_attr_from_target(graph_module: torch.fx.GraphModule, target: str) -> Any: + """Resolve a possibly dotted get_attr target (e.g. ``linear.weight``).""" + attr: Any = graph_module + for target_atom in target.split("."): + attr = getattr(attr, target_atom) + return attr + + +def set_attr_from_target( + graph_module: torch.fx.GraphModule, target: str, replacement: Any +) -> None: + """Assign to a possibly dotted get_attr target (e.g. ``linear.weight``).""" + attr: Any = graph_module + target_list = target.split(".") + for target_atom in target_list[:-1]: + attr = getattr(attr, target_atom) + setattr(attr, target_list[-1], replacement) + + def is_parameter( node: torch.fx.Node, edge_program: torch.export.ExportedProgram ) -> bool: diff --git a/backends/qualcomm/quantizer/quantizer.py b/backends/qualcomm/quantizer/quantizer.py index 6eac3fe7e79..0a5a580165b 100644 --- a/backends/qualcomm/quantizer/quantizer.py +++ b/backends/qualcomm/quantizer/quantizer.py @@ -400,6 +400,7 @@ def __init__( self.custom_quant_annotations: Sequence[Callable] = [] self.discard_nodes: Set[str] = set() self._recipe = None + self._convert_linear_to_conv2d = False @property def recipe(self): @@ -534,7 +535,9 @@ def transform_for_annotation(self, model: GraphModule) -> GraphModule: """ return get_qnn_pass_manager_cls( self.backend - )().transform_for_annotation_pipeline(model) + )().transform_for_annotation_pipeline( + model, convert_linear_to_conv2d=self._convert_linear_to_conv2d + ) def validate(self, model: GraphModule) -> None: # Validate: only for mapped nodes (qnn_op present); unmapped → skip validation @@ -689,6 +692,17 @@ def set_block_size_map(self, block_size_map: Dict[str, Tuple]) -> None: """ self.block_size_map = block_size_map + def set_convert_linear_to_conv2d(self, convert_linear_to_conv2d: bool) -> None: + """ + Convert linear to conv2d during the annotation pipeline. + + If this is enabled, quant_recipe will need to target conv node instead of linear node. + + Args: + convert_linear_to_conv2d (bool): True to convert during annotation. + """ + self._convert_linear_to_conv2d = convert_linear_to_conv2d + def set_default_quant_config( self, quant_dtype: QuantDtype, diff --git a/backends/qualcomm/tests/rework/passes/passes_helper.py b/backends/qualcomm/tests/rework/passes/passes_helper.py index d5bf2f7d2ba..5b5a1eec6c6 100644 --- a/backends/qualcomm/tests/rework/passes/passes_helper.py +++ b/backends/qualcomm/tests/rework/passes/passes_helper.py @@ -176,11 +176,16 @@ def lower_annotation_gm( sample_input: tuple[torch.Tensor, ...], target_pass: type[ExportPass], backend_type: QnnExecuTorchBackendType = QnnExecuTorchBackendType.kHtpBackend, + convert_linear_to_conv2d: bool = False, ) -> torch.fx.GraphModule: pm_cls = get_qnn_pass_manager_cls(backend_type) gm = torch.export.export(module, sample_input, strict=True).module() pass_classes = PassPipeline._slice_to_target( - pm_cls.get_annotation_passes(), target_pass, "annotation" + pm_cls.get_annotation_passes( + convert_linear_to_conv2d=convert_linear_to_conv2d, + ), + target_pass, + "annotation", ) instances = PassPipeline._instantiate( pass_classes, diff --git a/backends/qualcomm/tests/rework/passes/test.py b/backends/qualcomm/tests/rework/passes/test.py index 73e02c0a457..d3156597ce2 100644 --- a/backends/qualcomm/tests/rework/passes/test.py +++ b/backends/qualcomm/tests/rework/passes/test.py @@ -326,6 +326,12 @@ def test_fuse_consecutive_cast(request, kwargs): FuseConsecutiveCast.test(request, kwargs) # noqa: F405 +@enumerate_backends() +@repack_pass_fixtures +def test_fuse_consecutive_reshape(request, kwargs): + FuseConsecutiveReshape.test(request, kwargs) # noqa: F405 + + @enumerate_backends() @repack_pass_fixtures def test_fuse_consecutive_transpose(request, kwargs): diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 60b7fe254cb..8b8f7f666bf 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/backends/qualcomm/tests/rework/src/pattern.py @@ -33,6 +33,7 @@ QCOM_AXIS_ORDER, QCOM_PASS_ACTIVATE_KEY, QCOM_QUANT_ATTRS, + QCOM_REQUANTIZE, ) from executorch.exir.delegate import executorch_call_delegate from executorch.exir.dialects._ops import ops as exir_ops @@ -551,47 +552,86 @@ def test( pass_pipeline: PassPipeline, ): target_pass = _passes.ConvertLinearToConv2d - conv = exir_ops.edge.aten.convolution.default + if quantizer is None: + conv = exir_ops.edge.aten.convolution.default + linear = exir_ops.edge.aten.linear.default + view = exir_ops.edge.aten.view_copy.default + permute = exir_ops.edge.aten.permute_copy.default + + def lower(module, sample_input): + return pass_pipeline.lower_edge_ep( + module=module, + sample_input=sample_input, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=target_pass, + quantizer=quantizer, + convert_linear_to_conv2d=True, + ).graph_module + + else: + conv = torch.ops.aten.conv2d.default + linear = torch.ops.aten.linear.default + view = torch.ops.aten.reshape.default + permute = torch.ops.aten.permute.default + + def lower(module, sample_input): + aten_gm = pass_pipeline.lower_annotation_gm( + module=module, + sample_input=sample_input, + target_pass=target_pass, + backend_type=backend_type, + convert_linear_to_conv2d=True, + ) + + # This is an extra test to ensure nn_module_stack got propagated + # to edge graph. + prev_convert = quantizer._convert_linear_to_conv2d + quantizer.set_convert_linear_to_conv2d(True) + try: + # The quantizer fixture is cached and shared across tests, so the + # annotation-time flag is restored before returning. + edge_gm = pass_pipeline.lower_edge_ep( + module=module, + sample_input=sample_input, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=target_pass, + quantizer=quantizer, + convert_linear_to_conv2d=True, + ).graph_module + finally: + quantizer.set_convert_linear_to_conv2d(prev_convert) + + edge_conv = exir_ops.edge.aten.convolution.default + assertions.assert_target_count_at_least(edge_gm, edge_conv, 1) + for node in edge_gm.graph.nodes: + if node.target == edge_conv: + stack = node.meta["nn_module_stack"] + assert any( + path.endswith(".conv") for path, _ in stack.values() + ), f"{node} lost annotation-phase .conv module path: {stack}" + + return aten_gm with subtests.test(msg="basic"): - gm = pass_pipeline.lower_edge_ep( - module=ConvertLinearToConv2d._Basic(), - sample_input=(torch.randn(2, 8),), - backend_type=backend_type, - compile_spec=compile_spec, - target_pass=target_pass, - quantizer=quantizer, - convert_linear_to_conv2d=True, - ).graph_module - assertions.assert_no_target(gm, exir_ops.edge.aten.linear.default) + gm = lower(ConvertLinearToConv2d._Basic(), (torch.randn(2, 8),)) + assertions.assert_no_target(gm, linear) assertions.assert_target_count(gm, conv, 1) # rank-2 input: reshape×2 (input + output restore) + permute×2 (pre/post conv) - assertions.assert_target_count(gm, exir_ops.edge.aten.view_copy.default, 2) - assertions.assert_target_count( - gm, exir_ops.edge.aten.permute_copy.default, 2 - ) + assertions.assert_target_count(gm, view, 2) + assertions.assert_target_count(gm, permute, 2) with subtests.test(msg="shared_weight"): - gm = pass_pipeline.lower_edge_ep( - module=ConvertLinearToConv2d._SharedWeight(), - sample_input=(torch.randn(2, 8), torch.randn(2, 8)), - backend_type=backend_type, - compile_spec=compile_spec, - target_pass=target_pass, - quantizer=quantizer, - convert_linear_to_conv2d=True, - ).graph_module - assertions.assert_no_target(gm, exir_ops.edge.aten.linear.default) + gm = lower( + ConvertLinearToConv2d._SharedWeight(), + (torch.randn(2, 8), torch.randn(2, 8)), + ) + assertions.assert_no_target(gm, linear) assertions.assert_target_count(gm, conv, 2) - conv_weight_sources = set() - for node in gm.graph.nodes: - if node.target is conv: - weight_arg = node.args[1] - conv_weight_sources.add( - weight_arg.args[0] - if weight_arg.target in dq_ops - else weight_arg - ) + conv_weight_sources = { + node.args[1] for node in gm.graph.nodes if node.target is conv + } assert ( len(conv_weight_sources) == 1 ), f"expected both convolution nodes to share one weight source, got {conv_weight_sources}" @@ -3781,6 +3821,128 @@ def test( assertions.assert_target_count_at_most(gm, FuseConsecutiveCast._CAST_OPS, 1) +class FuseConsecutiveReshape: + class _ConsecutiveReshape(torch.nn.Module): + def forward(self, x): + a = torch.relu(x) + b = a.view(2, 3, 4) + c = b.view(6, 4) + d = c.view(24) + e = d.view(12, 2) + return torch.relu(e) + + @staticmethod + def _annotate_16a8w(wide_node_names): + """This is to trigger requantize""" + from executorch.backends.qualcomm.quantizer.qconfig import ( + get_16a8w_qnn_ptq_config, + ) + from executorch.backends.qualcomm.quantizer.rules import Q_ANNOTATION_KEY + from torchao.quantization.pt2e.quantizer import QuantizationAnnotation + + def annotate(gm: torch.fx.GraphModule): + config = get_16a8w_qnn_ptq_config() + for node in gm.graph.nodes: + if node.name not in wide_node_names: + continue + input_qspec_map = { + arg: config.input_activation + for arg in node.args + if isinstance(arg, torch.fx.Node) + } + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=config.output_activation, + _annotated=True, + ) + + return annotate + + @staticmethod + @unpack_pass_fixtures + def test( + quantizer, + compile_spec, + backend_type: QnnExecuTorchBackendType, + assertions: Assertions, + pass_pipeline: PassPipeline, + subtests, + ): + from executorch.backends.qualcomm.export_utils import make_quantizer + + module = FuseConsecutiveReshape._ConsecutiveReshape() + inputs = (torch.randn(4, 6),) + target_pass = _passes.FuseConsecutiveReshape + view = exir_ops.edge.aten.view_copy.default + + def lower(active_quantizer): + return pass_pipeline.lower_edge_ep( + module=module, + sample_input=inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=target_pass, + quantizer=active_quantizer, + ).graph_module + + def assert_requantize_intact(gm): + """Every QCOM_REQUANTIZE entry is keyed by consumer name, so each name + must still be a real user after the fuse rewires args.""" + found = 0 + for node in gm.graph.nodes: + requantize = node.meta.get(QCOM_REQUANTIZE) + if not requantize: + continue + found += 1 + users = {user.name for user in node.users} + missing = set(requantize) - users + assert not missing, ( + f"{node} has QCOM_REQUANTIZE naming non-users {missing}; " + f"actual users {users}" + ) + unnamed = users - set(requantize) + assert not unnamed, ( + f"{node} carries QCOM_REQUANTIZE but gained users {unnamed} " + f"that would read the un-requantized value" + ) + assert found > 0, "expected a QCOM_REQUANTIZE boundary to be annotated" + + with subtests.test(msg="default"): + gm = lower(quantizer) + assertions.assert_no_consecutive(gm, view) + assertions.assert_target_count(gm, view, 1) + + # Test requantize with FuseConsecutiveReshape + if quantizer is None: + return + + requantize_cases = [ + ( + "requantize_after_first_view", + {"view_1", "view_2", "view_3", "relu_1"}, + 4, + ), + ( + "requantize_at_first_relu", + {"view", "view_1", "view_2", "view_3", "relu_1"}, + 4, + ), + ("requantize_at_middle_view", {"view_2", "view_3", "relu_1"}, 3), + ("requantize_at_output", {"relu_1"}, 1), + ] + for name, wide_node_names, expected_views in requantize_cases: + with subtests.test(msg=name): + mixed_quantizer = make_quantizer( + backend=backend_type, + custom_annotations=( + FuseConsecutiveReshape._annotate_16a8w(wide_node_names), + ), + ) + gm = lower(mixed_quantizer) + assertions.assert_target_count(gm, view, expected_views) + assert_requantize_intact(gm) + + class FuseConsecutiveTranspose: class _ConsecutivePermute(torch.nn.Module): def forward(self, x): diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 1ebda343c8c..aba7ef0e123 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -5030,8 +5030,6 @@ def test_qnn_backend_linear(self): self.lower_module_and_test_output(qdq_module, sample_input) def test_qnn_backend_linear_to_conv2d(self): - from executorch.backends.qualcomm._passes import ConvertLinearToConv2d - test_comb = [ { QCOM_MODULE: [ @@ -5046,22 +5044,16 @@ def test_qnn_backend_linear_to_conv2d(self): }, ] - passes_job = get_qnn_pass_manager_cls().get_capture_program_passes() - passes_job[ConvertLinearToConv2d][QCOM_PASS_ACTIVATE_KEY] = True - passes_job[ConvertLinearToConv2d][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ - "edge_program" - ] = None - index = 0 for comb in test_comb: for module in comb[QCOM_MODULE]: for sample_input in comb[QCOM_SAMPLE_INPUTS]: with self.subTest(i=index): index += 1 - qdq_module = self.get_qdq_module(module, sample_input) - self.lower_module_and_test_output( - qdq_module, sample_input, passes_job=passes_job + qdq_module = self.get_qdq_module( + module, sample_input, convert_linear_to_conv2d=True ) + self.lower_module_and_test_output(qdq_module, sample_input) def test_qnn_backend_linear_shared_weights(self): modules = [ @@ -5078,28 +5070,20 @@ def test_qnn_backend_linear_shared_weights(self): self.lower_module_and_test_output(qdq_module, sample_input) def test_qnn_backend_linear_to_conv2d_shared_weights(self): - from executorch.backends.qualcomm._passes import ConvertLinearToConv2d - modules = [ LinearSharedWeight(512, 32), # noqa: F405 ] - passes_job = get_qnn_pass_manager_cls().get_capture_program_passes() - passes_job[ConvertLinearToConv2d][QCOM_PASS_ACTIVATE_KEY] = True - passes_job[ConvertLinearToConv2d][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ - "edge_program" - ] = None - sample_input = ( torch.randn([3, 512]), torch.randn([3, 512]), ) for i, module in enumerate(modules): with self.subTest(i=i): - qdq_module = self.get_qdq_module(module, sample_input) - self.lower_module_and_test_output( - qdq_module, sample_input, passes_job=passes_job + qdq_module = self.get_qdq_module( + module, sample_input, convert_linear_to_conv2d=True ) + self.lower_module_and_test_output(qdq_module, sample_input) @unittest.skipIf(is_qnn_sdk_version_less_than("2.30"), "UT pass after QNN 2.30") def test_qnn_backend_linear_block(self): @@ -5124,19 +5108,12 @@ def test_qnn_backend_linear_block(self): @unittest.skipIf(is_qnn_sdk_version_less_than("2.30"), "UT pass after QNN 2.30") def test_qnn_backend_linear_to_conv2d_block(self): - from executorch.backends.qualcomm._passes import ConvertLinearToConv2d modules = [ Linear(use_bias=False), # noqa: F405 Linear(use_bias=True), # noqa: F405 ] - passes_job = get_qnn_pass_manager_cls().get_capture_program_passes() - passes_job[ConvertLinearToConv2d][QCOM_PASS_ACTIVATE_KEY] = True - passes_job[ConvertLinearToConv2d][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ - "edge_program" - ] = None - sample_input = (torch.randn([3, 512]),) for i, module in enumerate(modules): with self.subTest(i=i): @@ -5148,10 +5125,9 @@ def test_qnn_backend_linear_to_conv2d_block(self): sample_input, quant_dtype=QuantDtype.use_16a4w_block, block_size_map={"linear": (1, 32)}, + convert_linear_to_conv2d=True, ) - self.lower_module_and_test_output( - module, sample_input, passes_job=passes_job - ) + self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_linear_qat(self): """ @@ -9337,17 +9313,26 @@ def test_hf_causal_lm(self): # This is the Hugging Face transformers flow, not the static llm flow. if not self.required_envs([]): self.skipTest("missing required envs") - prompt = "My favourite condiment is " + + # TODO: Robust testing framework to check accuracy and performance metrics. + golden_start_with = { + "llama3_2-1b": "Simply put, the theory of relativity states that the speed of light", + "qwen2_5-0_5b": "Simply put, the theory of relativity states that the laws of physics", + "qwen3-0_6b": "Simply put, the theory of relativity states that the laws of physics", + "smollm2_135m": "Simply put, the theory of relativity states that the speed of light", + "granite-3_3-2b": "Simply put, the theory of relativity states that the laws of physics", + } + assert ( + self.model_name in golden_start_with + ), f"{self.model_name} is not supported in test_hf_causal_lm. Currently support: {golden_start_with.keys()}" + prompt = "Simply put, the theory of relativity states that" cmds = [ "python", f"{self.executorch_root}/examples/qualcomm/oss_scripts/hf_causal_lm.py", "--prompt", prompt, "--decoder_model", - "qwen2_5-0_5b", - "--ptq", - "16a8w", - "--enable_spinquant_r3", + self.model_name, "--max_seq_len", "128", "--artifact", @@ -9357,7 +9342,6 @@ def test_hf_causal_lm(self): ] self.add_default_cmds(cmds) - golden_start_with = "My favourite condiment is iced tea." p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) with Listener((self.ip, self.port)) as listener: conn = listener.accept() @@ -9369,8 +9353,8 @@ def test_hf_causal_lm(self): if not self.compile_only: model_out = msg["result"][0] self.assertTrue( - model_out.startswith(golden_start_with), - f"Expected Output: '{golden_start_with}' Actual Output: '{model_out}'", + model_out.startswith(golden_start_with[self.model_name]), + f"Expected Output: '{golden_start_with[self.model_name]}' Actual Output: '{model_out}'", ) def test_static_llm_qat(self): diff --git a/backends/qualcomm/tests/utils.py b/backends/qualcomm/tests/utils.py index 40cad560366..8bcbcfa7b2e 100644 --- a/backends/qualcomm/tests/utils.py +++ b/backends/qualcomm/tests/utils.py @@ -765,6 +765,7 @@ def get_qdq_module( bypass_check: bool = False, block_size_map: Dict[str, Tuple] = None, submodule_qconfig_list: Optional[List[Tuple[Callable, ModuleQConfig]]] = None, + convert_linear_to_conv2d: bool = False, ) -> torch.fx.GraphModule: m = torch.export.export( module, inputs, dynamic_shapes=dynamic_shapes, strict=True @@ -780,6 +781,7 @@ def get_qdq_module( backend=get_backend_type(self.backend), soc_model=self.soc_model, ) + quantizer.set_convert_linear_to_conv2d(convert_linear_to_conv2d) if block_size_map is not None: quantizer.set_block_size_map(block_size_map) prepared = prepare_pt2e(m, quantizer) diff --git a/examples/qualcomm/oss_scripts/hf_causal_lm.py b/examples/qualcomm/oss_scripts/hf_causal_lm.py index 2ed7e49d2a8..5e6c95e4154 100644 --- a/examples/qualcomm/oss_scripts/hf_causal_lm.py +++ b/examples/qualcomm/oss_scripts/hf_causal_lm.py @@ -19,8 +19,6 @@ SimpleADB, ) -from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype - from executorch.examples.qualcomm.oss_scripts.llm_utils.qnn_decoder_model_manager import ( get_qnn_llm_edge_manager, HUGGING_FACE_REPO_IDS, @@ -35,40 +33,47 @@ PTE_FILENAME = "hf_causal_lm_qnn" +# Map the HF decoder_model keys (HUGGING_FACE_REPO_IDS) to the version strings +# the shared qnn_llama_runner understands (see runner.cpp Runner()). +DECODER_MODEL_VERSION = { + "llama3_2-1b": "llama3", + "qwen2_5-0_5b": "qwen2_5", + "qwen2_5-1_5b_instruct": "qwen2_5", + "qwen2_5-0_5b_instruct": "qwen2_5", + "qwen3-0_6b": "qwen3", + "smollm2_135m": "smollm2_135m", + "granite-3_3-2b": "granite", +} + def compile(args: argparse.Namespace, qnn_config: QnnConfig): # noqa: C901 # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - manager = get_qnn_llm_edge_manager( - args.decoder_model, args.max_seq_len, args.enable_spinquant_r3 - ) + manager = get_qnn_llm_edge_manager(args.decoder_model, args.max_seq_len) fixed_point_type = {} - if args.ptq: - if args.ptq == "8a8w": - fixed_point_type["io_type"] = torch.uint8 + if not args.use_fp16: + kv_bits = manager.quant_recipe.get_kv_io_bit_width() + if kv_bits == 8: fixed_point_type["kv_type"] = torch.uint8 - elif args.ptq in ( - "16a8w", - "16a4w", - "16a4w_block", - "16a16w", - ): - fixed_point_type["io_type"] = torch.uint16 + elif kv_bits == 16: fixed_point_type["kv_type"] = torch.uint16 else: - raise ValueError( - f"No support for quant type {args.ptq}. Support 8a8w, 16a8w, 16a4w and 16a4w_block." - ) - quant_dtype = getattr(QuantDtype, f"use_{args.ptq}") + raise RuntimeError(f"unknown kv io bit width {kv_bits}") + + logits_bits = manager.quant_recipe.get_logits_output_bit_width() + if logits_bits == 16: + fixed_point_type["io_type"] = torch.uint16 + else: + raise ValueError("Only support uint16 logits output for quantized hf llm.") + model_id = HUGGING_FACE_REPO_IDS[args.decoder_model] tokenizer = AutoTokenizer.from_pretrained(model_id) tokenizer_json_path = tokenizer.save_pretrained(args.artifact)[-1] manager.pt2e_quantize( - quant_dtype, fixed_point_type, args.calibration_tasks, args.calibration_limit, @@ -83,7 +88,7 @@ def compile(args: argparse.Namespace, qnn_config: QnnConfig): # noqa: C901 qnn_config.skip_delegate_node_ids, qnn_config.skip_delegate_node_ops, ) - if args.ptq: + if not args.use_fp16: logits_quant_attrs = manager.get_logits_quant_attrs() json.dump( { @@ -106,12 +111,30 @@ def inference(args: argparse, qnn_config: QnnConfig): def post_process(): with open(f"{args.artifact}/outputs/result.txt", "r") as f: - outputs.append(f.read()) + text = f.read() + # In tokenized-prompt mode the runner echoes the prompt-file path instead + # of the prompt text; drop it and prepend the real prompt for readability. + prefix = os.path.basename(tokenized_prompt_path) + if text.startswith(prefix): + text = text[len(prefix) :] + outputs.append(args.prompt + text) model_id = HUGGING_FACE_REPO_IDS[args.decoder_model] tokenizer = AutoTokenizer.from_pretrained(model_id) tokenizer_json_path = tokenizer.save_pretrained(args.artifact)[-1] seq_len = args.max_seq_len + runner_bin = "examples/qualcomm/oss_scripts/llama/qnn_llama_runner" + decoder_model_version = DECODER_MODEL_VERSION[args.decoder_model] + + # The base (non-instruct) HF models were not trained on the runner's chat + # template. Tokenize the raw prompt here (matching the Python calibration + # path) and feed it via --tokenized_prompt so the runner skips + # get_formatted_prompt. File format: raw little-endian uint64 tokens. + import numpy as np + + prompt_token_ids = tokenizer(args.prompt)["input_ids"] + tokenized_prompt_path = f"{args.artifact}/tokenized_prompt.raw" + np.asarray(prompt_token_ids, dtype=np.uint64).tofile(tokenized_prompt_path) if args.enable_x86_64: # x86 emulator is intended for CI and not performance. Check only the first few tokens. seq_len = min(seq_len, 16) @@ -121,13 +144,15 @@ def post_process(): runner_cmd = " ".join( [ f"export LD_LIBRARY_PATH={qnn_sdk}/lib/{target}/:{args.build_folder}/lib &&", - f"{args.build_folder}/examples/models/llama/llama_main", - f'--prompt "{args.prompt}"', + f"{args.build_folder}/{runner_bin}", + f"--tokenized_prompt {tokenized_prompt_path}", + f"--decoder_model_version {decoder_model_version}", + "--eval_mode 0", f"--tokenizer_path {tokenizer_json_path}", f"--model_path {pte_path}", f"--seq_len {seq_len}", "--temperature 0", - f" > {output_data_folder}/result.txt", + f"--output_path {output_data_folder}/result.txt", ] ) subprocess.run( @@ -141,23 +166,25 @@ def post_process(): runner_cmd = " ".join( [ f"cd {workspace} &&", - "./llama_main", - f'--prompt "{args.prompt}"', + "./qnn_llama_runner", + "--tokenized_prompt tokenized_prompt.raw", + f"--decoder_model_version {decoder_model_version}", + "--eval_mode 0", "--tokenizer_path tokenizer.json", f"--model_path {PTE_FILENAME}.pte", f"--seq_len {seq_len}", "--temperature 0", - " > outputs/result.txt", + "--output_path outputs/result.txt", ] ) adb = SimpleADB( qnn_config=qnn_config, pte_path=pte_path, workspace=workspace, - runner="examples/models/llama/llama_main", + runner=runner_bin, ) # No pregen inputs, input_list is not required - adb.push(inputs=[], files=[tokenizer_json_path]) + adb.push(inputs=[], files=[tokenizer_json_path, tokenized_prompt_path]) adb.execute(custom_runner_cmd=runner_cmd) adb.pull(host_output_path=args.artifact, callback=post_process) @@ -200,18 +227,18 @@ def main(args): ) parser.add_argument( - "-P", - "--ptq", - choices=["8a8w", "16a8w", "16a4w", "16a4w_block"], - help="If specified, will do PTQ quantization.", + "--prompt", + help="User prompts for LLM.", + required=True, type=str, ) parser.add_argument( - "--prompt", - help="User prompts for Qwen.", - required=True, - type=str, + "-F", + "--use_fp16", + help="If specified, will run in fp16 precision and discard ptq setting", + action="store_true", + default=False, ) parser.add_argument( @@ -240,11 +267,6 @@ def main(args): default=None, help="number of samples used for calibration from lm_eval", ) - parser.add_argument( - "--enable_spinquant_r3", - action="store_true", - help="Specify to enable spin quant R3", - ) try: args = parser.parse_args() diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py index 264f8b74480..1a0d45dd25e 100644 --- a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py +++ b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py @@ -12,7 +12,7 @@ import re from functools import partial -from typing import Any, Dict, List +from typing import Dict, List import torch @@ -21,7 +21,11 @@ from executorch.backends.qualcomm._passes.qnn_pass_manager import ( get_qnn_pass_manager_cls, ) -from executorch.backends.qualcomm.builders.utils import is_graph_output +from executorch.backends.qualcomm.builders.utils import ( + get_attr_from_target, + is_graph_output, + set_attr_from_target, +) from executorch.backends.qualcomm.export_utils import make_quantizer from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype @@ -951,25 +955,10 @@ def activation_override(quantized_node, unquantized_node): def parameter_override(quantized_node, unquantized_node): # Some parameters need to be iterated over to retrieve attributes such as static_llama.tok_embedding.weight - def _get_attr(graph_module: torch.fx.GraphModule, target: str) -> Any: - attr: Any = graph_module - for target_atom in target.split("."): - attr = getattr(attr, target_atom) - return attr - - def _set_attr( - graph_module: torch.fx.GraphModule, target: str, replacement: Any - ) -> Any: - attr: Any = graph_module - target_list = target.split(".") - for target_atom in target_list[:-1]: - attr = getattr(attr, target_atom) - setattr(attr, target_list[-1], replacement) - - _set_attr( + set_attr_from_target( unquantized_model, unquantized_node.target, - _get_attr(quantized_model, quantized_node.target), + get_attr_from_target(quantized_model, quantized_node.target), ) # scale / zero point are part of op's attributes if list(quantized_node.users)[0].target in ptq_target: diff --git a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py index 8dc334baf28..12406f9444a 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py +++ b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py @@ -4,20 +4,60 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import importlib import logging -import math from typing import Optional -import scipy import torch import transformers +from executorch.examples.qualcomm.oss_scripts.llama.model.apply_rope import ( + apply_rotary_emb_single, +) from transformers import GenerationConfig, PretrainedConfig - -from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM +from transformers.cache_utils import Cache, StaticLayer TRANSFORMERS_VERSION = "4.53.1" +class QnnCustomStaticLayer(StaticLayer): + """StaticLayer that returns cat(past, new) to attention (so the current layer + sees the full context, matching static_llama.py:494), and stashes the pre-cat + new K/V slice so the wrapper can return only the new slot as graph output.""" + + def __init__(self, max_cache_len): + super().__init__(max_cache_len=max_cache_len) + self.new_keys = None + self.new_values = None + + def update(self, key_states, value_states, cache_kwargs=None): + self.new_keys = key_states + self.new_values = value_states + keys = torch.cat([self.keys, key_states.transpose(2, 3)], dim=-1) + values = torch.cat([self.values, value_states], dim=-2) + return keys, values + + def get_mask_sizes(self, cache_position): + return self.max_cache_len, 0 + + +class QnnCustomStaticCache(Cache): + """StaticCache-shaped cache seeded from external past K/V tensors, one pair + per layer. `max_cache_len` is the full context length (past_len + ar_len).""" + + def __init__(self, past_k_list, past_v_list, max_cache_len): + layers = [] + for pk, pv in zip(past_k_list, past_v_list): + layer = QnnCustomStaticLayer(max_cache_len=max_cache_len) + layer.max_batch_size, layer.num_heads, _, layer.head_dim = pk.shape + layer.dtype = pk.dtype + layer.device = pk.device + layer.keys = pk + layer.values = pv + layer.is_initialized = True + layers.append(layer) + super().__init__(layers=layers) + + def save_config_to_constant_methods( config: PretrainedConfig, generation_config: Optional[GenerationConfig] = None, @@ -26,9 +66,10 @@ def save_config_to_constant_methods( # Initialize metadata with values from model config metadata = { "get_bos_id": getattr(config, "bos_token_id", None), - "get_eos_id": getattr(config, "eos_token_id", None), + "get_eos_ids": getattr(config, "eos_token_id", None), "get_vocab_size": getattr(config, "vocab_size", None), "get_max_seq_len": getattr(config, "max_position_embeddings", None), + "get_n_layers": getattr(config, "num_hidden_layers", None), "use_kv_cache": getattr(generation_config, "use_cache", None), "use_sdpa_with_kv_cache": False, } @@ -63,15 +104,6 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) -@torch._dynamo.assume_constant_result -def get_transposed_hadamard_matrix(head_dim): - r3_weight = torch.tensor( - scipy.linalg.hadamard(head_dim, dtype=float) / math.sqrt(head_dim), - dtype=torch.float32, - ) - return r3_weight.transpose(0, 1) - - def _qnn_attention( module: torch.nn.Module, query: torch.Tensor, @@ -82,17 +114,12 @@ def _qnn_attention( dropout: float = 0.0, **kwargs, ): - if getattr(module.config, "enable_spinquant_r3", False): - r3_weight = get_transposed_hadamard_matrix(module.head_dim) - query = torch.matmul(query, r3_weight) - key = torch.matmul(key, r3_weight) - key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + attn_weights = torch.matmul(query, key_states) * scaling if attention_mask is not None: - causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + causal_mask = attention_mask[:, :, :, : key_states.shape[-1]] attn_weights = attn_weights + causal_mask attn_weights = torch.nn.functional.softmax( @@ -115,16 +142,50 @@ def _qnn_attention_mask( ): kv_arange = torch.arange(kv_length, device=cache_position.device) reshaped_cache_position = cache_position.view(-1, 1) - - # Simplest and most efficient way to obtain a causal mask causal_mask = kv_arange <= reshaped_cache_position - atten_mask = torch.full((causal_mask.shape[0], kv_length), -65504.0) + atten_mask = torch.full((causal_mask.shape[0], kv_length), -65535.0) atten_mask = atten_mask.masked_fill(causal_mask, 0) atten_mask = atten_mask[None, None, :, :].expand(batch_size, -1, -1, -1) return atten_mask +def _qnn_apply_rotary_pos_emb( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, + unsqueeze_dim: int = 1, +): + return ( + apply_rotary_emb_single(q, cos, sin), + apply_rotary_emb_single(k, cos, sin), + ) + + +class QnnPrecomputedRotaryEmbedding(torch.nn.Module): + def __init__(self, rotary_emb: torch.nn.Module, max_seq_len: int, dtype): + super().__init__() + positions = torch.arange(max_seq_len, dtype=torch.long).unsqueeze(0) + dummy = torch.zeros(1, max_seq_len, 1, dtype=dtype) + with torch.no_grad(): + cos, sin = rotary_emb(dummy, positions) + half = cos.shape[-1] // 2 + self.register_buffer( + "cos_table", cos[0][:, :half].to(dtype).contiguous(), persistent=False + ) + self.register_buffer( + "sin_table", sin[0][:, :half].to(dtype).contiguous(), persistent=False + ) + + def forward(self, x, position_ids): + flat = position_ids.reshape(-1) + cos = self.cos_table.index_select(0, flat) + sin = self.sin_table.index_select(0, flat) + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + class QnnCausalLMExportableModule(torch.nn.Module): def __init__(self, model): super().__init__() @@ -134,32 +195,89 @@ def __init__(self, model): model.config, model.generation_config ) logging.info(f"Metadata to be recorded in PTE: {self._metadata}") - self.exportable_module = TorchExportableModuleForDecoderOnlyLM( - self.model, - batch_size=1, - max_cache_len=self._metadata.get("get_max_seq_len"), + + self.num_layers = self.config.num_hidden_layers + self.num_kv_heads = getattr( + self.config, "num_key_value_heads", self.config.num_attention_heads + ) + self.head_dim = self.config.head_dim + self.max_seq_len = self.config.max_seq_len + self.ar_len = self.config.ar_len + self.past_len = self.max_seq_len - self.ar_len + + self._register_attention_mask_for_4_53() + self._use_precomputed_rope() + + def _use_precomputed_rope(self): + decoder = self.model.model + rotary_emb = getattr(decoder, "rotary_emb", None) + if rotary_emb is None: + logging.warning("No rotary_emb found; skipping RoPE precompute.") + return + decoder.rotary_emb = QnnPrecomputedRotaryEmbedding( + rotary_emb, self.max_seq_len, self.model.dtype ) - self._register_attention_mask_for_4_53(self.exportable_module) - def _register_attention_mask_for_4_53(self, exportable_module: torch.nn.Module): + modeling = importlib.import_module(type(decoder).__module__) + if not hasattr(modeling, "apply_rotary_pos_emb"): + logging.warning( + f"{modeling.__name__} has no apply_rotary_pos_emb; " + "keeping HF's rotation." + ) + else: + modeling.apply_rotary_pos_emb = _qnn_apply_rotary_pos_emb + logging.info( + f"Patched {modeling.__name__}.apply_rotary_pos_emb with " + "static_llama's apply_rotary_emb_single." + ) + + logging.info( + f"Replaced in-graph RoPE with precomputed tables (max_seq_len={self.max_seq_len})." + ) + + def _register_attention_mask_for_4_53(self): if transformers.__version__ >= TRANSFORMERS_VERSION: from transformers.masking_utils import AttentionMaskInterface from transformers.modeling_utils import AttentionInterface AttentionInterface.register("qnn_attention", _qnn_attention) AttentionMaskInterface.register("qnn_attention", _qnn_attention_mask) - exportable_module.model.model.config._attn_implementation = "qnn_attention" + self.model.config._attn_implementation = "qnn_attention" self._metadata.update({"use_sdpa_with_kv_cache": False}) def get_example_inputs(self): - example_input_ids = torch.tensor([[1]], dtype=torch.long) - example_cache_position = torch.tensor([0], dtype=torch.long) - return (example_input_ids, example_cache_position) + input_tokens = torch.ones((1, self.ar_len), dtype=torch.int32) + # Explicit additive causal mask, matching static_llama / KVManager: + # 0.0 == attend, large-negative == masked. Shape [B, 1, ar_len, context_len]. + atten_mask = torch.zeros(1, 1, self.ar_len, self.max_seq_len) + pos_ids = torch.zeros((1, self.ar_len), dtype=torch.int32) + # K cache is transposed (seq last) to match static_llama: + # K: [B, H, head_dim, past_len] V: [B, H, past_len, head_dim] + past_k = [ + torch.zeros(1, self.num_kv_heads, self.head_dim, self.past_len) + for _ in range(self.num_layers) + ] + past_v = [ + torch.zeros(1, self.num_kv_heads, self.past_len, self.head_dim) + for _ in range(self.num_layers) + ] + return (input_tokens, atten_mask, pos_ids, past_k, past_v) + + def forward(self, input_tokens, atten_mask, pos_ids, past_k, past_v): + cache = QnnCustomStaticCache(past_k, past_v, max_cache_len=self.max_seq_len) - def forward(self, input_ids: torch.Tensor, cache_position: torch.Tensor): - return self.exportable_module( - input_ids=input_ids, cache_position=cache_position + outs = self.model( + input_ids=input_tokens, + attention_mask=atten_mask, + position_ids=pos_ids, + past_key_values=cache, + cache_position=pos_ids, + use_cache=True, ) + # Return only the new slice, transposing K back to static_llama layout. + new_k = [layer.new_keys.transpose(-1, -2) for layer in cache.layers] + new_v = [layer.new_values for layer in cache.layers] + return outs.logits, new_k, new_v def get_metadata(self): return self._metadata diff --git a/examples/qualcomm/oss_scripts/llm_utils/hf_llm_quant_recipe.py b/examples/qualcomm/oss_scripts/llm_utils/hf_llm_quant_recipe.py new file mode 100644 index 00000000000..78c6835e271 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llm_utils/hf_llm_quant_recipe.py @@ -0,0 +1,299 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# 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 typing import Optional + +import torch +from executorch.backends.qualcomm.quantizer.custom_annotation import annotate_kv_8bit +from executorch.backends.qualcomm.quantizer.quant_recipe import ( + QuantGranularity, + QuantRecipe, +) +from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype +from torchao.quantization.pt2e import MinMaxObserver + + +class HFLLMQuantRecipe: + """ + Qualcomm's HuggingFace-style LLM quantization recipe. + Mirrors ``StaticLLMQuantRecipe`` but the recipe is applied at *annotation* + time before ``convert_linear_to_conv2d`` runs, so strategies target + ``aten.linear.default`` with 2D block sizes, and regex patterns match the + HuggingFace module tree (``model.lm_head``, ``model.layers.N.mlp.*``). + """ + + def __init__(self): + self.recipe: Optional[QuantRecipe] = None + + # For IO bitwidth + self.default_quant_dtype = getattr(self, "default_quant_dtype", None) + if self.default_quant_dtype is None: + raise ValueError("default_quant_dtype must be defined in the recipe.") + + def get_kv_io_bit_width(self) -> int: + if self.default_quant_dtype is None: + return 32 + elif self.default_quant_dtype == QuantDtype.use_8a8w or annotate_kv_8bit in ( + getattr(c, "func", c) for c in self.recipe.custom_quant_annotations + ): + return 8 + else: + # If quantized but not 8a8w or mix_quantization, it has to be 16bit kv io. + return 16 + + def get_logits_output_bit_width(self) -> int: + # We use 16bit logits for all quant config + return 32 if self.default_quant_dtype is None else 16 + + +class Llama3_2_1B_HFQuantRecipe(HFLLMQuantRecipe): + default_quant_dtype = QuantDtype.use_16a4w + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation", + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 32, 1, 1)}, + note="Annotate with 16a4w block quantization since these layers are not sensitive.", + ) + .add_regex( + { + r"model\.lm_head\.conv", + r"model\.layers\.[0-3]\.mlp\.down_proj\.conv", + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="Head and early down_proj are sensitive and should be annotated with 16a8w.", + ) + ) + self.recipe.custom_quant_annotations.append(annotate_kv_8bit) + + +class Qwen2_5_0_5B_HFQuantRecipe(HFLLMQuantRecipe): + """Ported from Qwen2_5_0_5BQuantRecipe (static_llm_quant_recipe.py:545).""" + + default_quant_dtype = QuantDtype.use_16a4w + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 16, 1, 1)}, + ) + .add_regex( + {r"model\.lm_head\.conv"}, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) + ) + + +class Qwen2_5_1_5B_HFQuantRecipe(HFLLMQuantRecipe): + """Ported from Qwen2_5_1_5BQuantRecipe (static_llm_quant_recipe.py:569). + + static_llama's ``output\\.conv`` head is ``model.lm_head.conv`` in HF. + """ + + default_quant_dtype = QuantDtype.use_16a4w + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 16, 1, 1)}, + ) + .add_regex( + {r"model\.lm_head\.conv"}, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) + ) + + +class Qwen3_0_6B_HFQuantRecipe(HFLLMQuantRecipe): + """Ported from Qwen3_0_6BQuantRecipe (static_llm_quant_recipe.py:603). + + static_llama's ``layers\\..*\\.feed_forward\\..*w2_conv`` (the down + projection) is ``model.layers.N.mlp.down_proj.conv`` in HF. + """ + + default_quant_dtype = QuantDtype.use_16a4w + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 32, 1, 1)}, + ) + .add_regex( + { + r"model\.layers\..*\.mlp\.down_proj\.conv", + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) + ) + + +class Smollm2_HFQuantRecipe(HFLLMQuantRecipe): + """Ported from Smollm2QuantRecipe (static_llm_quant_recipe.py:676).""" + + default_quant_dtype = QuantDtype.use_16a8w + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ).add_node_target( + { + torch.ops.aten.conv2d.default, + }, + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) + + +class Granite_3_3_2B_Instruct_HFQuantRecipe(HFLLMQuantRecipe): + """Ported from Granite_3_3_2B_InstructQuantRecipe (static_llm_quant_recipe.py:403). + + static_llama's ``layers\\..*\\.attention\\..*wv.*`` (the value projection) is + ``model.layers.N.self_attn.v_proj.conv`` in HF. + """ + + default_quant_dtype = QuantDtype.use_16a4w + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 64, 1, 1)}, + ) + .add_regex( + { + r"model\.layers\..*\.self_attn\.v_proj\.conv", + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) + ) + self.recipe.custom_quant_annotations.append(annotate_kv_8bit) + + +class DefaultQuantRecipe(HFLLMQuantRecipe): + """When quant recipe is not provided, this will be used""" + + default_quant_dtype = QuantDtype.use_16a8w + + def __init__(self, verbose: bool = False): + super().__init__() + self.recipe = QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + ).add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + ) diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py index da1ff8e0aba..40dac6008e5 100644 --- a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -6,7 +6,7 @@ import logging from functools import partial -from typing import Callable, List +from typing import Callable, List, Optional import torch from executorch.backends.qualcomm._passes import TagQuantIO @@ -16,7 +16,6 @@ ) from executorch.backends.qualcomm.builders.utils import is_graph_output from executorch.backends.qualcomm.export_utils import make_quantizer -from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype from executorch.backends.qualcomm.utils.constants import ( QCOM_PASS_ACTIVATE_KEY, QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY, @@ -32,10 +31,19 @@ from executorch.examples.qualcomm.oss_scripts.llm_utils.decoder_model_wrapper import ( QnnCausalLMExportableModule, ) +from executorch.examples.qualcomm.oss_scripts.llm_utils.hf_llm_quant_recipe import ( + DefaultQuantRecipe, + Granite_3_3_2B_Instruct_HFQuantRecipe, + HFLLMQuantRecipe, + Llama3_2_1B_HFQuantRecipe, + Qwen2_5_0_5B_HFQuantRecipe, + Qwen2_5_1_5B_HFQuantRecipe, + Qwen3_0_6B_HFQuantRecipe, + Smollm2_HFQuantRecipe, +) from executorch.exir.capture._config import ExecutorchBackendConfig from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass from pytorch_tokenizers import get_tokenizer -from torchao.quantization.pt2e import MinMaxObserver from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from transformers import AutoConfig, AutoModelForCausalLM, GenerationConfig @@ -43,6 +51,10 @@ FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" logging.basicConfig(level=logging.INFO, format=FORMAT) +# Method name the shared qnn_llama_runner expects for the KV path +# (see examples/qualcomm/oss_scripts/llama/runner/runner.cpp). +KV_FORWARD = "kv_forward" + HUGGING_FACE_REPO_IDS = { "llama3_2-1b": "NousResearch/Llama-3.2-1B", "qwen2_5-0_5b": "Qwen/Qwen2.5-0.5B", @@ -53,8 +65,21 @@ "granite-3_3-2b": "ibm-granite/granite-3.3-2b-instruct", } +# TODO: This dict is temporary and will require a refactor later. +# Will create a file similar to executorch/examples/qualcomm/oss_scripts/llama/__init__.py +# and migrate model specific configs there. +HUGGING_FACE_QUANT_RECIPES = { + "llama3_2-1b": Llama3_2_1B_HFQuantRecipe, + "qwen2_5-0_5b": Qwen2_5_0_5B_HFQuantRecipe, + "qwen2_5-0_5b_instruct": Qwen2_5_0_5B_HFQuantRecipe, + "qwen2_5-1_5b_instruct": Qwen2_5_1_5B_HFQuantRecipe, + "qwen3-0_6b": Qwen3_0_6B_HFQuantRecipe, + "smollm2_135m": Smollm2_HFQuantRecipe, + "granite-3_3-2b": Granite_3_3_2B_Instruct_HFQuantRecipe, +} + -def get_qnn_llm_edge_manager(model_name, max_seq_len=128, enable_spinquant_r3=True): +def get_qnn_llm_edge_manager(model_name, max_seq_len=128): model_id = HUGGING_FACE_REPO_IDS[model_name] config = AutoConfig.from_pretrained(model_id) device = "cpu" @@ -67,7 +92,6 @@ def get_qnn_llm_edge_manager(model_name, max_seq_len=128, enable_spinquant_r3=Tr config.max_seq_len = max_seq_len config.ar_len = 1 # kv mode config.max_batch_size = batch_size - config.enable_spinquant_r3 = enable_spinquant_r3 config.use_cache = True # Some config has head_dim provided that is different from equation below(e.g., qwen3) @@ -106,6 +130,14 @@ def __init__(self, model_name, model_wrapper, config, verbose=True) -> None: self.passes_job = get_qnn_pass_manager_cls().get_capture_program_passes() self.edge_prog_mgr = None self.logits_quant_attrs = None + recipe_cls = HUGGING_FACE_QUANT_RECIPES.get(model_name, DefaultQuantRecipe) + if recipe_cls == DefaultQuantRecipe: + logging.warning( + f"{model_name} does not have customized quant recipe using default quant recipe." + ) + self.quant_recipe: Optional[HFLLMQuantRecipe] = ( + recipe_cls(verbose) if recipe_cls else None + ) def source_transform( self, transforms: List[Callable[[torch.nn.Module], torch.nn.Module]] @@ -126,13 +158,16 @@ def source_transform( return self def _tag_ios(self, node, fixed_point_type, config): - # shape of k caches and v caches + # static_llama layout: K is transposed (seq last), V is seq-major. + # K in: [B, H, head_dim, past_len] K out: [B, H, head_dim, ar_len] + # V in: [B, H, past_len, head_dim] V out: [B, H, ar_len, head_dim] + past_len = config.max_seq_len - config.ar_len kv_cache_shape = { - # single head, kv input - (config.head_dim, config.max_seq_len), - (config.max_seq_len, config.head_dim), - # single head, kv output + # K (head_dim, seq) + (config.head_dim, past_len), (config.head_dim, config.ar_len), + # V (seq, head_dim) + (past_len, config.head_dim), (config.ar_len, config.head_dim), } @@ -144,16 +179,30 @@ def _tag_ios(self, node, fixed_point_type, config): ) } + atten_mask_shape = { + ( + config.max_batch_size, + 1, + config.ar_len, + config.max_seq_len, + ) + } + quant_io_type = None if node.op == "placeholder": if ( - len(users := list(node.users)) == 1 - and users[0].meta["val"].size()[-2:] in kv_cache_shape + node.meta["val"].dim() == 4 + and node.meta["val"].size()[-2:] in kv_cache_shape ): quant_io_type = fixed_point_type["kv_type"] + elif node.meta["val"].size() in atten_mask_shape: + quant_io_type = fixed_point_type["io_type"] if is_graph_output(node): - if node.meta["val"].size()[-2:] in kv_cache_shape: + if ( + node.meta["val"].dim() == 4 + and node.meta["val"].size()[-2:] in kv_cache_shape + ): quant_io_type = fixed_point_type["kv_type"] elif node.meta["val"].size() in logit_out_shape: quant_io_type = fixed_point_type["io_type"] @@ -181,17 +230,62 @@ def pt2e_calibrate( f"Calibrating with tasks: {calibration_tasks}, limit: {calibration_limit}, calibration_data: {calibration_data}, tokenizer_path: {tokenizer_path}, seq_length: {self.config.max_seq_len}" ) + def _empty_past(): + past_k = [ + torch.zeros( + 1, + self.model_wrapper.num_kv_heads, + self.model_wrapper.head_dim, + self.model_wrapper.past_len, + ) + for _ in range(self.model_wrapper.num_layers) + ] + past_v = [ + torch.zeros( + 1, + self.model_wrapper.num_kv_heads, + self.model_wrapper.past_len, + self.model_wrapper.head_dim, + ) + for _ in range(self.model_wrapper.num_layers) + ] + return past_k, past_v + + def _build_mask(n_past, past_len, context_len): + mask = torch.full((1, 1, 1, context_len), -65535.0) + mask[..., :n_past] = 0.0 + mask[..., past_len:] = 0.0 + return mask + def calibrate_template( module: torch.fx.GraphModule, tokenizer, prompts: str, max_len: int ): - # TODO: change criteria & support batch inputs if necessary pos = 0 token_list = tokenizer.encode(prompts, bos=True, eos=False) + past_k, past_v = _empty_past() + past_len = self.model_wrapper.past_len + context_len = self.model_wrapper.max_seq_len + # The prefix buffer holds at most past_len slots, so we can advance + # the position at most past_len times (matching the runner, whose + # seq_len is clamped to context_len). + max_len = min(max_len, past_len) with torch.no_grad(): while token_list[-1] != tokenizer.eos_id and pos < max_len: - cur_pos = torch.tensor([pos], dtype=torch.long) - logits = module(torch.full((1, 1), token_list[pos]), cur_pos) + n_past = min(pos, past_len) + atten_mask = _build_mask(n_past, past_len, context_len) + input_pos = torch.tensor([[n_past]], dtype=torch.int32) + logits, new_k, new_v = module( + torch.full((1, 1), token_list[pos], dtype=torch.int32), + atten_mask, + input_pos, + past_k, + past_v, + ) + # Prefix append: write the new slot into buffer at slot n_past. + for layer in range(self.model_wrapper.num_layers): + past_k[layer][..., :, n_past] = new_k[layer][..., :, 0] + past_v[layer][..., n_past, :] = new_v[layer][..., 0, :] pos += 1 if pos >= len(token_list): token_list.append(torch.argmax(logits, dim=-1).item()) @@ -240,7 +334,6 @@ def calibrate_template( def pt2e_quantize( self, - quant_dtype, fixed_point_type, calibration_tasks, calibration_limit, @@ -251,27 +344,10 @@ def pt2e_quantize( ): self.export() - quantizer = make_quantizer( - quant_dtype=quant_dtype, - per_channel_linear=True, - per_channel_conv=True, - act_observer=MinMaxObserver, - backend=backend, - soc_model=soc_model, - ) - if quant_dtype == QuantDtype.use_16a4w_block: - - def extract_linear_nodes(graph): - linear_nodes = [] - for node in graph.nodes: - if node.target == torch.ops.aten.linear.default: - linear_nodes.append(node) # linear node - linear_nodes.append(node.args[1]) # weight node - return linear_nodes - - linear_nodes = extract_linear_nodes(self.graph_module.graph) - block_size_map = {n.name: (1, 16) for n in linear_nodes} - quantizer.set_block_size_map(block_size_map) + quantizer = make_quantizer(backend=backend, soc_model=soc_model) + quantizer.set_recipe(self.quant_recipe.recipe) + quantizer.set_convert_linear_to_conv2d(True) + self.graph_module = prepare_pt2e(self.graph_module, quantizer) self.pt2e_calibrate( calibration_tasks, @@ -297,11 +373,12 @@ def to_edge_transform_and_lower_to_qnn( compiler_spec = generate_qnn_executorch_compiler_spec( soc_model=get_soc_to_chipset_map()[soc_model], backend_options=backend_options, + use_mha2sha=True, ) with torch.no_grad(): self.edge_prog_mgr = to_edge_transform_and_lower_to_qnn( - self.graph_module, - self.model_wrapper.get_example_inputs(), + {KV_FORWARD: self.graph_module}, + {KV_FORWARD: self.model_wrapper.get_example_inputs()}, compiler_spec, constant_methods=self.model_wrapper.get_metadata(), passes_job=self.passes_job, @@ -310,7 +387,9 @@ def to_edge_transform_and_lower_to_qnn( convert_linear_to_conv2d=True, ) - print_delegation_info(self.edge_prog_mgr.exported_program().graph_module) + print_delegation_info( + self.edge_prog_mgr.exported_program(KV_FORWARD).graph_module + ) if not self.use_fp16: logit_out_shape = { ( @@ -319,7 +398,7 @@ def to_edge_transform_and_lower_to_qnn( self.config.vocab_size, ) } - for n in self.edge_prog_mgr.exported_program().graph.nodes: + for n in self.edge_prog_mgr.exported_program(KV_FORWARD).graph.nodes: if n.op == "output": for node, output_encoding in n.meta[QCOM_QUANT_ATTRS_MAP].items(): if node.meta["val"].size() in logit_out_shape: @@ -332,6 +411,7 @@ def to_executorch(self, artifact, pte_filename): executorch_config = ExecutorchBackendConfig( memory_planning_pass=MemoryPlanningPass( alloc_graph_input=False, + alloc_graph_output=False, ), passes=[BuildQuantIo()], )