Qualcomm AI Engine Direct - Keep the dtype of a recomposed RMSNorm

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,91 @@ 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.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):
+        """
+        Whether the pattern multiplies by the root mean square the same tensor
+        it squares, so it really is a norm of one tensor.
+        """
+        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 = square.args[1] is base
+        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
+        # The nearest tensor in the output dtype, past casts that keep it.
+        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 +202,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)

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

backends/qualcomm/tests/test_passes.py

@@ -578,6 +578,191 @@ 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):  # an upcast written with type_as
+            def __init__(self, dtype):
+                super().__init__()
+                self.weight = torch.nn.Parameter(torch.randn(32, dtype=dtype))
+
+            def forward(self, x):
+                h = x.float()
+                h = h * torch.rsqrt((h * h).mean(-1, keepdim=True) + 1e-6)
+                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
+        # in another dtype is fine when it is about to be quantized, since the
+        # quantizer gives it the input's quantized type. Edge graphs carry casts
+        # as _to_dim_order_copy, or as _to_copy without dim order, and captured
+        # graphs as to.dtype.
+        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)]),
+            (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_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):