diff --git a/devtools/inspector/_inspector_utils.py b/devtools/inspector/_inspector_utils.py index 556987e4bbf..11bb2b7d949 100644 --- a/devtools/inspector/_inspector_utils.py +++ b/devtools/inspector/_inspector_utils.py @@ -791,13 +791,15 @@ def _map_sequence_aot_output( def _find_matching_runtime_output_by_shape_and_dtype( aot_intermediate_output: torch.Tensor, runtime_intermediate_output: Sequence, + fallback_index: int, ) -> Any: """ Find the runtime output that matches the AOT output shape and dtype. - Used for multi-output operations (like native_layer_norm.out, native_dropout.out). + Used for multi-output operations (like native_layer_norm.out, native_dropout.out) + and for delegates that return more than one tensor. Returns: - The matching runtime output, or runtime_intermediate_output[-1] as fallback. + The matching runtime output, or runtime_intermediate_output[fallback_index]. """ # Find all runtime outputs that match the AOT shape matching_indices = [] @@ -824,14 +826,13 @@ def _find_matching_runtime_output_by_shape_and_dtype( # Exactly one dtype match - use it (e.g., dropout case where mask is bool) return runtime_intermediate_output[dtype_matching_indices[0]] - # No unique match found, return the last element as fallback - return runtime_intermediate_output[-1] + # No unique match found, fall back to positional pairing + return runtime_intermediate_output[fallback_index] def _map_non_sequence_aot_output( aot_intermediate_output: Any, runtime_intermediate_output: Any, - num_outputs: int, negative_index: int, ) -> Any: """ @@ -843,21 +844,33 @@ def _map_non_sequence_aot_output( if not isinstance(runtime_intermediate_output, Sequence): return runtime_intermediate_output - # Use the last element of the runtime output as fallback if no match is found + # Positional pairing is the fallback if no shape/dtype match is found. aot_mapped_runtime_intermediate_output = runtime_intermediate_output[negative_index] - # delegate runtime call and AOT intermediate is not a sequence. - # For multi-output operations (like native_layer_norm.out, native_dropout.out), - # the runtime captures all outputs but AOT only captures the primary output. - # We need to find the runtime output that matches the AOT output shape and dtype. - if ( - num_outputs == 1 - and len(runtime_intermediate_output) > 1 - and isinstance(aot_intermediate_output, torch.Tensor) + # The runtime event carries every tensor the op or delegate returned, while + # the AOT side has one tensor per node. Two cases land here: + # + # one runtime output per AOT node: a multi-output aten op + # (native_layer_norm.out, native_dropout.out) whose extra outputs AOT + # never captured. + # several runtime outputs: a delegate returning more than one tensor. The + # caller walks them positionally, taking the i-th-from-last debug handle + # in the delegate's handle tuple for the i-th-from-last runtime output. + # Nothing in the debug handles records which fused node produced which + # output, so handle order and output order agree only by luck -- with one + # output it is trivially true, with several it usually is not. + # + # Both cases are answered the same way: pick the runtime output whose shape + # and dtype match the AOT tensor, and only fall back to position when that + # search cannot single one out. + if len(runtime_intermediate_output) > 1 and isinstance( + aot_intermediate_output, torch.Tensor ): aot_mapped_runtime_intermediate_output = ( _find_matching_runtime_output_by_shape_and_dtype( - aot_intermediate_output, runtime_intermediate_output + aot_intermediate_output, + runtime_intermediate_output, + fallback_index=negative_index, ) ) @@ -903,7 +916,6 @@ def _process_single_runtime_output( aot_mapped_runtime_intermediate_output = _map_non_sequence_aot_output( aot_intermediate_output, runtime_intermediate_output, - num_outputs, negative_index, ) diff --git a/devtools/inspector/tests/inspector_utils_test.py b/devtools/inspector/tests/inspector_utils_test.py index cbdc557f405..1c9eee9493c 100644 --- a/devtools/inspector/tests/inspector_utils_test.py +++ b/devtools/inspector/tests/inspector_utils_test.py @@ -6,6 +6,7 @@ # pyre-unsafe +import math import tempfile import unittest from typing import Dict, Tuple @@ -44,6 +45,7 @@ propagate_back_debug_handle, TimeScale, ) +from executorch.devtools.inspector.numerical_comparator import SNRComparator from executorch.exir import to_edge from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY, UNSET_DEBUG_HANDLE @@ -423,6 +425,150 @@ def test_map_runtime_aot_intermediate_outputs_delegated(self): break self.assertTrue(found) + def _multi_output_delegate_fixture(self): + """A delegate whose output order does not match its debug handle order. + + A delegate is logged as one runtime event carrying every tensor it + returned, and the mapping walks those outputs positionally: the + i-th-from-last debug handle in the delegate's handle tuple is paired with + the i-th-from-last runtime output. Nothing in the debug handles records + which fused node produced which output, so the two orders line up only + when the delegate has a single output. Here they deliberately disagree, + which is the shape a real fused delegate has. + """ + aot_intermediate_outputs = { + (1,): torch.randn(1, 4, 512), + (2,): torch.randn(1, 4, 1408), + (3,): torch.randn(1, 4, 2048), + } + runtime_intermediate_outputs = { + (1, 2, 3): ( + [ + aot_intermediate_outputs[(2,)].clone(), + aot_intermediate_outputs[(3,)].clone(), + aot_intermediate_outputs[(1,)].clone(), + ], + 3, + ), + } + return aot_intermediate_outputs, runtime_intermediate_outputs + + def test_map_runtime_aot_intermediate_outputs_multi_output_delegate(self): + # Each AOT tensor must be paired with the runtime tensor of the same + # shape. Pairing by position instead hands the 512-wide AOT tensor the + # 1408-wide runtime one. + aot_intermediate_outputs, runtime_intermediate_outputs = ( + self._multi_output_delegate_fixture() + ) + + actual = map_runtime_aot_intermediate_outputs( + aot_intermediate_outputs, runtime_intermediate_outputs + ) + + self.assertEqual(len(actual), 3) + for (_, aot_output), (runtime_debug_handle, runtime_output) in actual.items(): + self.assertEqual(runtime_debug_handle, (1, 2, 3)) + self.assertEqual(aot_output.shape, runtime_output.shape) + self.assertTrue(torch.allclose(aot_output, runtime_output)) + + def test_snr_comparator_on_multi_output_delegate(self): + # The mispairing surfaced inside the comparator rather than the mapper: + # subtracting two differently shaped tensors raised + # Error computing SNR difference between tensors: The size of tensor a + # (512) must match the size of tensor b (1408) at non-singleton + # dimension 2 + # which aborted the whole comparison, not just the offending row. + aot_intermediate_outputs, runtime_intermediate_outputs = ( + self._multi_output_delegate_fixture() + ) + mapping = map_runtime_aot_intermediate_outputs( + aot_intermediate_outputs, runtime_intermediate_outputs + ) + + df = SNRComparator().compare(mapping, {}, {}) + + self.assertEqual(len(df), 3) + for gap in df["gap"]: + self.assertEqual(len(gap), 1) + self.assertFalse(math.isnan(gap[0])) + + def test_map_runtime_aot_intermediate_outputs_delegate_with_interior_nodes(self): + # A delegate that fused more nodes than it returns tensors, which is the + # usual shape once a whole block lands in one partition: + # + # 10 linear_gate -> [1, 4, 1408] + # 11 linear_up -> [1, 4, 1408] + # 12 silu -> [1, 4, 1408] + # 13 mul -> [1, 4, 1408] escapes the partition + # 14 linear_down -> [1, 4, 512] escapes the partition + # + # Handles 10-12 are interior. The two that escape are the last two in + # the handle tuple, so both positional AOT picks do land on real + # delegate outputs -- but the delegate's output list is ordered the + # other way round, because the partition's output spec is built + # independently of the backend's handle serialization order. Pairing by + # position therefore hands the 512-wide AOT tensor the 1408-wide runtime + # one and vice versa, and the comparator aborts on the shape mismatch. + aot_intermediate_outputs = { + (10,): torch.randn(1, 4, 1408), + (11,): torch.randn(1, 4, 1408), + (12,): torch.randn(1, 4, 1408), + (13,): torch.randn(1, 4, 1408), + (14,): torch.randn(1, 4, 512), + } + runtime_intermediate_outputs = { + (10, 11, 12, 13, 14): ( + [ + aot_intermediate_outputs[(14,)].clone(), + aot_intermediate_outputs[(13,)].clone(), + ], + 2, + ), + } + + actual = map_runtime_aot_intermediate_outputs( + aot_intermediate_outputs, runtime_intermediate_outputs + ) + + self.assertEqual(len(actual), 2) + for (_, aot_output), (_, runtime_output) in actual.items(): + self.assertEqual(aot_output.shape, runtime_output.shape) + self.assertTrue(torch.allclose(aot_output, runtime_output)) + + # The mismatch surfaces in the comparator, and one bad pair aborts the + # whole DataFrame rather than a single row. + df = SNRComparator().compare(actual, {}, {}) + self.assertEqual(len(df), 2) + for gap in df["gap"]: + self.assertEqual(len(gap), 1) + self.assertFalse(math.isnan(gap[0])) + + def test_map_runtime_aot_intermediate_outputs_multi_output_ambiguous_shapes(self): + # Shape and dtype cannot separate outputs that share both, so pairing + # stays positional there. Guards against the search falling back to the + # last runtime output, which would collapse every AOT tensor onto it. + aot_intermediate_outputs = { + (1,): torch.tensor([1.0, 2.0, 3.0]), + (2,): torch.tensor([4.0, 5.0, 6.0]), + } + runtime_intermediate_outputs = { + (1, 2): ( + [ + aot_intermediate_outputs[(1,)].clone(), + aot_intermediate_outputs[(2,)].clone(), + ], + 2, + ), + } + + actual = map_runtime_aot_intermediate_outputs( + aot_intermediate_outputs, runtime_intermediate_outputs + ) + + self.assertEqual(len(actual), 2) + for (_, aot_output), (_, runtime_output) in actual.items(): + self.assertTrue(torch.allclose(aot_output, runtime_output)) + def test_convert_input_to_tensor_convertible_inputs(self): # Scalar -> tensor actual_output1 = convert_to_float_tensor(5) diff --git a/devtools/inspector/tests/targets.bzl b/devtools/inspector/tests/targets.bzl index c08ac0595db..47df4727bee 100644 --- a/devtools/inspector/tests/targets.bzl +++ b/devtools/inspector/tests/targets.bzl @@ -48,6 +48,7 @@ def define_common_targets(is_fbcode = False): "//executorch/devtools/etdump:schema_flatcc", "//executorch/devtools/etrecord/tests:etrecord_test_library", "//executorch/devtools/inspector:inspector_utils", + "//executorch/devtools/inspector/numerical_comparator:lib", ], )