diff --git a/backends/qualcomm/_passes/recompose_rms_norm.py b/backends/qualcomm/_passes/recompose_rms_norm.py index f9cdd676345..d4013fbe8da 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,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) 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 2c1818f8a17..fa9130e7509 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/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):