From 0a06661203e25492e79fd597c00568bd58194168 Mon Sep 17 00:00:00 2001 From: Xiyue Shi Date: Mon, 26 Oct 2026 15:58:37 -0700 Subject: [PATCH 1/2] Qualcomm AI Engine Direct - Support Gemma RMSNorm recomposition Co-authored-by: Tanya-Motwani --- backends/qualcomm/_passes/layout_transform.py | 37 ++ .../qualcomm/_passes/recompose_rms_norm.py | 236 ++++++++++--- backends/qualcomm/builders/op_rms_norm.py | 10 +- backends/qualcomm/tests/rework/src/pattern.py | 331 +++++++++++++++++- backends/qualcomm/tests/test_passes.py | 106 ++++++ 5 files changed, 659 insertions(+), 61 deletions(-) diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index f2bb7a45f3c..7cbcf1e1238 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -17,6 +17,7 @@ ) 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 from executorch.exir.sym_util import eval_shape @@ -337,7 +338,43 @@ def conditional_sensitive_check(self, node): return True return False + def bypass_view_copy_for_rms_norm(self, gm): + # RecomposeRmsNorm may reshape gamma for broadcasting, but QNN cannot + # partition a view_copy directly feeding the RMSNorm weight. Bypass it + # only when the source parameter already matches normalized_shape. + graph = gm.graph + changed = False + for node in graph.nodes: + if node.op != "call_function": + continue + if node.target != exir_ops.edge.aten.rms_norm.default: + continue + + if len(node.args) < 2: + continue + normalized_shape = node.args[1] + weight = node.args[2] if len(node.args) > 2 else node.kwargs.get("weight") + if ( + weight is not None + and weight.op == "call_function" + and weight.target == exir_ops.edge.aten.view_copy.default + ): + source_weight = weight.args[0] + if not is_parameter(source_weight, self.edge_program): + continue + if tuple(source_weight.meta["val"].shape) != tuple(normalized_shape): + continue + if len(node.args) > 2: + node.args = (*node.args[:2], source_weight, *node.args[3:]) + else: + node.kwargs = {**node.kwargs, "weight": source_weight} + changed = True + + return changed + def call(self, graph_module: torch.fx.GraphModule): + if self.bypass_view_copy_for_rms_norm(graph_module): + dead_code_elimination_pass(graph_module) graph = graph_module.graph sensitive_nodes = [ node diff --git a/backends/qualcomm/_passes/recompose_rms_norm.py b/backends/qualcomm/_passes/recompose_rms_norm.py index 0572fe6b4ee..2a142805943 100644 --- a/backends/qualcomm/_passes/recompose_rms_norm.py +++ b/backends/qualcomm/_passes/recompose_rms_norm.py @@ -3,13 +3,29 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging +import math +import numbers + import torch -from executorch.backends.qualcomm._passes.utils import find_patterns +from executorch.backends.qualcomm._passes.utils import copy_meta, find_patterns from executorch.backends.qualcomm.builders.node_visitor import dq_ops 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 +logger = logging.getLogger(__name__) + +# Export wraps lifted scalar tensor constants before they reach Tensor_Tensor ops. +SCALAR_WRAPPER_TARGETS = { + exir_ops.edge.aten.detach.default, + exir_ops.edge.aten.detach_.default, + exir_ops.edge.aten.lift_fresh_copy.default, + torch.ops.aten.detach.default, + torch.ops.aten.detach_.default, + torch.ops.aten.lift_fresh_copy.default, +} + def _is_node(node): return isinstance(node, torch.fx.Node) @@ -58,6 +74,7 @@ def _is_pow(node): return _is_call(node) and node.target in [ exir_ops.edge.aten.pow.Tensor_Tensor, exir_ops.edge.aten.pow.Tensor_Scalar, + torch.ops.aten.pow.Tensor_Tensor, torch.ops.aten.pow.Tensor_Scalar, ] @@ -69,18 +86,60 @@ def _is_rsqrt(node): ] +def _get_scalar(value, lifted_tensor_constants=None): + while _is_node(value) and value.target in SCALAR_WRAPPER_TARGETS: + value = value.args[0] + if _is_node(value): + node = value + value = ( + lifted_tensor_constants.get(node.name) + if lifted_tensor_constants is not None + else None + ) + if value is None: + value = getattr(node.meta.get("val"), "constant", None) + if isinstance(value, torch.Tensor): + if value.numel() != 1: + return None + value = value.item() + if isinstance(value, numbers.Real) and not isinstance(value, bool): + return value + return None + + +def _is_pow_with_exponent(node, exponent, lifted_tensor_constants=None): + return _is_pow(node) and ( + _get_scalar(node.args[1], lifted_tensor_constants) == exponent + ) + + +def _is_pow_neg_half(node, lifted_tensor_constants=None): + return _is_pow_with_exponent(node, -0.5, lifted_tensor_constants) + + class RecomposeRmsNorm(ExportPass): """ Merge decomposed operators back to one super node. """ - def __init__(self, quantization_capture=False): + def __init__(self, quantization_capture=False, edge_program=None): super(RecomposeRmsNorm, self).__init__() self.rms_norm_target = exir_ops.edge.aten.rms_norm.default self.skip_targets = [ exir_ops.edge.aten.to.dtype, ] self.quantization_capture = quantization_capture + self.lifted_tensor_constants = {} + if edge_program is not None: + # FP export keeps lifted tensor literals in ExportedProgram.constants, + # rather than attaching them to the placeholder's FakeTensor metadata. + for ( + node_name, + constant_name, + ) in edge_program.graph_signature.inputs_to_lifted_tensor_constants.items(): + constant = edge_program.constants.get(constant_name) + if constant is not None: + self.lifted_tensor_constants[node_name] = constant if quantization_capture: self.rms_norm_target = torch.ops.aten.rms_norm.default self.skip_targets = [ @@ -89,10 +148,98 @@ def __init__(self, quantization_capture=False): def _get_input_node(self, node): input_node = node.args[0] - while input_node.target in self.skip_targets: + while input_node.target in self.skip_targets or self._is_fp32_upcast( + input_node + ): input_node = input_node.args[0] return input_node + @staticmethod + def _is_fp32_upcast(node): + # Fold only the no-layout-change fp16-to-fp32 RMSNorm compute upcast. + # Other dim-order copies may change the input's logical layout. + if node.target != exir_ops.edge.dim_order_ops._to_dim_order_copy.default: + return False + + source = node.args[0] + if not _is_node(source): + return False + source_val = source.meta.get("val") + result_val = node.meta.get("val") + if source_val is None or result_val is None: + return False + if source_val.dtype != torch.float16 or result_val.dtype != torch.float32: + return False + if node.kwargs.get("dtype") != torch.float32: + return False + + dim_order = node.kwargs.get("dim_order") + return dim_order is None or list(dim_order) == list(range(source_val.dim())) + + def _is_valid_rms_norm(self, pattern): + _, _, outer_mul, inverse_sqrt, add, mean, square = pattern + if not ( + _is_rsqrt(inverse_sqrt) + or _is_pow_neg_half(inverse_sqrt, self.lifted_tensor_constants) + ): + return False + + if _is_mul(square): + if len(square.args) != 2 or square.args[0] is not square.args[1]: + return False + elif not _is_pow_with_exponent(square, 2, self.lifted_tensor_constants): + return False + + outer_inputs = [arg for arg in outer_mul.args if arg is not inverse_sqrt] + if len(outer_inputs) != 1 or outer_inputs[0] is not square.args[0]: + return False + + dim = mean.args[1] if len(mean.args) > 1 else mean.kwargs.get("dim") + keepdim = ( + mean.args[2] if len(mean.args) > 2 else mean.kwargs.get("keepdim", False) + ) + if isinstance(dim, int): + dims = [dim] + elif isinstance(dim, (list, tuple)): + dims = list(dim) + else: + return False + if len(dims) != 1 or keepdim is not True: + return False + if dims[0] != -1: + input_val = square.args[0].meta.get("val") + if input_val is None or dims[0] != input_val.dim() - 1: + return False + alpha = add.kwargs.get("alpha", add.args[2] if len(add.args) > 2 else 1) + if alpha != 1: + return False + + eps = _get_scalar(add.args[1], self.lifted_tensor_constants) + return eps is not None and math.isfinite(eps) and eps >= 0 + + def _find_rms_norm_pattern(self, node, patterns): + for matched_patterns in find_patterns(node, patterns): + if matched_patterns is None: + continue + candidate = matched_patterns[0] + if self._is_valid_rms_norm(candidate): + return candidate + logger.debug( + "Rejected RMS norm candidate %s: invalid normalization semantics", + candidate[0].name, + ) + return None + + @staticmethod + def _get_rms_norm_weight(last_mul_node): + for arg in last_mul_node.args: + if not (_is_get_attr(arg) or _is_placeholder(arg) or _is_dq(arg)): + continue + arg_val = arg.meta.get("val") + if arg_val is not None and len(arg_val.shape) >= 1: + return arg, list(arg_val.shape) + return None, None + def call(self, graph_module: torch.fx.GraphModule): graph = graph_module.graph @@ -102,54 +249,55 @@ def call(self, graph_module: torch.fx.GraphModule): [_is_mul, "*", _is_mul, _is_rsqrt, _is_add, _is_mean, _is_pow], # executorch.examples.models.llama.norm.RMSNorm [_is_mul, "*", _is_mul, _is_rsqrt, _is_add, _is_mean, _is_mul], + # transformers.models.gemma4.modeling_gemma4.Gemma4RMSNorm + [_is_mul, "*", _is_mul, _is_pow, _is_add, _is_mean, _is_pow], ] for node in graph.nodes: if not _is_mul(node): continue - rms_norm_patterns = [ - pattern - for pattern in find_patterns(node, patterns) - if pattern is not None - ] + rms_norm_pattern = self._find_rms_norm_pattern(node, patterns) + if rms_norm_pattern is None: + continue + + last_mul_node = rms_norm_pattern[0] + # Export represents model weights as get_attrs or placeholders. + # Require a non-scalar operand to exclude scalar multipliers. + gamma_node, gamma_shape = self._get_rms_norm_weight(last_mul_node) + if gamma_node is None: + logger.debug( + "Rejected RMS norm candidate %s: no non-scalar weight", + last_mul_node.name, + ) + continue - if len(rms_norm_patterns) > 0: - # Use first matched pattern - rms_norm_pattern = rms_norm_patterns[0][0] - last_mul_node = rms_norm_pattern[0] - gamma_node = None - # weight should be a constant - for arg in last_mul_node.args: - if _is_get_attr(arg) or _is_placeholder(arg) or _is_dq(arg): - gamma_node = arg - - if gamma_node is None: - continue - - eps = rms_norm_pattern[4].args[1] - if isinstance(eps, torch.fx.Node): - eps = eps.meta["val"].constant.item() - input_node = self._get_input_node(rms_norm_pattern[6]) - - with graph.inserting_before(last_mul_node): - # args schema - # (Tensor input, int[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor - rms_node = graph.create_node( - "call_function", - self.rms_norm_target, - ( - input_node, - list(gamma_node.meta["val"].shape), - gamma_node, - eps, - ), - ) - users = last_mul_node.users.copy() - for user in users: - user.replace_input_with(last_mul_node, rms_node) - # copy metadata - rms_node.meta = last_mul_node.meta + eps = _get_scalar(rms_norm_pattern[4].args[1], self.lifted_tensor_constants) + if eps is None: + logger.debug( + "Rejected RMS norm candidate %s: epsilon is not scalar", + last_mul_node.name, + ) + continue + input_node = self._get_input_node(rms_norm_pattern[6]) + + with graph.inserting_before(last_mul_node): + # args schema + # (Tensor input, int[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor + rms_node = graph.create_node( + "call_function", + self.rms_norm_target, + ( + input_node, + gamma_shape, + gamma_node, + eps, + ), + ) + rms_node.meta = copy_meta(last_mul_node.meta) + users = last_mul_node.users.copy() + for user in users: + user.replace_input_with(last_mul_node, rms_node) dead_code_elimination_pass(graph_module) return PassResult(graph_module, True) diff --git a/backends/qualcomm/builders/op_rms_norm.py b/backends/qualcomm/builders/op_rms_norm.py index e34efed61f3..d4161499fe3 100644 --- a/backends/qualcomm/builders/op_rms_norm.py +++ b/backends/qualcomm/builders/op_rms_norm.py @@ -61,9 +61,10 @@ def define_node( axes = [node.args[0].meta["val"].dim() - 1] axes_shape = [len(axes)] - has_weight = len(node.args) > 2 and node.args[2] is not None + weight = node.args[2] if len(node.args) > 2 else node.kwargs.get("weight") + has_weight = weight is not None if has_weight: - weight_node = self.get_node(node.args[2]) + weight_node = self.get_node(weight) weight_tensor = get_parameter(weight_node, self.edge_program) else: # elementwise_affine=False: use all-ones weight as identity @@ -111,8 +112,9 @@ def define_node( ) epsilon = torch.finfo(torch.float32).eps - if len(node.args) > 3: - epsilon = node.args[3] + eps = node.args[3] if len(node.args) > 3 else node.kwargs.get("eps") + if eps is not None: + epsilon = eps output_tensor = self.get_tensor(node, node) output_tensor_wrapper = self.define_tensor( node, diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 7b90fde61e2..3cbeea0f661 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/backends/qualcomm/tests/rework/src/pattern.py @@ -4393,6 +4393,34 @@ def __init__(self): def forward(self, x): return self.relu(self.conv(x)) + class _RmsNormParameterView(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return torch.ops.aten.rms_norm.default( + x, [self.weight.numel()], self.weight.view_as(self.weight), 1e-6 + ) + + class _RmsNormReshapedParameterView(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2, 8)) + + def forward(self, x): + return torch.ops.aten.rms_norm.default(x, [16], self.weight.view(16), 1e-6) + + class _RmsNormRuntimeWeightView(torch.nn.Module): + def forward(self, x, weight): + return torch.ops.aten.rms_norm.default( + x, [weight.numel()], weight.view_as(weight), 1e-6 + ) + + class _RmsNormWithoutOptionalArgs(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.rms_norm.default(x, [x.shape[-1]]) + @staticmethod @unpack_pass_fixtures def test( @@ -4439,6 +4467,70 @@ def test( assertions.assert_target_count(gm, permute, 2) assertions.assert_meta_key_on_targets(gm, permute, QCOM_INSERTED_PERMUTE) + if quantizer is None: + with subtests.test(msg="rms_norm_parameter_view"): + gm = pass_pipeline.lower_edge_ep( + module=LayoutTransform._RmsNormParameterView(16), + sample_input=(torch.randn(1, 4, 16),), + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=target_pass, + quantizer=quantizer, + ).graph_module + rms_norm = next( + node + for node in gm.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ) + assert rms_norm.args[1] == [16] + assert rms_norm.args[2].op == "placeholder" + + with subtests.test(msg="rms_norm_reshaped_parameter_view"): + gm = pass_pipeline.lower_edge_ep( + module=LayoutTransform._RmsNormReshapedParameterView(), + sample_input=(torch.randn(1, 4, 16),), + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=target_pass, + quantizer=quantizer, + ).graph_module + rms_norm = next( + node + for node in gm.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ) + assert rms_norm.args[1] == [16] + assert rms_norm.args[2].target == exir_ops.edge.aten.view_copy.default + + with subtests.test(msg="rms_norm_runtime_weight_view"): + gm = pass_pipeline.lower_edge_ep( + module=LayoutTransform._RmsNormRuntimeWeightView(), + sample_input=(torch.randn(1, 4, 16), torch.ones(16)), + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=target_pass, + quantizer=quantizer, + ).graph_module + rms_norm = next( + node + for node in gm.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ) + assert rms_norm.args[2].target == exir_ops.edge.aten.view_copy.default + + with subtests.test(msg="rms_norm_without_optional_args"): + gm = pass_pipeline.lower_edge_ep( + module=LayoutTransform._RmsNormWithoutOptionalArgs(), + sample_input=(torch.randn(1, 4, 16),), + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=target_pass, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count( + gm, exir_ops.edge.aten.rms_norm.default, 1 + ) + class LiftConstantScalarOperands: class _Pow(torch.nn.Module): @@ -4828,15 +4920,73 @@ def test( class RecomposeRmsNorm: class _RmsNorm(torch.nn.Module): - def __init__(self, dim: int, eps: float = 1e-6): + def __init__( + self, + dim: int, + eps: float = 1e-6, + square_exponent: int = 2, + reduction_dim: int = -1, + keepdim: bool = True, + ): super().__init__() self.weight = torch.nn.Parameter(torch.ones(dim)) self.eps = eps + self.square_exponent = square_exponent + self.reduction_dim = reduction_dim + self.keepdim = keepdim def forward(self, x): - norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + norm = x * torch.rsqrt( + x.pow(self.square_exponent).mean( + self.reduction_dim, keepdim=self.keepdim + ) + + self.eps + ) + return norm * self.weight + + class _GemmaRmsNorm(_RmsNorm): + def __init__( + self, + dim: int, + square_exponent: int = 2, + reduction_dim: int = -1, + negative_eps: bool = False, + tensor_eps: bool = False, + tensor_exponent: bool = False, + ): + super().__init__(dim) + self.square_exponent = square_exponent + self.reduction_dim = reduction_dim + self.negative_eps = negative_eps + self.tensor_eps = tensor_eps + self.tensor_exponent = tensor_exponent + + def forward(self, x, statistic_input=None): + eps = torch.tensor(self.eps) if self.tensor_eps else self.eps + exponent = torch.tensor(-0.5) if self.tensor_exponent else -0.5 + statistic_input = x if statistic_input is None else statistic_input + statistic = statistic_input.pow(self.square_exponent).mean( + self.reduction_dim, keepdim=True + ) + eps = -eps if self.negative_eps else eps + statistic = statistic + eps + norm = x * torch.pow(statistic, exponent) + return norm * self.weight + + class _RsqrtTensorSquareExponent(_RmsNorm): + def forward(self, x): + norm = x * torch.rsqrt( + torch.pow(x, torch.tensor(4)).mean(-1, keepdim=True) + self.eps + ) return norm * self.weight + class _Fp32RmsNorm(_RmsNorm): + def forward(self, x): + input_dtype = x.dtype + x = x.float() + norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return self.weight * norm.to(input_dtype) + @staticmethod @unpack_pass_fixtures def test( @@ -4845,20 +4995,175 @@ def test( backend_type: QnnExecuTorchBackendType, assertions: Assertions, pass_pipeline: PassPipeline, + subtests, ): dim = 16 inputs = (torch.randn(1, 4, 8, dim),) - gm = pass_pipeline.lower_edge_ep( - module=RecomposeRmsNorm._RmsNorm(dim), - sample_input=inputs, - backend_type=backend_type, - compile_spec=compile_spec, - target_pass=_passes.RecomposeRmsNorm, - quantizer=quantizer, - ).graph_module - assertions.assert_target_count(gm, exir_ops.edge.aten.rms_norm.default, 1) - assertions.assert_no_target(gm, exir_ops.edge.aten.rsqrt.default) - assertions.assert_no_target(gm, exir_ops.edge.aten.mean.dim) + for name, module in ( + ("rsqrt", RecomposeRmsNorm._RmsNorm(dim)), + ("gemma_pow_neg_half", RecomposeRmsNorm._GemmaRmsNorm(dim)), + ): + with subtests.test(msg=name): + gm = pass_pipeline.lower_edge_ep( + module=module, + sample_input=inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count( + gm, exir_ops.edge.aten.rms_norm.default, 1 + ) + assertions.assert_no_target(gm, exir_ops.edge.aten.mean.dim) + if name == "rsqrt": + assertions.assert_no_target(gm, exir_ops.edge.aten.rsqrt.default) + + with subtests.test(msg="gemma_positive_last_reduction_dim"): + gm = pass_pipeline.lower_edge_ep( + module=RecomposeRmsNorm._GemmaRmsNorm(dim, reduction_dim=3), + sample_input=inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.rms_norm.default, 1) + assertions.assert_no_target(gm, exir_ops.edge.aten.mean.dim) + + with subtests.test(msg="gemma_tensor_exponent"): + gm = pass_pipeline.lower_edge_ep( + module=RecomposeRmsNorm._GemmaRmsNorm(dim, tensor_exponent=True), + sample_input=inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.rms_norm.default, 1) + assertions.assert_no_target(gm, exir_ops.edge.aten.mean.dim) + + with subtests.test(msg="gemma_tensor_eps"): + gm = pass_pipeline.lower_edge_ep( + module=RecomposeRmsNorm._GemmaRmsNorm(dim, tensor_eps=True), + sample_input=inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.rms_norm.default, 1) + assertions.assert_no_target(gm, exir_ops.edge.aten.mean.dim) + + with subtests.test(msg="rsqrt_tensor_square_exponent"): + gm = pass_pipeline.lower_edge_ep( + module=RecomposeRmsNorm._RsqrtTensorSquareExponent(dim), + sample_input=inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.rms_norm.default, 0) + assertions.assert_target_count(gm, exir_ops.edge.aten.mean.dim, 1) + + rsqrt_invalid_inputs = (torch.randn(1, 1, 1, dim),) + for name, module in ( + ( + "rsqrt_wrong_square_exponent", + RecomposeRmsNorm._RmsNorm(dim, square_exponent=4), + ), + ( + "rsqrt_wrong_reduction_dim", + RecomposeRmsNorm._RmsNorm(dim, reduction_dim=1), + ), + ( + "rsqrt_without_keepdim", + RecomposeRmsNorm._RmsNorm(dim, keepdim=False), + ), + ): + with subtests.test(msg=name): + gm = pass_pipeline.lower_edge_ep( + module=module, + sample_input=rsqrt_invalid_inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count( + gm, exir_ops.edge.aten.rms_norm.default, 0 + ) + assertions.assert_target_count(gm, exir_ops.edge.aten.mean.dim, 1) + + for name, module in ( + ( + "gemma_wrong_square_exponent", + RecomposeRmsNorm._GemmaRmsNorm(dim, square_exponent=4), + ), + ( + "gemma_wrong_reduction_dim", + RecomposeRmsNorm._GemmaRmsNorm(dim, reduction_dim=0), + ), + ( + "gemma_negative_eps", + RecomposeRmsNorm._GemmaRmsNorm(dim, negative_eps=True), + ), + ): + with subtests.test(msg=name): + gm = pass_pipeline.lower_edge_ep( + module=module, + sample_input=inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count( + gm, exir_ops.edge.aten.rms_norm.default, 0 + ) + assertions.assert_target_count(gm, exir_ops.edge.aten.mean.dim, 1) + if name == "gemma_negative_eps": + assertions.assert_target_count( + gm, + { + exir_ops.edge.aten.add.Scalar, + exir_ops.edge.aten.add.Tensor, + }, + 1, + ) + + with subtests.test(msg="gemma_different_square_input"): + gm = pass_pipeline.lower_edge_ep( + module=RecomposeRmsNorm._GemmaRmsNorm(dim), + sample_input=(torch.randn(1, 4, 8, dim), torch.randn(1, 4, 8, dim)), + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.rms_norm.default, 0) + assertions.assert_target_count(gm, exir_ops.edge.aten.mean.dim, 1) + + if quantizer is None: + with subtests.test(msg="fp32_upcast"): + gm = pass_pipeline.lower_edge_ep( + module=RecomposeRmsNorm._Fp32RmsNorm(dim), + sample_input=(torch.randn(1, 4, 8, dim, dtype=torch.float16),), + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + rms_norm = next( + node + for node in gm.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ) + assert rms_norm.args[0].op == "placeholder" + assert rms_norm.args[0].meta["val"].dtype == torch.float16 + assert rms_norm.args[2].meta["val"].dtype == torch.float32 + assert rms_norm.meta["val"].dtype == torch.float32 class RemoveRedundancy: diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 606fed01d72..8d6b1adcc4a 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -44,6 +44,7 @@ Unfold, ) from executorch.backends.qualcomm.utils.constants import ( + QCOM_DATA, QCOM_ENCODING, QCOM_QUANT_ATTRS, QCOM_SCALE, @@ -172,6 +173,111 @@ def test_context_loader_op_lowers_with_ir_validation(self): ) self.assertEqual(1, len(context_loader_nodes)) + def test_rms_norm_builder_reads_keyword_weight_and_eps(self): + from executorch.backends.qualcomm.builders.node_visitor_manager import ( + _node_visitor_dict, + ) + from executorch.backends.qualcomm.builders.qnn_constants import OpRmsNorm + + for eps, expected_eps in ( + (0.125, 0.125), + (None, torch.finfo(torch.float32).eps), + ): + with self.subTest(eps=eps): + graph = torch.fx.Graph() + input_node = graph.placeholder("input") + input_node.meta["val"] = torch.empty(1, 4, 16) + weight_node = graph.placeholder("weight") + node = graph.call_function( + exir_ops.edge.aten.rms_norm.default, + (input_node, [16]), + {"weight": weight_node, "eps": eps}, + ) + node.meta["val"] = torch.empty(1, 4, 16) + edge_program = MagicMock() + visitor = _node_visitor_dict["aten.rms_norm.default"]( + {}, edge_program, False + ) + visitor.get_node = MagicMock(side_effect=lambda value: value) + visitor.get_tensor = MagicMock(side_effect=[torch.empty(1, 4, 16)] * 2) + visitor.define_tensor = MagicMock(side_effect=[MagicMock()] * 4) + + with ( + patch( + "executorch.backends.qualcomm.builders.op_rms_norm." + "get_parameter", + return_value=torch.arange(1, 17, dtype=torch.float32), + ) as get_parameter, + patch( + "executorch.backends.qualcomm.builders.op_rms_norm." + "PyQnnManager.PyQnnOpWrapper", + ) as op_wrapper, + ): + visitor.define_node(node, {}) + + get_parameter.assert_called_once_with(weight_node, edge_program) + epsilon_call = next( + call + for call in op_wrapper.return_value.AddScalarParam.call_args_list + if call.args[0] == OpRmsNorm.param_epsilon + ) + self.assertEqual(expected_eps, epsilon_call.args[2][QCOM_DATA]) + + def test_rms_norm_recomposition_rejects_missing_upcast_metadata(self): + from executorch.backends.qualcomm._passes.recompose_rms_norm import ( + RecomposeRmsNorm, + ) + + graph = torch.fx.Graph() + source = graph.placeholder("source") + upcast = graph.call_function( + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + (source,), + {"dtype": torch.float32}, + ) + + self.assertFalse(RecomposeRmsNorm._is_fp32_upcast(upcast)) + + def test_rms_norm_recomposition_rejects_scaled_eps(self): + from executorch.backends.qualcomm._passes.recompose_rms_norm import ( + RecomposeRmsNorm, + ) + + graph = torch.fx.Graph() + input_node = graph.placeholder("input") + square = graph.call_function(torch.ops.aten.pow.Tensor_Scalar, (input_node, 2)) + mean = graph.call_function(torch.ops.aten.mean.dim, (square, [-1], True)) + add = graph.call_function( + torch.ops.aten.add.Tensor, + (mean, 1e-6), + {"alpha": 2}, + ) + inverse_sqrt = graph.call_function( + torch.ops.aten.pow.Tensor_Scalar, (add, -0.5) + ) + outer_mul = graph.call_function( + torch.ops.aten.mul.Tensor, (input_node, inverse_sqrt) + ) + last_mul = graph.call_function( + torch.ops.aten.mul.Tensor, (outer_mul, input_node) + ) + + self.assertFalse( + RecomposeRmsNorm()._is_valid_rms_norm( + (last_mul, None, outer_mul, inverse_sqrt, add, mean, square) + ) + ) + + def test_rms_norm_recomposition_unwraps_literal_scalar(self): + from executorch.backends.qualcomm._passes.recompose_rms_norm import _get_scalar + + graph = torch.fx.Graph() + scalar = graph.call_function( + torch.ops.aten.lift_fresh_copy.default, (torch.tensor(-0.5),) + ) + + self.assertEqual(_get_scalar(scalar), -0.5) + def _build_quantized_graph(self): """Build a quantized graph through AnnotateQuantAttrs + FoldQDQ.""" From ea96a9bb9fda6bc250debceedd8b2ed2b5ec0a34 Mon Sep 17 00:00:00 2001 From: Xiyue Shi Date: Fri, 30 Oct 2026 10:23:09 -0700 Subject: [PATCH 2/2] Address code review feedbacks --- backends/qualcomm/_passes/layout_transform.py | 37 -- .../qualcomm/_passes/recompose_rms_norm.py | 282 ++++++++---- .../qualcomm/_passes/remove_redundancy.py | 21 +- backends/qualcomm/builders/op_rms_norm.py | 13 +- backends/qualcomm/tests/models.py | 12 + backends/qualcomm/tests/rework/passes/test.py | 6 + backends/qualcomm/tests/rework/src/pattern.py | 414 +++++++++++++----- backends/qualcomm/tests/test_passes.py | 57 +-- backends/qualcomm/tests/test_qnn_delegate.py | 9 + 9 files changed, 573 insertions(+), 278 deletions(-) diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index 7cbcf1e1238..f2bb7a45f3c 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -17,7 +17,6 @@ ) 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 from executorch.exir.sym_util import eval_shape @@ -338,43 +337,7 @@ def conditional_sensitive_check(self, node): return True return False - def bypass_view_copy_for_rms_norm(self, gm): - # RecomposeRmsNorm may reshape gamma for broadcasting, but QNN cannot - # partition a view_copy directly feeding the RMSNorm weight. Bypass it - # only when the source parameter already matches normalized_shape. - graph = gm.graph - changed = False - for node in graph.nodes: - if node.op != "call_function": - continue - if node.target != exir_ops.edge.aten.rms_norm.default: - continue - - if len(node.args) < 2: - continue - normalized_shape = node.args[1] - weight = node.args[2] if len(node.args) > 2 else node.kwargs.get("weight") - if ( - weight is not None - and weight.op == "call_function" - and weight.target == exir_ops.edge.aten.view_copy.default - ): - source_weight = weight.args[0] - if not is_parameter(source_weight, self.edge_program): - continue - if tuple(source_weight.meta["val"].shape) != tuple(normalized_shape): - continue - if len(node.args) > 2: - node.args = (*node.args[:2], source_weight, *node.args[3:]) - else: - node.kwargs = {**node.kwargs, "weight": source_weight} - changed = True - - return changed - def call(self, graph_module: torch.fx.GraphModule): - if self.bypass_view_copy_for_rms_norm(graph_module): - dead_code_elimination_pass(graph_module) graph = graph_module.graph sensitive_nodes = [ node diff --git a/backends/qualcomm/_passes/recompose_rms_norm.py b/backends/qualcomm/_passes/recompose_rms_norm.py index 2a142805943..ea063d64e2b 100644 --- a/backends/qualcomm/_passes/recompose_rms_norm.py +++ b/backends/qualcomm/_passes/recompose_rms_norm.py @@ -10,9 +10,11 @@ import torch from executorch.backends.qualcomm._passes.utils import copy_meta, find_patterns from executorch.backends.qualcomm.builders.node_visitor import dq_ops +from executorch.backends.qualcomm.builders.utils import is_parameter 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 +from torch._subclasses.fake_tensor import FakeTensor logger = logging.getLogger(__name__) @@ -26,6 +28,22 @@ torch.ops.aten.lift_fresh_copy.default, } +TO_DTYPE_TARGETS = { + exir_ops.edge.aten.to.dtype, + torch.ops.aten.to.dtype, +} + +DIM_ORDER_COPY_TARGETS = { + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, +} + +VIEW_TARGETS = { + exir_ops.edge.aten.view_copy.default, + torch.ops.aten.view.default, + torch.ops.aten.view_as.default, + torch.ops.aten.view_copy.default, +} + def _is_node(node): return isinstance(node, torch.fx.Node) @@ -35,10 +53,6 @@ def _is_call(node): return _is_node(node) and node.op == "call_function" -def _is_placeholder(node): - return _is_node(node) and node.op == "placeholder" - - def _is_get_attr(node): return _is_node(node) and node.op == "get_attr" @@ -86,6 +100,30 @@ def _is_rsqrt(node): ] +def _tensor_dtype(node): + value = node.meta.get("val") if _is_node(node) else None + return value.dtype if value is not None else None + + +def _is_cast_wrapper(node): + if not _is_call(node) or node.target not in ( + TO_DTYPE_TARGETS | DIM_ORDER_COPY_TARGETS + ): + return False + source = node.args[0] + source_value = source.meta.get("val") if _is_node(source) else None + result_value = node.meta.get("val") + if source_value is None or result_value is None: + return False + if source_value.shape != result_value.shape: + return False + if node.target in DIM_ORDER_COPY_TARGETS: + dim_order = node.kwargs.get("dim_order") + if dim_order is not None and list(dim_order) != list(range(source_value.dim())): + return False + return True + + def _get_scalar(value, lifted_tensor_constants=None): while _is_node(value) and value.target in SCALAR_WRAPPER_TARGETS: value = value.args[0] @@ -97,7 +135,9 @@ def _get_scalar(value, lifted_tensor_constants=None): else None ) if value is None: - value = getattr(node.meta.get("val"), "constant", None) + node_val = node.meta.get("val") + if isinstance(node_val, FakeTensor): + value = node_val.constant if isinstance(value, torch.Tensor): if value.numel() != 1: return None @@ -113,10 +153,6 @@ def _is_pow_with_exponent(node, exponent, lifted_tensor_constants=None): ) -def _is_pow_neg_half(node, lifted_tensor_constants=None): - return _is_pow_with_exponent(node, -0.5, lifted_tensor_constants) - - class RecomposeRmsNorm(ExportPass): """ Merge decomposed operators back to one super node. @@ -125,10 +161,8 @@ class RecomposeRmsNorm(ExportPass): def __init__(self, quantization_capture=False, edge_program=None): super(RecomposeRmsNorm, self).__init__() self.rms_norm_target = exir_ops.edge.aten.rms_norm.default - self.skip_targets = [ - exir_ops.edge.aten.to.dtype, - ] self.quantization_capture = quantization_capture + self.edge_program = edge_program self.lifted_tensor_constants = {} if edge_program is not None: # FP export keeps lifted tensor literals in ExportedProgram.constants, @@ -142,58 +176,57 @@ def __init__(self, quantization_capture=False, edge_program=None): self.lifted_tensor_constants[node_name] = constant if quantization_capture: self.rms_norm_target = torch.ops.aten.rms_norm.default - self.skip_targets = [ - torch.ops.aten.to.dtype, - ] - - def _get_input_node(self, node): - input_node = node.args[0] - while input_node.target in self.skip_targets or self._is_fp32_upcast( - input_node - ): - input_node = input_node.args[0] - return input_node - - @staticmethod - def _is_fp32_upcast(node): - # Fold only the no-layout-change fp16-to-fp32 RMSNorm compute upcast. - # Other dim-order copies may change the input's logical layout. - if node.target != exir_ops.edge.dim_order_ops._to_dim_order_copy.default: - return False - - source = node.args[0] - if not _is_node(source): - return False - source_val = source.meta.get("val") - result_val = node.meta.get("val") - if source_val is None or result_val is None: - return False - if source_val.dtype != torch.float16 or result_val.dtype != torch.float32: - return False - if node.kwargs.get("dtype") != torch.float32: - return False - dim_order = node.kwargs.get("dim_order") - return dim_order is None or list(dim_order) == list(range(source_val.dim())) + def _get_rms_norm_epsilon(self, add, mean): + if len(add.args) < 2: + return None + if add.args[0] is mean: + eps_arg = add.args[1] + elif add.args[1] is mean: + eps_arg = add.args[0] + else: + return None + return _get_scalar(eps_arg, self.lifted_tensor_constants) - def _is_valid_rms_norm(self, pattern): - _, _, outer_mul, inverse_sqrt, add, mean, square = pattern - if not ( - _is_rsqrt(inverse_sqrt) - or _is_pow_neg_half(inverse_sqrt, self.lifted_tensor_constants) - ): + def _is_fp32_upcast(self, node): + if not _is_cast_wrapper(node): return False + return ( + _tensor_dtype(node.args[0]) in {torch.float16, torch.bfloat16} + and _tensor_dtype(node) == torch.float32 + ) - if _is_mul(square): - if len(square.args) != 2 or square.args[0] is not square.args[1]: - return False - elif not _is_pow_with_exponent(square, 2, self.lifted_tensor_constants): - return False + def _get_wildcard_path(self, last_mul, outer_mul): + def find_path(node, remaining): + if node is outer_mul: + return [] + if remaining == 0 or not _is_cast_wrapper(node): + return None + path = find_path(node.args[0], remaining - 1) + return None if path is None else [node, *path] + + for arg in last_mul.args: + path = find_path(arg, 3) + if path is not None: + return path + return None - outer_inputs = [arg for arg in outer_mul.args if arg is not inverse_sqrt] - if len(outer_inputs) != 1 or outer_inputs[0] is not square.args[0]: + def _is_known_output_cast_path(self, path, square): + if not path: + return True + normalized_input = square.args[0] + if not self._is_fp32_upcast(normalized_input): + return False + input_dtype = _tensor_dtype(normalized_input.args[0]) + if _tensor_dtype(path[0]) != input_dtype: return False + return all( + _tensor_dtype(node.args[0]) in {torch.float32, input_dtype} + and _tensor_dtype(node) in {torch.float32, input_dtype} + for node in path + ) + def _is_last_dim_reduction(self, mean, input_node): dim = mean.args[1] if len(mean.args) > 1 else mean.kwargs.get("dim") keepdim = ( mean.args[2] if len(mean.args) > 2 else mean.kwargs.get("keepdim", False) @@ -206,38 +239,103 @@ def _is_valid_rms_norm(self, pattern): return False if len(dims) != 1 or keepdim is not True: return False - if dims[0] != -1: - input_val = square.args[0].meta.get("val") - if input_val is None or dims[0] != input_val.dim() - 1: - return False + if dims[0] == -1: + return True + input_val = input_node.meta.get("val") + return input_val is not None and dims[0] == input_val.dim() - 1 + + def _validate_rms_norm(self, pattern): + last_mul, _, outer_mul, inverse_sqrt, add, mean, square = pattern + if not ( + _is_rsqrt(inverse_sqrt) + or _is_pow_with_exponent(inverse_sqrt, -0.5, self.lifted_tensor_constants) + ): + return None + + if _is_mul(square): + if len(square.args) != 2 or square.args[0] is not square.args[1]: + return None + elif not _is_pow_with_exponent(square, 2, self.lifted_tensor_constants): + return None + + outer_inputs = [arg for arg in outer_mul.args if arg is not inverse_sqrt] + if len(outer_inputs) != 1 or outer_inputs[0] is not square.args[0]: + return None + + wildcard_path = self._get_wildcard_path(last_mul, outer_mul) + if wildcard_path is None or not self._is_known_output_cast_path( + wildcard_path, square + ): + return None + + if not self._is_last_dim_reduction(mean, square.args[0]): + return None alpha = add.kwargs.get("alpha", add.args[2] if len(add.args) > 2 else 1) if alpha != 1: - return False + return None + + eps = self._get_rms_norm_epsilon(add, mean) + if eps is None or not math.isfinite(eps) or eps < 0: + return None + return eps - eps = _get_scalar(add.args[1], self.lifted_tensor_constants) - return eps is not None and math.isfinite(eps) and eps >= 0 + def _is_valid_rms_norm(self, pattern): + return self._validate_rms_norm(pattern) is not None def _find_rms_norm_pattern(self, node, patterns): for matched_patterns in find_patterns(node, patterns): if matched_patterns is None: continue - candidate = matched_patterns[0] - if self._is_valid_rms_norm(candidate): - return candidate - logger.debug( - "Rejected RMS norm candidate %s: invalid normalization semantics", - candidate[0].name, - ) + for candidate in matched_patterns: + eps = self._validate_rms_norm(candidate) + if eps is not None: + return candidate, eps + logger.debug( + "Rejected RMS norm candidate %s: invalid normalization semantics", + candidate[0].name, + ) return None - @staticmethod - def _get_rms_norm_weight(last_mul_node): + def _is_static_weight(self, node): + source = node.args[0] if _is_dq(node) else node + return _is_get_attr(source) or ( + self.edge_program is not None and is_parameter(source, self.edge_program) + ) + + def _get_static_weight(self, node): + if self._is_static_weight(node): + return node + if not _is_call(node) or node.target not in VIEW_TARGETS: + return None + + source = node.args[0] + source_val = source.meta.get("val") if _is_node(source) else None + node_val = node.meta.get("val") + if ( + source_val is None + or node_val is None + or source_val.shape != node_val.shape + or not self._is_static_weight(source) + ): + return None + return source + + def _get_rms_norm_weight(self, last_mul_node, outer_mul, input_node): for arg in last_mul_node.args: - if not (_is_get_attr(arg) or _is_placeholder(arg) or _is_dq(arg)): + if arg is outer_mul or not _is_node(arg): + continue + weight = self._get_static_weight(arg) + if weight is None: continue - arg_val = arg.meta.get("val") - if arg_val is not None and len(arg_val.shape) >= 1: - return arg, list(arg_val.shape) + arg_val = weight.meta.get("val") + input_val = input_node.meta.get("val") + if ( + arg_val is not None + and input_val is not None + and len(arg_val.shape) == 1 + and list(arg_val.shape) == [input_val.shape[-1]] + ): + return weight, list(arg_val.shape) return None, None def call(self, graph_module: torch.fx.GraphModule): @@ -245,7 +343,8 @@ def call(self, graph_module: torch.fx.GraphModule): # Root Mean Square normalization math equivalent implementation patterns = [ - # transformers.models.qwen2.modeling_qwen2.Qwen2RMSNorm + # transformers.models.qwen2.modeling_qwen2.Qwen2RMSNorm and + # transformers.models.qwen3.modeling_qwen3.Qwen3RMSNorm [_is_mul, "*", _is_mul, _is_rsqrt, _is_add, _is_mean, _is_pow], # executorch.examples.models.llama.norm.RMSNorm [_is_mul, "*", _is_mul, _is_rsqrt, _is_add, _is_mean, _is_mul], @@ -257,14 +356,22 @@ def call(self, graph_module: torch.fx.GraphModule): if not _is_mul(node): continue - rms_norm_pattern = self._find_rms_norm_pattern(node, patterns) - if rms_norm_pattern is None: + match = self._find_rms_norm_pattern(node, patterns) + if match is None: continue + rms_norm_pattern, eps = match last_mul_node = rms_norm_pattern[0] - # Export represents model weights as get_attrs or placeholders. - # Require a non-scalar operand to exclude scalar multipliers. - gamma_node, gamma_shape = self._get_rms_norm_weight(last_mul_node) + outer_mul = rms_norm_pattern[2] + normalized_input = rms_norm_pattern[-1].args[0] + input_node = ( + normalized_input.args[0] + if self._is_fp32_upcast(normalized_input) + else normalized_input + ) + gamma_node, gamma_shape = self._get_rms_norm_weight( + last_mul_node, outer_mul, input_node + ) if gamma_node is None: logger.debug( "Rejected RMS norm candidate %s: no non-scalar weight", @@ -272,15 +379,6 @@ def call(self, graph_module: torch.fx.GraphModule): ) continue - eps = _get_scalar(rms_norm_pattern[4].args[1], self.lifted_tensor_constants) - if eps is None: - logger.debug( - "Rejected RMS norm candidate %s: epsilon is not scalar", - last_mul_node.name, - ) - continue - input_node = self._get_input_node(rms_norm_pattern[6]) - with graph.inserting_before(last_mul_node): # args schema # (Tensor input, int[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor diff --git a/backends/qualcomm/_passes/remove_redundancy.py b/backends/qualcomm/_passes/remove_redundancy.py index 7589d40b745..2bbcaf17853 100644 --- a/backends/qualcomm/_passes/remove_redundancy.py +++ b/backends/qualcomm/_passes/remove_redundancy.py @@ -4,6 +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. import torch +from executorch.backends.qualcomm.builders.utils import is_parameter 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 @@ -14,8 +15,9 @@ class RemoveRedundancy(ExportPass): Trim certain operators to reduce unnecessary overhead. """ - def __init__(self, quantization_capture=False): + def __init__(self, quantization_capture=False, edge_program=None): super(RemoveRedundancy, self).__init__() + self.edge_program = edge_program self.redundant_ops_general = { torch.clone: self._default_condition, torch.ops.aten.clone.default: self._default_condition, @@ -26,6 +28,7 @@ def __init__(self, quantization_capture=False): exir_ops.edge.aten.expand_copy.default: self._same_shape_condition, exir_ops.edge.aten.lift_fresh_copy.default: self._default_condition, exir_ops.edge.aten.repeat.default: self._same_shape_condition, + exir_ops.edge.aten.view_copy.default: self._rms_norm_weight_view_condition, # remove this target if '_skip_dim_order' is set to False exir_ops.edge.dim_order_ops._to_dim_order_copy.default: self._dim_order_op_condition, # remove channel_last / contiguous _to_copy if '_skip_dim_order' is set to True @@ -48,6 +51,22 @@ def _dim_order_op_condition(self, node): def _same_shape_condition(self, node): return node.args[0].meta["val"].shape == node.meta["val"].shape + def _rms_norm_weight_view_condition(self, node): + source = node.args[0] + is_static = source.op == "get_attr" or ( + self.edge_program is not None and is_parameter(source, self.edge_program) + ) + if not is_static or not self._same_shape_condition(node): + return False + return all( + user.target == exir_ops.edge.aten.rms_norm.default + and ( + (len(user.args) > 2 and user.args[2] is node) + or user.kwargs.get("weight") is node + ) + for user in node.users + ) + def _to_copy_op_condition(self, node): return "memory_format" in node.kwargs diff --git a/backends/qualcomm/builders/op_rms_norm.py b/backends/qualcomm/builders/op_rms_norm.py index d4161499fe3..d9cebe8a5f3 100644 --- a/backends/qualcomm/builders/op_rms_norm.py +++ b/backends/qualcomm/builders/op_rms_norm.py @@ -19,7 +19,7 @@ ) from executorch.exir.dialects._ops import ops as exir_ops -from .node_visitor import NodeVisitor +from .node_visitor import NodeVisitor, q_ops from .node_visitor_manager import register_node_visitor from .qnn_constants import OpRmsNorm, QNN_OP_PACKAGE_NAME_QTI_AISW @@ -51,7 +51,7 @@ def define_node( normalized_shapes = node.args[1] if ( len(normalized_shapes) != 1 - and normalized_shapes[0] != input_tensor.shape[-1] + or normalized_shapes[0] != input_tensor.shape[-1] ): warnings.warn( "[QNN Delegate Op Builder]: Only supports normalization with last input dimension.", @@ -111,10 +111,17 @@ def define_node( nodes_to_wrappers, ) - epsilon = torch.finfo(torch.float32).eps eps = node.args[3] if len(node.args) > 3 else node.kwargs.get("eps") if eps is not None: epsilon = eps + else: + logical_input = node.args[0] + input_dtype = logical_input.meta["val"].dtype + if logical_input.target in q_ops: + input_dtype = logical_input.args[0].meta["val"].dtype + if not input_dtype.is_floating_point: + input_dtype = node.meta["val"].dtype + epsilon = torch.finfo(input_dtype).eps output_tensor = self.get_tensor(node, node) output_tensor_wrapper = self.define_tensor( node, diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 94d875617c7..b944e8cb345 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -2420,6 +2420,18 @@ def forward(self, x): return self.rms(x) +class Fp32RmsNorm(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def forward(self, x): + input_dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) + return x.to(input_dtype) * self.weight.to(input_dtype) + + class Roll(torch.nn.Module): def __init__(self, shifts, dims=None): super().__init__() diff --git a/backends/qualcomm/tests/rework/passes/test.py b/backends/qualcomm/tests/rework/passes/test.py index 1bf2c76d28c..f442a287258 100644 --- a/backends/qualcomm/tests/rework/passes/test.py +++ b/backends/qualcomm/tests/rework/passes/test.py @@ -398,6 +398,12 @@ def test_recompose_rms_norm(request, kwargs): RecomposeRmsNorm.test(request, kwargs) # noqa: F405 +@enumerate_backends_quantized() +@repack_pass_fixtures +def test_recompose_rms_norm_annotation(request, kwargs): + RecomposeRmsNorm.test_annotation(request, kwargs) # noqa: F405 + + @enumerate_backends() @repack_pass_fixtures def test_remove_0d_tensor(request, kwargs): diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 3cbeea0f661..3df3de8459d 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/backends/qualcomm/tests/rework/src/pattern.py @@ -4393,34 +4393,6 @@ def __init__(self): def forward(self, x): return self.relu(self.conv(x)) - class _RmsNormParameterView(torch.nn.Module): - def __init__(self, dim: int): - super().__init__() - self.weight = torch.nn.Parameter(torch.ones(dim)) - - def forward(self, x): - return torch.ops.aten.rms_norm.default( - x, [self.weight.numel()], self.weight.view_as(self.weight), 1e-6 - ) - - class _RmsNormReshapedParameterView(torch.nn.Module): - def __init__(self): - super().__init__() - self.weight = torch.nn.Parameter(torch.ones(2, 8)) - - def forward(self, x): - return torch.ops.aten.rms_norm.default(x, [16], self.weight.view(16), 1e-6) - - class _RmsNormRuntimeWeightView(torch.nn.Module): - def forward(self, x, weight): - return torch.ops.aten.rms_norm.default( - x, [weight.numel()], weight.view_as(weight), 1e-6 - ) - - class _RmsNormWithoutOptionalArgs(torch.nn.Module): - def forward(self, x): - return torch.ops.aten.rms_norm.default(x, [x.shape[-1]]) - @staticmethod @unpack_pass_fixtures def test( @@ -4467,70 +4439,6 @@ def test( assertions.assert_target_count(gm, permute, 2) assertions.assert_meta_key_on_targets(gm, permute, QCOM_INSERTED_PERMUTE) - if quantizer is None: - with subtests.test(msg="rms_norm_parameter_view"): - gm = pass_pipeline.lower_edge_ep( - module=LayoutTransform._RmsNormParameterView(16), - sample_input=(torch.randn(1, 4, 16),), - backend_type=backend_type, - compile_spec=compile_spec, - target_pass=target_pass, - quantizer=quantizer, - ).graph_module - rms_norm = next( - node - for node in gm.graph.nodes - if node.target == exir_ops.edge.aten.rms_norm.default - ) - assert rms_norm.args[1] == [16] - assert rms_norm.args[2].op == "placeholder" - - with subtests.test(msg="rms_norm_reshaped_parameter_view"): - gm = pass_pipeline.lower_edge_ep( - module=LayoutTransform._RmsNormReshapedParameterView(), - sample_input=(torch.randn(1, 4, 16),), - backend_type=backend_type, - compile_spec=compile_spec, - target_pass=target_pass, - quantizer=quantizer, - ).graph_module - rms_norm = next( - node - for node in gm.graph.nodes - if node.target == exir_ops.edge.aten.rms_norm.default - ) - assert rms_norm.args[1] == [16] - assert rms_norm.args[2].target == exir_ops.edge.aten.view_copy.default - - with subtests.test(msg="rms_norm_runtime_weight_view"): - gm = pass_pipeline.lower_edge_ep( - module=LayoutTransform._RmsNormRuntimeWeightView(), - sample_input=(torch.randn(1, 4, 16), torch.ones(16)), - backend_type=backend_type, - compile_spec=compile_spec, - target_pass=target_pass, - quantizer=quantizer, - ).graph_module - rms_norm = next( - node - for node in gm.graph.nodes - if node.target == exir_ops.edge.aten.rms_norm.default - ) - assert rms_norm.args[2].target == exir_ops.edge.aten.view_copy.default - - with subtests.test(msg="rms_norm_without_optional_args"): - gm = pass_pipeline.lower_edge_ep( - module=LayoutTransform._RmsNormWithoutOptionalArgs(), - sample_input=(torch.randn(1, 4, 16),), - backend_type=backend_type, - compile_spec=compile_spec, - target_pass=target_pass, - quantizer=quantizer, - ).graph_module - assertions.assert_target_count( - gm, exir_ops.edge.aten.rms_norm.default, 1 - ) - class LiftConstantScalarOperands: class _Pow(torch.nn.Module): @@ -4944,7 +4852,7 @@ def forward(self, x): ) return norm * self.weight - class _GemmaRmsNorm(_RmsNorm): + class _Gemma4RmsNorm(_RmsNorm): def __init__( self, dim: int, @@ -4953,6 +4861,7 @@ def __init__( negative_eps: bool = False, tensor_eps: bool = False, tensor_exponent: bool = False, + eps_first: bool = False, ): super().__init__(dim) self.square_exponent = square_exponent @@ -4960,6 +4869,7 @@ def __init__( self.negative_eps = negative_eps self.tensor_eps = tensor_eps self.tensor_exponent = tensor_exponent + self.eps_first = eps_first def forward(self, x, statistic_input=None): eps = torch.tensor(self.eps) if self.tensor_eps else self.eps @@ -4969,10 +4879,57 @@ def forward(self, x, statistic_input=None): self.reduction_dim, keepdim=True ) eps = -eps if self.negative_eps else eps - statistic = statistic + eps + statistic = eps + statistic if self.eps_first else statistic + eps norm = x * torch.pow(statistic, exponent) return norm * self.weight + class _RmsNormWithIntermediate(_RmsNorm): + def __init__(self, dim: int, intermediate: str): + super().__init__(dim) + self.intermediate = intermediate + + def forward(self, x): + norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + if self.intermediate == "add": + norm = norm + 1 + else: + norm = torch.relu(norm) + return norm * self.weight + + class _RmsNormWithWeightShape(_RmsNorm): + def __init__(self, dim: int, weight_shape: tuple[int, ...]): + super().__init__(dim) + self.weight = torch.nn.Parameter(torch.ones(weight_shape)) + + class _RmsNormWithParameterView(_RmsNorm): + def forward(self, x): + norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return norm * self.weight.view_as(self.weight) + + class _RmsNormWithReshapedParameterView(_RmsNorm): + def __init__(self, dim: int): + super().__init__(dim) + self.weight = torch.nn.Parameter(torch.ones(2, dim // 2)) + + def forward(self, x): + norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return norm * self.weight.view(-1) + + class _RmsNormWithRuntimeWeight(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.dim = dim + self.eps = 1e-6 + + def forward(self, x, weight): + norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return norm * weight + + class _RmsNormWithProductSquare(_RmsNorm): + def forward(self, x, y): + norm = x * torch.rsqrt((x * y).mean(-1, keepdim=True) + self.eps) + return norm * self.weight + class _RsqrtTensorSquareExponent(_RmsNorm): def forward(self, x): norm = x * torch.rsqrt( @@ -4987,6 +4944,53 @@ def forward(self, x): norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return self.weight * norm.to(input_dtype) + class _Gemma4Fp32RmsNorm(_Gemma4RmsNorm): + def forward(self, x): + input_dtype = x.dtype + x = x.float() + statistic = x.pow(2).mean(-1, keepdim=True) + self.eps + norm = x * torch.pow(statistic, -0.5) + return norm.to(input_dtype) * self.weight.to(input_dtype) + + class _LlamaFp32RmsNorm(_RmsNorm): + def forward(self, x): + output = x.float() + output = output * torch.rsqrt( + output.pow(2).mean(-1, keepdim=True) + self.eps + ) + return output.type_as(x) * self.weight.type_as(x) + + @staticmethod + @unpack_pass_fixtures + def test_annotation( + backend_type: QnnExecuTorchBackendType, + assertions: Assertions, + pass_pipeline: PassPipeline, + ): + dim = 16 + gm = pass_pipeline.lower_annotation_gm( + module=RecomposeRmsNorm._RmsNormWithParameterView(dim), + sample_input=(torch.randn(1, 4, 8, dim),), + target_pass=_passes.RecomposeRmsNorm, + backend_type=backend_type, + ) + assertions.assert_target_count(gm, torch.ops.aten.rms_norm.default, 1) + assertions.assert_no_target(gm, torch.ops.aten.mean.dim) + assertions.assert_no_target( + gm, + { + torch.ops.aten.view.default, + torch.ops.aten.view_as.default, + torch.ops.aten.view_copy.default, + }, + ) + rms_norm = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten.rms_norm.default + ) + assert rms_norm.args[2].op == "get_attr" + @staticmethod @unpack_pass_fixtures def test( @@ -5001,7 +5005,11 @@ def test( inputs = (torch.randn(1, 4, 8, dim),) for name, module in ( ("rsqrt", RecomposeRmsNorm._RmsNorm(dim)), - ("gemma_pow_neg_half", RecomposeRmsNorm._GemmaRmsNorm(dim)), + ("gemma4_pow_neg_half", RecomposeRmsNorm._Gemma4RmsNorm(dim)), + ( + "shape_preserving_weight_view", + RecomposeRmsNorm._RmsNormWithParameterView(dim), + ), ): with subtests.test(msg=name): gm = pass_pipeline.lower_edge_ep( @@ -5016,12 +5024,23 @@ def test( gm, exir_ops.edge.aten.rms_norm.default, 1 ) assertions.assert_no_target(gm, exir_ops.edge.aten.mean.dim) + rms_norm = next( + node + for node in gm.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ) + assert rms_norm.args[2].meta["val"].shape == (dim,) + assert math.isclose(rms_norm.args[3], module.eps, rel_tol=0.02) + if name == "shape_preserving_weight_view": + assertions.assert_no_target( + gm, exir_ops.edge.aten.view_copy.default + ) if name == "rsqrt": assertions.assert_no_target(gm, exir_ops.edge.aten.rsqrt.default) with subtests.test(msg="gemma_positive_last_reduction_dim"): gm = pass_pipeline.lower_edge_ep( - module=RecomposeRmsNorm._GemmaRmsNorm(dim, reduction_dim=3), + module=RecomposeRmsNorm._Gemma4RmsNorm(dim, reduction_dim=3), sample_input=inputs, backend_type=backend_type, compile_spec=compile_spec, @@ -5033,7 +5052,7 @@ def test( with subtests.test(msg="gemma_tensor_exponent"): gm = pass_pipeline.lower_edge_ep( - module=RecomposeRmsNorm._GemmaRmsNorm(dim, tensor_exponent=True), + module=RecomposeRmsNorm._Gemma4RmsNorm(dim, tensor_exponent=True), sample_input=inputs, backend_type=backend_type, compile_spec=compile_spec, @@ -5045,7 +5064,19 @@ def test( with subtests.test(msg="gemma_tensor_eps"): gm = pass_pipeline.lower_edge_ep( - module=RecomposeRmsNorm._GemmaRmsNorm(dim, tensor_eps=True), + module=RecomposeRmsNorm._Gemma4RmsNorm(dim, tensor_eps=True), + sample_input=inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.rms_norm.default, 1) + assertions.assert_no_target(gm, exir_ops.edge.aten.mean.dim) + + with subtests.test(msg="gemma_commuted_eps"): + gm = pass_pipeline.lower_edge_ep( + module=RecomposeRmsNorm._Gemma4RmsNorm(dim, eps_first=True), sample_input=inputs, backend_type=backend_type, compile_spec=compile_spec, @@ -5099,15 +5130,15 @@ def test( for name, module in ( ( "gemma_wrong_square_exponent", - RecomposeRmsNorm._GemmaRmsNorm(dim, square_exponent=4), + RecomposeRmsNorm._Gemma4RmsNorm(dim, square_exponent=4), ), ( "gemma_wrong_reduction_dim", - RecomposeRmsNorm._GemmaRmsNorm(dim, reduction_dim=0), + RecomposeRmsNorm._Gemma4RmsNorm(dim, reduction_dim=0), ), ( "gemma_negative_eps", - RecomposeRmsNorm._GemmaRmsNorm(dim, negative_eps=True), + RecomposeRmsNorm._Gemma4RmsNorm(dim, negative_eps=True), ), ): with subtests.test(msg=name): @@ -5133,9 +5164,63 @@ def test( 1, ) + for name, module, module_inputs in ( + ( + "intermediate_add", + RecomposeRmsNorm._RmsNormWithIntermediate(dim, "add"), + inputs, + ), + ( + "intermediate_relu", + RecomposeRmsNorm._RmsNormWithIntermediate(dim, "relu"), + inputs, + ), + ( + "wrong_weight_shape_scalar_like", + RecomposeRmsNorm._RmsNormWithWeightShape(dim, (1,)), + inputs, + ), + ( + "wrong_weight_shape_multidimensional", + RecomposeRmsNorm._RmsNormWithWeightShape(dim, (1, 1, dim)), + inputs, + ), + ( + "shape_changing_weight_view", + RecomposeRmsNorm._RmsNormWithReshapedParameterView(dim), + inputs, + ), + ( + "runtime_weight", + RecomposeRmsNorm._RmsNormWithRuntimeWeight(dim), + (inputs[0], torch.ones(dim)), + ), + ( + "product_square", + RecomposeRmsNorm._RmsNormWithProductSquare(dim), + ( + torch.rand_like(inputs[0]) + 0.5, + torch.rand_like(inputs[0]) + 0.5, + ), + ), + ): + with subtests.test(msg=name): + gm = pass_pipeline.lower_edge_ep( + module=module, + sample_input=module_inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + assertions.assert_target_count( + gm, exir_ops.edge.aten.rms_norm.default, 0 + ) + assertions.assert_target_count(gm, exir_ops.edge.aten.mean.dim, 1) + with subtests.test(msg="gemma_different_square_input"): gm = pass_pipeline.lower_edge_ep( - module=RecomposeRmsNorm._GemmaRmsNorm(dim), + module=RecomposeRmsNorm._Gemma4RmsNorm(dim), sample_input=(torch.randn(1, 4, 8, dim), torch.randn(1, 4, 8, dim)), backend_type=backend_type, compile_spec=compile_spec, @@ -5146,24 +5231,31 @@ def test( assertions.assert_target_count(gm, exir_ops.edge.aten.mean.dim, 1) if quantizer is None: - with subtests.test(msg="fp32_upcast"): - gm = pass_pipeline.lower_edge_ep( - module=RecomposeRmsNorm._Fp32RmsNorm(dim), - sample_input=(torch.randn(1, 4, 8, dim, dtype=torch.float16),), - backend_type=backend_type, - compile_spec=compile_spec, - target_pass=_passes.RecomposeRmsNorm, - quantizer=quantizer, - ).graph_module - rms_norm = next( - node - for node in gm.graph.nodes - if node.target == exir_ops.edge.aten.rms_norm.default - ) - assert rms_norm.args[0].op == "placeholder" - assert rms_norm.args[0].meta["val"].dtype == torch.float16 - assert rms_norm.args[2].meta["val"].dtype == torch.float32 - assert rms_norm.meta["val"].dtype == torch.float32 + for name, module in ( + ("qwen2_fp32_cast", RecomposeRmsNorm._Fp32RmsNorm(dim).half()), + ( + "gemma4_fp32_pow_neg_half_cast", + RecomposeRmsNorm._Gemma4Fp32RmsNorm(dim).half(), + ), + ("llama_fp32_cast", RecomposeRmsNorm._LlamaFp32RmsNorm(dim).half()), + ): + with subtests.test(msg=name): + gm = pass_pipeline.lower_edge_ep( + module=module, + sample_input=(torch.randn(1, 4, 8, dim, dtype=torch.float16),), + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RecomposeRmsNorm, + quantizer=quantizer, + ).graph_module + rms_norm = next( + node + for node in gm.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ) + assert rms_norm.args[0].meta["val"].dtype == torch.float16 + assert rms_norm.args[2].meta["val"].dtype == torch.float16 + assert rms_norm.meta["val"].dtype == torch.float16 class RemoveRedundancy: @@ -5171,6 +5263,38 @@ class _Clone(torch.nn.Module): def forward(self, x): return x.clone() + 1 + class _RmsNormParameterView(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return torch.ops.aten.rms_norm.default( + x, [self.weight.numel()], self.weight.view_as(self.weight), 1e-6 + ) + + class _RmsNormReshapedParameterView(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2, 8)) + + def forward(self, x): + return torch.ops.aten.rms_norm.default(x, [16], self.weight.view(16), 1e-6) + + class _RmsNormRuntimeWeightView(torch.nn.Module): + def forward(self, x, weight): + return torch.ops.aten.rms_norm.default( + x, [weight.numel()], weight.view_as(weight), 1e-6 + ) + + class _StaticWeightView(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def forward(self, x): + return x + self.weight.view_as(self.weight) + @staticmethod @unpack_pass_fixtures def test( @@ -5179,6 +5303,7 @@ def test( backend_type: QnnExecuTorchBackendType, assertions: Assertions, pass_pipeline: PassPipeline, + subtests, ): inputs = (torch.randn(3, 4),) gm = pass_pipeline.lower_edge_ep( @@ -5194,6 +5319,55 @@ def test( ) assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1) + if quantizer is None: + for name, module, module_inputs, expect_removed in ( + ( + "rms_norm_parameter_view", + RemoveRedundancy._RmsNormParameterView(16), + (torch.randn(1, 4, 16),), + True, + ), + ( + "rms_norm_reshaped_parameter_view", + RemoveRedundancy._RmsNormReshapedParameterView(), + (torch.randn(1, 4, 16),), + False, + ), + ( + "rms_norm_runtime_weight_view", + RemoveRedundancy._RmsNormRuntimeWeightView(), + (torch.randn(1, 4, 16), torch.ones(16)), + False, + ), + ( + "static_view_outside_rms_norm", + RemoveRedundancy._StaticWeightView(16), + (torch.randn(1, 4, 16),), + False, + ), + ): + with subtests.test(msg=name): + gm = pass_pipeline.lower_edge_ep( + module=module, + sample_input=module_inputs, + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=_passes.RemoveRedundancy, + quantizer=quantizer, + ).graph_module + if expect_removed: + rms_norm = next( + node + for node in gm.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ) + assert rms_norm.args[1] == [16] + assert rms_norm.args[2].op == "placeholder" + else: + assertions.assert_target_count( + gm, exir_ops.edge.aten.view_copy.default, 1 + ) + class ReplaceArangeArgs: class _Arange(torch.nn.Module): diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 8d6b1adcc4a..bc227b6f0e8 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -179,28 +179,50 @@ def test_rms_norm_builder_reads_keyword_weight_and_eps(self): ) from executorch.backends.qualcomm.builders.qnn_constants import OpRmsNorm - for eps, expected_eps in ( - (0.125, 0.125), - (None, torch.finfo(torch.float32).eps), + for logical_dtype, storage_dtype, eps, expected_eps in ( + (torch.float32, torch.float32, 0.125, 0.125), + (torch.float32, torch.float32, None, torch.finfo(torch.float32).eps), + (torch.float16, torch.float16, None, torch.finfo(torch.float16).eps), + (torch.float32, torch.uint8, 0.125, 0.125), + (torch.float32, torch.uint8, None, torch.finfo(torch.float32).eps), ): - with self.subTest(eps=eps): + with self.subTest( + logical_dtype=logical_dtype, storage_dtype=storage_dtype, eps=eps + ): graph = torch.fx.Graph() input_node = graph.placeholder("input") - input_node.meta["val"] = torch.empty(1, 4, 16) + input_node.meta["val"] = torch.empty(1, 4, 16, dtype=logical_dtype) + rms_input_node = input_node + if storage_dtype != logical_dtype: + quantize_node = graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + (input_node, 1.0, 0, 0, 255, storage_dtype), + ) + quantize_node.meta["val"] = torch.empty( + 1, 4, 16, dtype=storage_dtype + ) + rms_input_node = graph.call_function( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + (quantize_node, 1.0, 0, 0, 255, storage_dtype), + ) + rms_input_node.meta["val"] = torch.empty( + 1, 4, 16, dtype=logical_dtype + ) weight_node = graph.placeholder("weight") node = graph.call_function( exir_ops.edge.aten.rms_norm.default, - (input_node, [16]), + (rms_input_node, [16]), {"weight": weight_node, "eps": eps}, ) - node.meta["val"] = torch.empty(1, 4, 16) + node.meta["val"] = torch.empty(1, 4, 16, dtype=logical_dtype) edge_program = MagicMock() visitor = _node_visitor_dict["aten.rms_norm.default"]( {}, edge_program, False ) - visitor.get_node = MagicMock(side_effect=lambda value: value) - visitor.get_tensor = MagicMock(side_effect=[torch.empty(1, 4, 16)] * 2) - visitor.define_tensor = MagicMock(side_effect=[MagicMock()] * 4) + visitor.get_tensor = MagicMock( + side_effect=lambda input_node, *_: input_node.meta["val"] + ) + visitor.define_tensor = MagicMock(return_value=MagicMock()) with ( patch( @@ -223,21 +245,6 @@ def test_rms_norm_builder_reads_keyword_weight_and_eps(self): ) self.assertEqual(expected_eps, epsilon_call.args[2][QCOM_DATA]) - def test_rms_norm_recomposition_rejects_missing_upcast_metadata(self): - from executorch.backends.qualcomm._passes.recompose_rms_norm import ( - RecomposeRmsNorm, - ) - - graph = torch.fx.Graph() - source = graph.placeholder("source") - upcast = graph.call_function( - exir_ops.edge.dim_order_ops._to_dim_order_copy.default, - (source,), - {"dtype": torch.float32}, - ) - - self.assertFalse(RecomposeRmsNorm._is_fp32_upcast(upcast)) - def test_rms_norm_recomposition_rejects_scaled_eps(self): from executorch.backends.qualcomm._passes.recompose_rms_norm import ( RecomposeRmsNorm, diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 221b7b4dbac..42b14c6e85c 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -2348,6 +2348,15 @@ def test_qnn_backend_rms_norm(self): with self.subTest(i=i): self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_fp32_rms_norm(self): + if get_backend_type(self.backend) != QnnExecuTorchBackendType.kHtpBackend: + self.skipTest("validates HTP FP32 RMSNorm accumulation") + module = Fp32RmsNorm(16).half() # noqa: F405 + sample_input = ( + torch.linspace(-4096, 4096, 16, dtype=torch.float16).reshape(1, 1, 1, 16), + ) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_roll(self): modules = [ Roll(shifts=(3, 3), dims=(1, 2)), # noqa: F405