From 9ca3886834bddb3e9dd15987b7e9a9a4de0167b8 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 11:45:12 -0800 Subject: [PATCH 1/3] Qualcomm AI Engine Direct - Keep the dtype of a recomposed RMSNorm RecomposeRmsNorm finds the norm's input by skipping casts, but it only skipped aten.to.dtype. Edge graphs carry the cast as _to_copy or _to_dim_order_copy, so for a Qwen-style norm (upcast to float32, normalize, cast back) the rebuilt rms_norm read the float32 cast and the cast back was lost. Its metadata was copied from the original output, so the graph still claimed the old dtype until a later pass recomputed it. A float16 Qwen norm then failed to lower with "rms_norm was not decomposed or delegated". A bfloat16 Qwen norm followed by masked_scatter failed with "expected self and source to have same dtypes", and when the rebuilt norm stayed on the CPU, to_executorch failed with "Missing out variants: aten::rms_norm". The pass assumes QNN's RmsNorm takes one dtype throughout, so it now rebuilds the norm only when its input and its output share a dtype, and no cast in between rounds below that dtype: - It skips casts that change only the dtype: to.dtype, to.device or type_as in captured graphs, _to_copy or _to_dim_order_copy in edge graphs, whose output keeps the input's layout, device and dim order. - It walks back to the nearest tensor in the output dtype, but not past a cast to a narrower dtype or one that changes the dim order. - Only such casts may sit between the normalization and the weight. The "*" in the patterns can match any op, so before, a skipped norm let a later multiply be rebuilt as the norm with the wrong weight. - Outside quantization capture, the dtype also needs a QNN type. - The tensor it scales has to be the tensor it squares, past such casts (T5's norm squares a float32 copy), and the square has to be pow(2) or h * h, past such casts. Before, any multiply by a root mean square was rebuilt as a norm of the squared tensor, and the scaled tensor was dropped. So a float16 Qwen norm is rebuilt on its float16 input. A norm that rounds to float16 before a float32 weight, and a bfloat16 norm, stay decomposed, unless the bfloat16 norm is about to be quantized. A new test covers these cases, with and without dim order, with a channels-last input and a cast that changes its dim order, a Llama-style norm, a T5 norm, float16 rounding before and inside a float32 norm, a multiply after the norm, a square of two casts, and quantization capture, including type_as, to with a device, and a weight narrower than the input there, and a wider one, which stays decomposed since rms_norm returns its input's dtype. Edge graphs first go through the QNN export pipeline, as in to_edge_transform_and_lower_to_qnn, so the exponent of pow is a lifted constant. The test checks which tensor the norm reads and compares the output with eager, with random weights. Another test checks that a different scaled tensor, a cube and a product of two tensors stay decomposed, with and without the lifted exponent. --- .../qualcomm/_passes/recompose_rms_norm.py | 87 +++++++++- backends/qualcomm/tests/test_passes.py | 160 ++++++++++++++++++ 2 files changed, 241 insertions(+), 6 deletions(-) diff --git a/backends/qualcomm/_passes/recompose_rms_norm.py b/backends/qualcomm/_passes/recompose_rms_norm.py index f9cdd676345..f3e9fe83c48 100644 --- a/backends/qualcomm/_passes/recompose_rms_norm.py +++ b/backends/qualcomm/_passes/recompose_rms_norm.py @@ -5,7 +5,10 @@ # LICENSE file in the root directory of this source tree. import torch from executorch.backends.qualcomm._passes.utils import find_patterns -from executorch.backends.qualcomm.builders.node_visitor import dq_ops +from executorch.backends.qualcomm.builders.node_visitor import ( + dq_ops, + QNN_TENSOR_TYPE_MAP, +) 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 @@ -79,19 +82,89 @@ def __init__(self, quantization_capture=False): self.rms_norm_target = exir_ops.edge.aten.rms_norm.default self.skip_targets = [ exir_ops.edge.aten.to.dtype, + exir_ops.edge.aten._to_copy.default, + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, ] self.quantization_capture = quantization_capture if quantization_capture: self.rms_norm_target = torch.ops.aten.rms_norm.default self.skip_targets = [ torch.ops.aten.to.dtype, + torch.ops.aten.to.device, + torch.ops.aten.type_as.default, ] - def _get_input_node(self, node): - input_node = node.args[0] - while input_node.target in self.skip_targets: + def _is_cast(self, node, dtype): + """ + Whether the node is a cast that changes nothing but the dtype, to a + dtype at least as wide as `dtype`. A narrower cast rounds the values, + which the rebuilt rms_norm in `dtype` would not do. + """ + if not _is_call(node) or node.target not in self.skip_targets: + return False + source, cast = node.args[0].meta["val"], node.meta["val"] + return torch.promote_types(cast.dtype, dtype) == cast.dtype and ( + cast.layout, + cast.device, + cast.dim_order(), + ) == (source.layout, source.device, source.dim_order()) + + def _strip_casts(self, node, dtype): + while self._is_cast(node, dtype): + node = node.args[0] + return node + + def _squares_the_scaled_tensor(self, pattern, dtype): + square, scale_mul = pattern[6], pattern[2] + base = square.args[0] + if _is_pow(square): + exponent = square.args[1] + if _is_node(exponent): + # LiftConstantScalarOperands may have turned it into a constant. + constant = exponent.meta["val"].constant + exponent = None if constant is None else constant.item() + squares = exponent == 2 + else: + squares = self._strip_casts(square.args[1], dtype) is self._strip_casts( + base, dtype + ) + scaled = [arg for arg in scale_mul.args if arg is not pattern[3]] + # T5's norm squares a float32 copy of the tensor it scales. + return ( + squares + and len(scaled) == 1 + and self._strip_casts(scaled[0], dtype) is self._strip_casts(base, dtype) + ) + + def _get_norm_input(self, pattern, gamma_node): + """ + The tensor to rebuild the norm on, or None if it has to stay + decomposed. rms_norm returns its input's dtype, so the input has to be + in the pattern's output dtype. Unless the norm is about to be + quantized, QNN also needs a type for that dtype. + """ + last_mul_node = pattern[0] + dtype = last_mul_node.meta["val"].dtype + if not self._squares_the_scaled_tensor(pattern, dtype): + return None + input_node = pattern[6].args[0] + while self._is_cast(input_node, dtype) and ( + input_node.meta["val"].dtype != dtype + or input_node.args[0].meta["val"].dtype == dtype + ): input_node = input_node.args[0] - return input_node + if input_node.meta["val"].dtype != dtype or ( + not self.quantization_capture and dtype not in QNN_TENSOR_TYPE_MAP + ): + return None + # The "*" in the patterns can match any op, but only casts may sit + # between the normalization and the weight. + normalized = next( + (arg for arg in last_mul_node.args if arg is not gamma_node), None + ) + return ( + input_node if self._strip_casts(normalized, dtype) is pattern[2] else None + ) def call(self, graph_module: torch.fx.GraphModule): graph = graph_module.graph @@ -127,7 +200,9 @@ def call(self, graph_module: torch.fx.GraphModule): if gamma_node is None: continue - input_node = self._get_input_node(rms_norm_pattern[6]) + input_node = self._get_norm_input(rms_norm_pattern, gamma_node) + if input_node is None: + continue # Only the last-dim form the RmsNorm builder supports is recomposed. # Anything else, e.g. a 0-d scale, is left decomposed. input_shape = list(input_node.meta["val"].shape) diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index a607bcb4d81..af7e8bb965a 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -747,6 +747,166 @@ def forward(self, x): RecomposeRmsNorm()(ep.graph_module) torch.testing.assert_close(ep.module()(x), model(x)) + def test_recompose_rms_norm_keeps_the_norm_dtype(self): + class UpcastNorm(torch.nn.Module): # transformers' Qwen2RMSNorm + def __init__(self, dtype, compute_dtype=torch.float32): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(32, dtype=dtype)) + self.compute_dtype = compute_dtype + + def upcast(self, x): + return x.to(self.compute_dtype) + + def forward(self, x): + h = self.upcast(x) + h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + 1e-6) + return self.weight * h.to(x.dtype) + + class TypeAsNorm(torch.nn.Module): # type_as, to.device, a square of casts + def __init__(self, dtype): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(32, dtype=dtype)) + + def forward(self, x): + square = x.float() * x.float() + rms = torch.rsqrt(square.mean(-1, keepdim=True) + 1e-6) + h = x.to(x.device, torch.float32) * rms + return h.type_as(x) * self.weight + + class T5Norm(UpcastNorm): # transformers' T5LayerNorm, half weights + def forward(self, x): + variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + 1e-6) + return self.weight * x.to(self.weight.dtype) + + class ContiguousUpcastNorm(UpcastNorm): + def upcast(self, x): + return x.to(self.compute_dtype, memory_format=torch.contiguous_format) + + class RoundTripNorm(UpcastNorm): + def forward(self, x): + return super().forward(x.half().float()) + + class ScaledAfterNorm(UpcastNorm): + def __init__(self, dtype): + super().__init__(dtype) + self.register_buffer("scale", torch.randn(32)) + + def forward(self, x): + return super().forward(x) * self.scale + + # The norm is rebuilt on the nearest input in the output dtype, past + # casts that keep it or widen it. It stays decomposed when a cast rounds + # below the output dtype or changes the dim order, when another op sits + # between the normalization and the weight, or when QNN has no type for + # it (bfloat16) and it is not about to be quantized. A weight narrower + # than the input is fine when the norm is about to be quantized, since + # the quantizer gives it the input's quantized type, but a wider one is + # not, since rms_norm returns the input's dtype. Edge graphs carry casts + # as _to_dim_order_copy, or as _to_copy without dim order, and captured + # graphs as to.dtype or, with a device, to.device. + f16, bf16, f32 = torch.float16, torch.bfloat16, torch.float32 + c, cl = torch.contiguous_format, torch.channels_last + for model, x_dtype, layout, capture, dim_order, norm_inputs in ( + (UpcastNorm(f16), f16, c, False, True, [(f16, True)]), + (UpcastNorm(f16), f16, cl, False, True, [(f16, True)]), + (ContiguousUpcastNorm(f16), f16, cl, False, True, []), + (UpcastNorm(f16), f16, c, False, False, [(f16, True)]), + (TypeAsNorm(f16), f16, c, False, True, [(f16, True)]), + (T5Norm(f16), f16, c, False, True, [(f16, True)]), + (UpcastNorm(f32), f16, c, False, True, []), + (UpcastNorm(f32, compute_dtype=f16), f32, c, False, True, []), + (RoundTripNorm(f32), f32, c, False, True, [(f32, False)]), + (UpcastNorm(bf16), bf16, c, False, True, []), + (ScaledAfterNorm(bf16), bf16, c, False, True, []), + (UpcastNorm(bf16), bf16, c, True, True, [(bf16, True)]), + (UpcastNorm(f32), f32, c, True, True, [(f32, True)]), + (UpcastNorm(bf16), f32, c, True, True, [(f32, True)]), + (UpcastNorm(f32), bf16, c, True, True, []), + (TypeAsNorm(bf16), bf16, c, True, True, [(bf16, True)]), + (TypeAsNorm(f32), f32, c, True, True, [(f32, True)]), + (T5Norm(bf16), bf16, c, True, True, [(bf16, True)]), + ): + with self.subTest( + model=type(model).__name__, + weight=model.weight.dtype, + x=x_dtype, + layout=layout, + capture=capture, + dim_order=dim_order, + ): + x = torch.randn(2, 4, 8, 32, dtype=x_dtype) + self._check_recomposed_rms_norm( + model, x.to(memory_format=layout), capture, dim_order, norm_inputs + ) + + def _check_recomposed_rms_norm(self, model, x, capture, dim_order, norm_inputs): + """ + Runs RecomposeRmsNorm on the captured or edge graph of `model`, checks + the dtype of each rebuilt norm's input and whether it is the graph + input, and compares the output with eager. + """ + ep = torch.export.export(model, (x,), strict=True) + if not capture: + # As in to_edge_transform_and_lower_to_qnn, which lifts the exponent of pow. + ep = get_qnn_pass_manager_cls()().transform_for_export_pipeline(ep) + config = EdgeCompileConfig(_skip_dim_order=not dim_order) + ep = to_edge(ep, compile_config=config).exported_program() + graph_module = ep.module() if capture else ep.graph_module + RecomposeRmsNorm(quantization_capture=capture)(graph_module) + rms_norms = ( + exir_ops.edge.aten.rms_norm.default, + torch.ops.aten.rms_norm.default, + ) + self.assertEqual( + [ + (node.args[0].meta["val"].dtype, node.args[0].op == "placeholder") + for node in graph_module.graph.nodes + if node.target in rms_norms + ], + norm_inputs, + ) + module = graph_module if capture else ep.module() + torch.testing.assert_close(module(x), model(x)) + + def test_recompose_rms_norm_needs_the_norm_of_one_tensor(self): + class AlmostNorm(torch.nn.Module): + def __init__(self, scaled, squared): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(32)) + self.register_buffer("other", torch.rand(32, 32) + 0.5) + self.scaled, self.squared = scaled, squared + + def forward(self, x): + other = self.other + squared = {"pow3": x.pow(3), "mixed": x * other}.get( + self.squared, x.pow(2) + ) + scaled = other if self.scaled == "other" else x + rms = torch.rsqrt(squared.mean(-1, keepdim=True) + 1e-6) + return self.weight * (scaled * rms) + + # Each model looks like a norm, but the tensor it scales is not the + # one it squares, or it does not square it, so it stays decomposed. + # The export pipeline lifts the exponent of pow into a constant. + for lift in (False, True): + for scaled, squared in (("other", "pow2"), ("x", "pow3"), ("x", "mixed")): + with self.subTest(scaled=scaled, squared=squared, lift=lift): + model = AlmostNorm(scaled, squared) + # Positive, so a cube or a product still has a real square root. + x = torch.rand(32, 32) + 0.5 + ep = torch.export.export(model, (x,), strict=True) + if lift: + pm = get_qnn_pass_manager_cls()() + ep = pm.transform_for_export_pipeline(ep) + ep = to_edge(ep).exported_program() + RecomposeRmsNorm()(ep.graph_module) + self.assertNotIn( + exir_ops.edge.aten.rms_norm.default, + [node.target for node in ep.graph.nodes], + ) + torch.testing.assert_close(ep.module()(x), model(x)) + def test_i64_to_i32_casts_a_top_of_graph_constant_after_placeholders(self): # DecomposeRemainder puts the int64 constant 3 at the top of the graph. class IntRemainder(torch.nn.Module): From 8096fcf3f905e903484dea180748ab65b598b858 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 11:45:12 -0800 Subject: [PATCH 2/3] Qualcomm AI Engine Direct - Keep a _to_copy that changes the dtype Without dim order, a layout change in an edge graph is a _to_copy with a memory_format argument, and RemoveRedundancy drops every such _to_copy, because QNN does not need layout changes. A cast like x.to(torch.float16, memory_format=torch.contiguous_format) is also a _to_copy with memory_format, so its dtype change was dropped too, and later ops read the old dtype while their metadata said the new one. RemoveRedundancy now drops a _to_copy with memory_format only when it keeps the dtype, as it already does for _to_dim_order_copy. The new test exports such casts without dim order. A layout-only cast is still removed, and a cast to float16 stays, with a float16 output. --- .../qualcomm/_passes/remove_redundancy.py | 2 +- backends/qualcomm/tests/test_passes.py | 28 +++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/backends/qualcomm/_passes/remove_redundancy.py b/backends/qualcomm/_passes/remove_redundancy.py index 7589d40b745..8174968e23f 100644 --- a/backends/qualcomm/_passes/remove_redundancy.py +++ b/backends/qualcomm/_passes/remove_redundancy.py @@ -49,7 +49,7 @@ def _same_shape_condition(self, node): return node.args[0].meta["val"].shape == node.meta["val"].shape def _to_copy_op_condition(self, node): - return "memory_format" in node.kwargs + return "memory_format" in node.kwargs and self._dim_order_op_condition(node) def _default_condition(self, ndoe): return True diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index af7e8bb965a..2253508bff8 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -907,6 +907,34 @@ def forward(self, x): ) torch.testing.assert_close(ep.module()(x), model(x)) + def test_remove_redundancy_keeps_a_to_copy_that_changes_the_dtype(self): + class Copy(torch.nn.Module): + def __init__(self, dtype): + super().__init__() + self.dtype = dtype + + def forward(self, x): + y = x.to(self.dtype, memory_format=torch.contiguous_format) + return y + y + + # Without dim order, a layout change is a _to_copy with memory_format, + # which QNN drops. A _to_copy that also changes the dtype has to stay. + x = torch.randn(2, 4, 8, 8).to(memory_format=torch.channels_last) + for dtype, kept in ((torch.float32, False), (torch.float16, True)): + with self.subTest(dtype=dtype): + config = EdgeCompileConfig(_skip_dim_order=True) + ep = to_edge( + torch.export.export(Copy(dtype), (x,), strict=True), + compile_config=config, + ).exported_program() + RemoveRedundancy(quantization_capture=False)(ep.graph_module) + self.assertEqual( + exir_ops.edge.aten._to_copy.default + in [node.target for node in ep.graph.nodes], + kept, + ) + self.assertEqual(ep.module()(x).dtype, dtype) + def test_i64_to_i32_casts_a_top_of_graph_constant_after_placeholders(self): # DecomposeRemainder puts the int64 constant 3 at the top of the graph. class IntRemainder(torch.nn.Module): From 1fc8832fd41a6eb19d9da85689b8e1d282e4a092 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 06:30:02 -0800 Subject: [PATCH 3/3] Qualcomm AI Engine Direct - Rebuild the float32 part of an RMSNorm RecomposeRmsNorm rebuilds a norm into one rms_norm only when its input, its weight and its output share a dtype. A Llama-style norm with a bfloat16 input or weight upcasts to float32, normalizes, casts back and multiplies by the weight, so it stayed decomposed. With HTP's float16 precision, QNN then ran its float32 pow, mean and rsqrt in float16. Squaring x passes the largest float16 value (65504) once |x| > 255.9, so rows with large values, like the massive activations of a language model, came out wrong (relative error up to 4.5 on SM8650 for values of 1000 to 2000). When the whole norm cannot be rebuilt and the graph is lowered for HTP's float16 precision, the pass now rebuilds only its float32 normalization, as an rms_norm without a weight, and leaves the casts and the weight multiply after it as they are. rms_norm computes in float32, so this matches what the model computes. A float32 norm whose weight is not one value per channel, such as a 0-d scale, gets the same. A normalization in float16 stays decomposed, since rms_norm would compute it in float32 instead. Other backends and precisions, and quantization capture, keep the normalization decomposed as before, since every op in it has a CPU kernel. The pass takes the compiler specs to know the backend and precision, so this needs a caller that passes them, as to_edge_transform_and_lower_to_qnn does. rms_norm has no CPU kernel, so a rebuilt norm that QNN refuses fails the lowering. The pass now rebuilds only what QNN's RmsNorm takes: the last of 1 to 4 dims, with static sizes and a constant eps, both for a whole norm and for its normalization, and also during quantization capture. Anything else stays decomposed. This also lets a rank-5 or a dynamic-shape float32 norm lower, which failed before. Outside quantization capture, the norm's input also has to read something at runtime. The pass cannot see the partitions, and HTP drops a partition that reads only constants, so a norm of constants cut off from runtime ops failed to lower. Such a norm now stays decomposed, also when it would have shared a partition with runtime ops. That still lowers, but not as one node, and with HTP's float16 precision its square can overflow for constants above 255.9. The pass takes the edge program to tell constants from runtime inputs, as the partitioner does. For the same reason a graph input is no longer taken as the weight: QNN's RmsNorm needs a constant one, and lowering failed when it tried to read the input as a parameter. On the QNN SDK's x86 simulator for SM8650 (float16), a bfloat16 norm and a float32 norm with a bfloat16 weight now lower to RmsNorm, with a worst relative error of 0.005 and 0.0007 on rows with values up to 2000, like a float32 norm (0.0008). The test now also checks whether each rebuilt norm keeps its weight. It expects the float32 normalization alone to be rebuilt for a bfloat16 norm, a float32 input with a bfloat16 weight, a float16 cast before a float32 weight, a cast that changes the dim order, and an op between the normalization and the weight, and no normalization-only rebuild in float16 or during quantization capture. A new test checks that a rank-5 input, a dynamic dim and an eps buffer, with and without a weight, and a mean that is not over the last dim with keepdim stay decomposed, that a norm with a 0-d scale is rebuilt without it, and that a norm of a buffer is not, with or without a weight, unless the model updates the buffer. It replaces the old last-dim test, which only compared outputs. Two more tests check that only HTP's float16 precision, passed in by the pass manager, rebuilds a normalization alone, and that a graph input is not taken as the weight. --- .../qualcomm/_passes/recompose_rms_norm.py | 130 ++++++-- backends/qualcomm/tests/test_passes.py | 278 ++++++++++++++---- 2 files changed, 322 insertions(+), 86 deletions(-) diff --git a/backends/qualcomm/_passes/recompose_rms_norm.py b/backends/qualcomm/_passes/recompose_rms_norm.py index f3e9fe83c48..5ff2fd8c6d8 100644 --- a/backends/qualcomm/_passes/recompose_rms_norm.py +++ b/backends/qualcomm/_passes/recompose_rms_norm.py @@ -9,6 +9,18 @@ dq_ops, QNN_TENSOR_TYPE_MAP, ) +from executorch.backends.qualcomm.builders.utils import ( + is_mutable_buffer_input, + is_parameter, +) +from executorch.backends.qualcomm.partition.utils import generate_qnn_executorch_option +from executorch.backends.qualcomm.serialization.qc_schema import ( + QnnExecuTorchBackendType, + QnnExecuTorchHtpPrecision, +) +from executorch.backends.qualcomm.serialization.qc_schema_serialize import ( + flatbuffer_to_option, +) 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 @@ -77,8 +89,21 @@ 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, compiler_specs=None + ): super(RecomposeRmsNorm, self).__init__() + self.edge_program = edge_program + self.htp_fp16 = False + if compiler_specs is not None: + options = flatbuffer_to_option( + generate_qnn_executorch_option(compiler_specs) + ).backend_options + self.htp_fp16 = ( + options.backend_type == QnnExecuTorchBackendType.kHtpBackend + and options.htp_options is not None + and options.htp_options.precision == QnnExecuTorchHtpPrecision.kHtpFp16 + ) self.rms_norm_target = exir_ops.edge.aten.rms_norm.default self.skip_targets = [ exir_ops.edge.aten.to.dtype, @@ -136,15 +161,15 @@ def _squares_the_scaled_tensor(self, pattern, dtype): and self._strip_casts(scaled[0], dtype) is self._strip_casts(base, dtype) ) - def _get_norm_input(self, pattern, gamma_node): + def _get_input(self, pattern, dtype): """ - The tensor to rebuild the norm on, or None if it has to stay - decomposed. rms_norm returns its input's dtype, so the input has to be - in the pattern's output dtype. Unless the norm is about to be - quantized, QNN also needs a type for that dtype. + The tensor to rebuild a norm in `dtype` on, or None if there is none. + rms_norm returns its input's dtype, so the input has to be in `dtype`. + Unless the norm is about to be quantized, QNN also needs a type for it, + and the input has to read something at runtime: the pass cannot see + the partitions, HTP drops one that reads only constants, and rms_norm + has no CPU kernel to fall back on. """ - last_mul_node = pattern[0] - dtype = last_mul_node.meta["val"].dtype if not self._squares_the_scaled_tensor(pattern, dtype): return None input_node = pattern[6].args[0] @@ -154,17 +179,51 @@ def _get_norm_input(self, pattern, gamma_node): ): input_node = input_node.args[0] if input_node.meta["val"].dtype != dtype or ( - not self.quantization_capture and dtype not in QNN_TENSOR_TYPE_MAP + not self.quantization_capture + and ( + dtype not in QNN_TENSOR_TYPE_MAP + or not self._reads_runtime_input(input_node) + ) ): return None + return input_node + + def _reads_runtime_input(self, node): + # Placeholders count as QNN's partitioner counts them. Without the + # program, every placeholder counts. + nodes, seen = [node], set() + while nodes: + node = nodes.pop() + if node in seen: + continue + seen.add(node) + if node.op == "placeholder" and ( + self.edge_program is None + or not is_parameter(node, self.edge_program) + or is_mutable_buffer_input(node, self.edge_program) + ): + return True + nodes.extend(node.all_input_nodes) + return False + + def _get_norm_input(self, pattern, gamma_node): + last_mul_node = pattern[0] + dtype = last_mul_node.meta["val"].dtype # The "*" in the patterns can match any op, but only casts may sit # between the normalization and the weight. normalized = next( (arg for arg in last_mul_node.args if arg is not gamma_node), None ) - return ( - input_node if self._strip_casts(normalized, dtype) is pattern[2] else None - ) + if self._strip_casts(normalized, dtype) is not pattern[2]: + return None + input_node = self._get_input(pattern, dtype) + if ( + input_node is None + or list(gamma_node.meta["val"].shape) + != list(input_node.meta["val"].shape)[-1:] + ): + return None + return input_node def call(self, graph_module: torch.fx.GraphModule): graph = graph_module.graph @@ -194,35 +253,54 @@ def call(self, graph_module: torch.fx.GraphModule): 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): + if ( + _is_get_attr(arg) + or (_is_placeholder(arg) and not self._reads_runtime_input(arg)) + or _is_dq(arg) + ): gamma_node = arg - if gamma_node is None: - continue - - input_node = self._get_norm_input(rms_norm_pattern, gamma_node) + input_node = ( + self._get_norm_input(rms_norm_pattern, gamma_node) + if gamma_node is not None + else None + ) + dtype = rms_norm_pattern[2].meta["val"].dtype + if input_node is None and self.htp_fp16 and dtype == torch.float32: + # Rebuild only the normalization and keep the casts and the + # weight after it. rms_norm computes in float32, so this is + # exact. Left decomposed, HTP's float16 precision would + # square in float16, which overflows above 65504. + output_node, gamma_node = rms_norm_pattern[2], None + input_node = self._get_input(rms_norm_pattern, dtype) + else: + output_node = last_mul_node if input_node is None: continue - # Only the last-dim form the RmsNorm builder supports is recomposed. - # Anything else, e.g. a 0-d scale, is left decomposed. + # Only the form QNN's RmsNorm takes is recomposed: the last of + # 1 to 4 dims, all with static sizes. Anything else is left + # decomposed, since rms_norm has no CPU kernel to fall back on. input_shape = list(input_node.meta["val"].shape) rank = len(input_shape) mean_node = rms_norm_pattern[5] mean_dims = mean_node.args[1] or range(rank) keepdim = len(mean_node.args) > 2 and mean_node.args[2] if ( - rank == 0 + not 0 < rank <= 4 + or not all(isinstance(dim, int) for dim in input_shape) or [dim % rank for dim in mean_dims] != [rank - 1] or not (keepdim or rank == 1) - or list(gamma_node.meta["val"].shape) != input_shape[-1:] ): continue eps = rms_norm_pattern[4].args[1] if isinstance(eps, torch.fx.Node): - eps = eps.meta["val"].constant.item() + eps = eps.meta["val"].constant + if eps is None: + continue + eps = eps.item() - with graph.inserting_before(last_mul_node): + with graph.inserting_before(output_node): # args schema # (Tensor input, int[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor rms_node = graph.create_node( @@ -235,11 +313,11 @@ def call(self, graph_module: torch.fx.GraphModule): eps, ), ) - users = last_mul_node.users.copy() + users = output_node.users.copy() for user in users: - user.replace_input_with(last_mul_node, rms_node) + user.replace_input_with(output_node, rms_node) # copy metadata - rms_node.meta = last_mul_node.meta + rms_node.meta = output_node.meta dead_code_elimination_pass(graph_module) return PassResult(graph_module, True) diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 2253508bff8..92cb4511e36 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -720,33 +720,6 @@ def forward(self, x): ) self.assertTrue(torch.equal(ep.module()(x), model(x))) - def test_recompose_rms_norm_only_takes_last_dim_with_weight(self): - class ScaledNorm(torch.nn.Module): - def __init__(self, dim, scale, keepdim=True): - super().__init__() - self.dim = dim - self.keepdim = keepdim - self.register_buffer("scale", scale) - - def forward(self, x): - ms = x.pow(2).mean(self.dim, keepdim=self.keepdim) - return x * torch.rsqrt(ms + 1e-6) * self.scale - - x = torch.randn(32, 32) - # A 0-d scale is not a weight, and a mean over dim 0, over all dims, or - # without keepdim is not a last-dim RMSNorm. - for model in ( - ScaledNorm(-1, torch.tensor(2.0)), - ScaledNorm(0, torch.randn(32)), - ScaledNorm(None, torch.randn(32)), - ScaledNorm(-1, torch.randn(32), keepdim=False), - ): - with self.subTest(dim=model.dim, keepdim=model.keepdim): - ep = to_edge(torch.export.export(model, (x,), strict=True)) - ep = ep.exported_program() - RecomposeRmsNorm()(ep.graph_module) - torch.testing.assert_close(ep.module()(x), model(x)) - def test_recompose_rms_norm_keeps_the_norm_dtype(self): class UpcastNorm(torch.nn.Module): # transformers' Qwen2RMSNorm def __init__(self, dtype, compute_dtype=torch.float32): @@ -787,45 +760,50 @@ class RoundTripNorm(UpcastNorm): def forward(self, x): return super().forward(x.half().float()) - class ScaledAfterNorm(UpcastNorm): + class ShiftedNorm(UpcastNorm): def __init__(self, dtype): super().__init__(dtype) - self.register_buffer("scale", torch.randn(32)) + self.register_buffer("shift", torch.randn(32, dtype=dtype)) def forward(self, x): - return super().forward(x) * self.scale + h = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) + return self.weight * (h + self.shift) # The norm is rebuilt on the nearest input in the output dtype, past - # casts that keep it or widen it. It stays decomposed when a cast rounds - # below the output dtype or changes the dim order, when another op sits - # between the normalization and the weight, or when QNN has no type for - # it (bfloat16) and it is not about to be quantized. A weight narrower - # than the input is fine when the norm is about to be quantized, since - # the quantizer gives it the input's quantized type, but a wider one is - # not, since rms_norm returns the input's dtype. Edge graphs carry casts - # as _to_dim_order_copy, or as _to_copy without dim order, and captured - # graphs as to.dtype or, with a device, to.device. + # casts that keep it or widen it. Only its float32 normalization is + # rebuilt when a cast rounds below the output dtype or changes the dim + # order, when the weight is cast or another op sits between the + # normalization and the weight, or when QNN has no type for the output + # dtype (bfloat16). When the norm is about to be quantized, it is + # rebuilt whole or not at all. A weight narrower than the input is fine + # there, since the quantizer gives it the input's quantized type, but a + # wider one is not, since rms_norm returns the input's dtype. Edge + # graphs carry casts as _to_dim_order_copy, or as _to_copy without dim + # order, and captured graphs as to.dtype or, with a device, to.device. f16, bf16, f32 = torch.float16, torch.bfloat16, torch.float32 c, cl = torch.contiguous_format, torch.channels_last for model, x_dtype, layout, capture, dim_order, norm_inputs in ( - (UpcastNorm(f16), f16, c, False, True, [(f16, True)]), - (UpcastNorm(f16), f16, cl, False, True, [(f16, True)]), - (ContiguousUpcastNorm(f16), f16, cl, False, True, []), - (UpcastNorm(f16), f16, c, False, False, [(f16, True)]), - (TypeAsNorm(f16), f16, c, False, True, [(f16, True)]), - (T5Norm(f16), f16, c, False, True, [(f16, True)]), - (UpcastNorm(f32), f16, c, False, True, []), + (UpcastNorm(f16), f16, c, False, True, [(f16, True, True)]), + (UpcastNorm(f16), f16, cl, False, True, [(f16, True, True)]), + (ContiguousUpcastNorm(f16), f16, cl, False, True, [(f32, False, False)]), + (UpcastNorm(f16), f16, c, False, False, [(f16, True, True)]), + (TypeAsNorm(f16), f16, c, False, True, [(f16, True, True)]), + (T5Norm(f16), f16, c, False, True, [(f16, True, True)]), + (UpcastNorm(f32), f16, c, False, True, [(f32, False, False)]), (UpcastNorm(f32, compute_dtype=f16), f32, c, False, True, []), - (RoundTripNorm(f32), f32, c, False, True, [(f32, False)]), - (UpcastNorm(bf16), bf16, c, False, True, []), - (ScaledAfterNorm(bf16), bf16, c, False, True, []), - (UpcastNorm(bf16), bf16, c, True, True, [(bf16, True)]), - (UpcastNorm(f32), f32, c, True, True, [(f32, True)]), - (UpcastNorm(bf16), f32, c, True, True, [(f32, True)]), + (RoundTripNorm(f32), f32, c, False, True, [(f32, False, True)]), + (UpcastNorm(bf16), bf16, c, False, True, [(f32, False, False)]), + (UpcastNorm(bf16), f32, c, False, True, [(f32, True, False)]), + (ShiftedNorm(f32), f32, c, False, True, [(f32, True, False)]), + (UpcastNorm(bf16), bf16, c, True, True, [(bf16, True, True)]), + (UpcastNorm(f32), f32, c, True, True, [(f32, True, True)]), + (UpcastNorm(bf16), f32, c, True, True, [(f32, True, True)]), (UpcastNorm(f32), bf16, c, True, True, []), - (TypeAsNorm(bf16), bf16, c, True, True, [(bf16, True)]), - (TypeAsNorm(f32), f32, c, True, True, [(f32, True)]), - (T5Norm(bf16), bf16, c, True, True, [(bf16, True)]), + (ShiftedNorm(f32), f32, c, True, True, []), + (ContiguousUpcastNorm(f16), f16, cl, True, True, []), + (TypeAsNorm(bf16), bf16, c, True, True, [(bf16, True, True)]), + (TypeAsNorm(f32), f32, c, True, True, [(f32, True, True)]), + (T5Norm(bf16), bf16, c, True, True, [(bf16, True, True)]), ): with self.subTest( model=type(model).__name__, @@ -840,11 +818,19 @@ def forward(self, x): model, x.to(memory_format=layout), capture, dim_order, norm_inputs ) + @staticmethod + def _htp_compiler_specs(use_fp16): + return generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=use_fp16), + ) + def _check_recomposed_rms_norm(self, model, x, capture, dim_order, norm_inputs): """ Runs RecomposeRmsNorm on the captured or edge graph of `model`, checks - the dtype of each rebuilt norm's input and whether it is the graph - input, and compares the output with eager. + the dtype of each rebuilt norm's input, whether it is the graph input + and whether the norm keeps the weight, and compares the output with + eager. """ ep = torch.export.export(model, (x,), strict=True) if not capture: @@ -853,14 +839,22 @@ def _check_recomposed_rms_norm(self, model, x, capture, dim_order, norm_inputs): config = EdgeCompileConfig(_skip_dim_order=not dim_order) ep = to_edge(ep, compile_config=config).exported_program() graph_module = ep.module() if capture else ep.graph_module - RecomposeRmsNorm(quantization_capture=capture)(graph_module) + RecomposeRmsNorm( + quantization_capture=capture, + edge_program=None if capture else ep, + compiler_specs=None if capture else self._htp_compiler_specs(True), + )(graph_module) rms_norms = ( exir_ops.edge.aten.rms_norm.default, torch.ops.aten.rms_norm.default, ) self.assertEqual( [ - (node.args[0].meta["val"].dtype, node.args[0].op == "placeholder") + ( + node.args[0].meta["val"].dtype, + node.args[0].op == "placeholder", + node.args[2] is not None, + ) for node in graph_module.graph.nodes if node.target in rms_norms ], @@ -869,6 +863,170 @@ def _check_recomposed_rms_norm(self, model, x, capture, dim_order, norm_inputs): module = graph_module if capture else ep.module() torch.testing.assert_close(module(x), model(x)) + def test_recompose_rms_norm_only_takes_what_qnn_supports(self): + class Norm(torch.nn.Module): # a Llama-style norm + def __init__(self, dtype=torch.bfloat16): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(32, dtype=dtype)) + self.eps = 1e-6 + + def forward(self, x): + h = x.float() + h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.eps) + return self.weight * h.to(x.dtype) + + class BufferEpsNorm(Norm): + def __init__(self, dtype=torch.bfloat16): + super().__init__(dtype) + del self.eps + self.register_buffer("eps", torch.tensor(1e-6)) + + class ScaledNorm(torch.nn.Module): # a norm times a scale or a gate + def __init__(self, scale, dim=-1, keepdim=True, gate=False): + super().__init__() + self.dim, self.keepdim, self.gate = dim, keepdim, gate + self.register_buffer("scale", scale) + + def forward(self, x): + ms = x.pow(2).mean(self.dim, keepdim=self.keepdim) + h = x * torch.rsqrt(ms + 1e-6) + return h * (torch.sigmoid(x) if self.gate else self.scale) + + class TableNorm(ScaledNorm): # a norm of a buffer, added to the input + def __init__(self, scale): + super().__init__(scale) + self.register_buffer("table", torch.randn(4, 32)) + + def forward(self, x): + return x + super().forward(self.table) + + class StateNorm(ScaledNorm): # a norm of a buffer the model updates + def __init__(self): + super().__init__(torch.tensor(2.0)) + self.register_buffer("state", torch.zeros(4, 32)) + + def forward(self, x): + self.state.fill_(1.0) + return x + super().forward(self.state) + + # QNN's RmsNorm takes a mean over the last of 1 to 4 dims, kept as a + # dim, with static sizes and a constant eps. Anything else stays + # decomposed, since rms_norm has no CPU kernel. A norm whose weight is + # not one value per channel is rebuilt without it. A norm that reads + # only constants stays decomposed, since HTP may drop a partition that + # reads nothing at runtime. + bf16 = torch.bfloat16 + dynamic = {"x": {0: torch.export.Dim("n", min=2)}} + scale, weight = torch.tensor(2.0), torch.randn(32) + for i, (model, x, dynamic_shapes, kept_weights) in enumerate( + ( + (Norm(), torch.randn(1, 2, 3, 4, 32, dtype=bf16), None, []), + (ScaledNorm(weight), torch.randn(1, 2, 3, 4, 32), None, []), + (BufferEpsNorm(), torch.randn(4, 32, dtype=bf16), None, []), + (BufferEpsNorm(torch.float32), torch.randn(4, 32), None, []), + (ScaledNorm(weight), torch.randn(4, 32), dynamic, []), + (ScaledNorm(scale, gate=True), torch.randn(4, 32), dynamic, []), + (ScaledNorm(scale), torch.randn(4, 32), None, [False]), + (ScaledNorm(weight, dim=0), torch.randn(32, 32), None, []), + (ScaledNorm(weight, dim=None), torch.randn(32, 32), None, []), + (ScaledNorm(weight, keepdim=False), torch.randn(32, 32), None, []), + (TableNorm(scale), torch.randn(4, 32), None, []), + (TableNorm(weight), torch.randn(4, 32), None, []), + (StateNorm(), torch.randn(4, 32), None, [False]), + ) + ): + with self.subTest(row=i, model=type(model).__name__): + ep = torch.export.export( + model, (x,), dynamic_shapes=dynamic_shapes, strict=True + ) + ep = to_edge(ep).exported_program() + RecomposeRmsNorm( + edge_program=ep, compiler_specs=self._htp_compiler_specs(True) + )(ep.graph_module) + self.assertEqual( + [ + node.args[2] is not None + for node in ep.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ], + kept_weights, + ) + torch.testing.assert_close(ep.module()(x), model(x)) + + def test_recompose_rms_norm_rebuilds_a_normalization_for_htp_fp16_only(self): + class ScaledNorm(torch.nn.Module): # a norm times a 0-d scale + def __init__(self): + super().__init__() + self.register_buffer("scale", torch.tensor(2.0)) + + def forward(self, x): + rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) + return x * rms * self.scale + + # The normalization alone is rebuilt to keep HTP's float16 precision + # from overflowing. Elsewhere it stays decomposed, as every op in it + # has a CPU kernel. + x = torch.randn(4, 32) + for name, compiler_specs, kept_weights in ( + ("htp_fp16", self._htp_compiler_specs(True), [False]), + ("htp_quantized", self._htp_compiler_specs(False), []), + ("none", None, []), + ): + with self.subTest(specs=name): + ep = torch.export.export(ScaledNorm(), (x,), strict=True) + ep = to_edge(ep).exported_program() + # As in to_edge_transform_and_lower_to_qnn, which passes the specs. + passes = get_qnn_pass_manager_cls()().get_to_edge_transform_passes( + ep, compiler_specs=compiler_specs + ) + next(p for p in passes if isinstance(p, RecomposeRmsNorm))( + ep.graph_module + ) + self.assertEqual( + [ + node.args[2] is not None + for node in ep.graph.nodes + if node.target == exir_ops.edge.aten.rms_norm.default + ], + kept_weights, + ) + + def test_recompose_rms_norm_takes_no_runtime_weight(self): + class InputScaledNorm(torch.nn.Module): + def forward(self, x, scale): + rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) + return x * rms * scale + + # QNN's RmsNorm needs a constant weight, so a graph input is not taken + # as one. Edge graphs rebuild the normalization alone instead. + inputs = (torch.randn(4, 32), torch.randn(32)) + for capture, kept_weights in ((False, [False]), (True, [])): + with self.subTest(capture=capture): + ep = torch.export.export(InputScaledNorm(), inputs, strict=True) + if capture: + graph_module = ep.module() + RecomposeRmsNorm(quantization_capture=True)(graph_module) + else: + ep = to_edge(ep).exported_program() + graph_module = ep.graph_module + RecomposeRmsNorm( + edge_program=ep, compiler_specs=self._htp_compiler_specs(True) + )(graph_module) + self.assertEqual( + [ + node.args[2] is not None + for node in graph_module.graph.nodes + if node.target + in ( + exir_ops.edge.aten.rms_norm.default, + torch.ops.aten.rms_norm.default, + ) + ], + kept_weights, + ) + module = graph_module if capture else ep.module() + torch.testing.assert_close(module(*inputs), InputScaledNorm()(*inputs)) + def test_recompose_rms_norm_needs_the_norm_of_one_tensor(self): class AlmostNorm(torch.nn.Module): def __init__(self, scaled, squared): @@ -900,7 +1058,7 @@ def forward(self, x): pm = get_qnn_pass_manager_cls()() ep = pm.transform_for_export_pipeline(ep) ep = to_edge(ep).exported_program() - RecomposeRmsNorm()(ep.graph_module) + RecomposeRmsNorm(edge_program=ep)(ep.graph_module) self.assertNotIn( exir_ops.edge.aten.rms_norm.default, [node.target for node in ep.graph.nodes],