From 8598730fc15e63b62df1808d9b0d82965bf13cc8 Mon Sep 17 00:00:00 2001 From: Gyanendra Sinha Date: Sat, 5 Sep 2026 11:16:21 -0700 Subject: [PATCH] Fix the debugging searches that pruned work they never checked Six related defects in the debugging package, all of which made a tool report a clean or unchanged result over work it had not done. Search narrowing (search_strategy, validator, comparator): * `_depth_range` was computed and logged but never consulted when choosing what to yield, so bisection visited the whole graph at one model execution per side per batch. Both `_group_unchecked_nodes_by_depth` and `_get_next_batch_from_current_level` now filter on it. * The lower bound could advance past unchecked siblings when a level was larger than `batch_size`; the guard now tests `<=` the passed depth. * `Status.EXCLUDED` was projected to UNKNOWN, which strategies narrow on, so a graph opening with placeholders and views narrowed to depth 0 before a single value was compared. Adds `ValidationResult.SKIPPED`, which is neither evidence of a fault nor evidence against one. FX graph diff (graph_diff, graph_match): * Only top-level `fx.Node` args were followed, so `aten.cat`/`aten.stack` got no incoming edges and their inputs looked unused. Arguments are now flattened through lists, tuples and dicts, and kwargs are wired by name. * An FX node's identity was its op and target alone -- `_attr_digest` reads an `ir_object` an FX node does not have -- so `cat(dim=0)` and `cat(dim=1)` hashed identically and diffed as isomorphic. Constant args and the exported result type now feed the digest. Crashes: * `ComputeDevice.__missing__` is the dict hook, never called by Enum, so an unrecognised residency raised instead of degrading to UNKNOWN. * `_body_counts` recorded a call for a callee with no `coreai.graph`, and `_histogram` then raised KeyError on the whole histogram. * Benchmarker: interval id 0 was both the "not timed" sentinel and a real id; the sentinel now sits outside the counter's range and is rejected before lookup. `_state` transitions moved under the callback lock. Tests cover each behaviour change. Pre-existing failures in test_benchmarker (module timings) and test_intermediates are unaffected. --- coreai_torch/debugging/benchmarker.py | 38 ++-- coreai_torch/debugging/comparator.py | 25 ++- coreai_torch/debugging/compute_plan.py | 2 +- coreai_torch/debugging/graph_diff.py | 202 ++++++++++++++++++++-- coreai_torch/debugging/graph_match.py | 8 +- coreai_torch/debugging/histogram.py | 10 +- coreai_torch/debugging/search_strategy.py | 108 +++++++++--- coreai_torch/debugging/validator.py | 3 +- docs/api/debugging.md | 15 +- tests/debugging/test_compute_plan.py | 11 +- tests/debugging/test_graph_diff.py | 69 ++++++++ tests/debugging/test_search_strategy.py | 109 ++++++++++++ 12 files changed, 536 insertions(+), 64 deletions(-) diff --git a/coreai_torch/debugging/benchmarker.py b/coreai_torch/debugging/benchmarker.py index c248c49..65710d1 100644 --- a/coreai_torch/debugging/benchmarker.py +++ b/coreai_torch/debugging/benchmarker.py @@ -123,6 +123,10 @@ class _BenchmarkerState(Enum): """Benchmark completed.""" +_NO_INTERVAL = 0 +"""Interval id standing for "this begin event opened nothing".""" + + @dataclass(frozen=True) class Statistics: """Statistical summary of measurements.""" @@ -522,11 +526,6 @@ def _timing_annotation_callback( figure to give, and annotating each one repeated a single measurement as many times as the dispatch had members. - Grouped rather than indexed, because a representative is not unique. Two - dispatches regularly cover the same operation, and ``{timing.op_id: timing}`` - kept only the last -- 22 annotations from 34 dispatches on one model, silently, - with the survivor chosen by insertion order. - Args: timings: The dispatches to report. @@ -1302,7 +1301,7 @@ def __init__( self._timings: dict[tuple[int, tuple[int, ...]], list[float]] = defaultdict( list ) - self._interval_counter = 0 + self._interval_counter = _NO_INTERVAL + 1 self._debug_info_records: list[DebugInfoRecord] = [] # Maps a compiled op identifier (odix_id, delegate_id) to the list of # coreai op IDs fused into it. Keyed on the same pair carried by a runtime @@ -1386,7 +1385,7 @@ def _reset_state(self: Self) -> None: self._intervals.clear() self._timings.clear() self._pending_durations.clear() - self._interval_counter = 0 + self._interval_counter = _NO_INTERVAL + 1 def _wait_for_intervals_to_complete(self: Self, timeout_s: float = 10.0) -> None: """ @@ -1563,23 +1562,24 @@ def _on_log_event_begin(self: Self, event: LogEvent) -> int: event: LogEvent from the profiler Returns: - Interval ID for tracking this event + Interval ID for tracking this event, or :data:`_NO_INTERVAL` when this + event is not being timed. """ # Only process events when actively running benchmark if self._state != _BenchmarkerState.RUNNING: - return 0 # Return dummy interval_id + return _NO_INTERVAL # Only process inference phase events phase = _LogEventPhase(event.phase) if phase != _LogEventPhase.INFERENCE: - return 0 # Return dummy interval_id for non-inference events + return _NO_INTERVAL # Opening an interval for the host-timestamped twin would pool it with the # hardware measurement of the same encoder; its end event then finds no open # interval and is ignored. if _is_gpu_interval(event) and not _is_hardware_timestamped(event): - return 0 + return _NO_INTERVAL with self._lock: interval_id = self._interval_counter @@ -1602,6 +1602,12 @@ def _on_log_event_end(self: Self, event: LogEvent, interval_id: int) -> None: interval_id: Interval ID from the begin callback """ + # Nothing was opened for this event, so there is nothing to close. Checked + # before the lookup: the sentinel must never reach `self._intervals`, where it + # would find whichever interval happened to be open under that id. + if interval_id == _NO_INTERVAL: + return + # Only process events when actively running benchmark if self._state != _BenchmarkerState.RUNNING: return @@ -1819,8 +1825,11 @@ async def benchmark( logger.debug("Warmup run (untimed)") await function(nd_inputs) - # Transition to RUNNING state to start collecting timing data - self._state = _BenchmarkerState.RUNNING + # Transition to RUNNING state to start collecting timing data. Written + # under the lock the callbacks take, so the transition is ordered against + # their reads rather than racing them. + with self._lock: + self._state = _BenchmarkerState.RUNNING for i in range(num_runs): logger.debug("Benchmark run %d/%d", i + 1, num_runs) @@ -1833,7 +1842,8 @@ async def benchmark( await asyncio.to_thread(self._wait_for_intervals_to_complete) # All intervals have closed; stop collecting and build the result. - self._state = _BenchmarkerState.COMPLETED + with self._lock: + self._state = _BenchmarkerState.COMPLETED result = self._create_result() logger.info( diff --git a/coreai_torch/debugging/comparator.py b/coreai_torch/debugging/comparator.py index feebed9..cf76887 100644 --- a/coreai_torch/debugging/comparator.py +++ b/coreai_torch/debugging/comparator.py @@ -417,8 +417,13 @@ def __init__( source: Source debug graph containing computation graph and inspector target: Target debug graph containing computation graph and inspector id_map: Mapping from source node IDs to target node IDs - strategy: Search strategy to use on source graph. Defaults to bisection (batch_size=10) + strategy: Search strategy to use on source graph. Defaults to + :class:`~coreai_torch.debugging.search_strategy.ExhaustiveStrategy`, + which checks every operation in one batch; see the note below for why + that rather than bisection. show_progress: Whether to show progress bar during comparison (default: True) + exclude_ops: Torch operation names treated as deliberate exclusions when + explaining what the id_map dropped. """ self.source = source @@ -890,11 +895,23 @@ def _validation_result_to_status(vr: SearchStrategy.ValidationResult) -> Status: def _status_to_validation_result( status: Status, ) -> SearchStrategy.ValidationResult: - """Convert Status to ValidationResult.""" + """Project a Status onto what a search strategy acts on. + + Lossy by design, but the loss has to preserve one distinction: whether the + operation was ever a candidate. EXCLUDED was projected to UNKNOWN along with + everything else, and a strategy narrows on UNKNOWN -- so a graph opening with + placeholders and `aten.view`s, both EXCLUDED, narrowed a level-order search to + depth 0 before a single value was compared. + """ if status == Comparator.Status.PASS: return SearchStrategy.ValidationResult.PASS elif status == Comparator.Status.FAIL: return SearchStrategy.ValidationResult.FAIL + elif status == Comparator.Status.EXCLUDED: + # Never a candidate: a policy exclusion, or a node that computes nothing. + # Reporting it as unverified would have the search hunt for a fault in an + # operation there was never anything to check. + return SearchStrategy.ValidationResult.SKIPPED elif status == Comparator.Status.SHAPE_AMBIGUOUS: # Unverified, not clean: the values were never compared. Reporting it as PASS # would let a search prune a subgraph on the strength of a comparison that @@ -1513,7 +1530,9 @@ async def create_comparator_for_programs( target_program: AIProgram (target compiled model) target_entry_point: Name of the coreai.graph in target program inspector_type: Type of inspector for the target program - strategy: Search strategy for source graph (defaults to bisection) + strategy: Search strategy for source graph. Defaults to `ExhaustiveStrategy`, + which checks every operation in one batch and reports every + divergence; see :class:`Comparator` for why that rather than bisection. use_caching: Whether to use caching inspectors (default: True) exclude_ops: Frozenset of torch operation names to exclude from comparison. Defaults to _DEFAULT_EXCLUDED_OPS which includes view/reshape diff --git a/coreai_torch/debugging/compute_plan.py b/coreai_torch/debugging/compute_plan.py index 5bea98d..b7a71c6 100644 --- a/coreai_torch/debugging/compute_plan.py +++ b/coreai_torch/debugging/compute_plan.py @@ -162,7 +162,7 @@ class ComputeDevice(Enum): """Unknown compute device.""" @classmethod - def __missing__(cls, value: object) -> "ComputeDevice": + def _missing_(cls, value: object) -> "ComputeDevice": """Return UNKNOWN for unrecognized device values.""" return cls.UNKNOWN diff --git a/coreai_torch/debugging/graph_diff.py b/coreai_torch/debugging/graph_diff.py index 67bd197..1e28ff7 100644 --- a/coreai_torch/debugging/graph_diff.py +++ b/coreai_torch/debugging/graph_diff.py @@ -16,7 +16,7 @@ import re import sys -from collections.abc import Collection +from collections.abc import Collection, Iterator from dataclasses import dataclass, field, fields from enum import Enum from io import StringIO @@ -1546,6 +1546,123 @@ def format_multi_graph_diff( # --------------------------------------------------------------------------- +def _fx_nodes_in(value: Any) -> Iterator[torch.fx.Node]: + """ + Every FX node reachable inside an argument, in a stable order. + + An FX argument is not always a node: `aten.cat` takes its inputs as a *list*, and a + dict or a nested tuple is equally legal. Only top-level `fx.Node` args used to be + followed, so a concatenation, a stack or a `_native_batch_norm` tuple contributed no + edges at all, and the nodes feeding them looked unused to the comparison. + + Args: + value: An FX argument, of any shape. + + Yields: + The nodes inside it, outer to inner, left to right. + + """ + if isinstance(value, torch.fx.Node): + yield value + elif isinstance(value, (list, tuple)): + for item in value: + yield from _fx_nodes_in(item) + elif isinstance(value, dict): + for _, item in sorted(value.items(), key=lambda entry: str(entry[0])): + yield from _fx_nodes_in(item) + + +def _fx_constant(value: Any) -> Any: + """ + An FX argument with its node references replaced, for the attribute digest. + + Nodes are edges, not configuration, so they are collapsed to a placeholder: their + identity is carried by the graph, and their *names* shift whenever anything upstream + changes. A tensor is reduced to its shape and dtype for the reason + `WeightPolicy.IGNORE` exists -- a rebuild re-initialises parameters, so comparing + values reports every weight of an unchanged model as changed. + + Args: + value: An FX argument, of any shape. + + Returns: + The same argument as plain, comparable values. + + """ + if isinstance(value, torch.fx.Node): + return "" + if isinstance(value, (list, tuple)): + return [_fx_constant(item) for item in value] + if isinstance(value, dict): + return {str(key): _fx_constant(item) for key, item in sorted(value.items())} + if isinstance(value, torch.Tensor): + return f"tensor<{tuple(value.shape)}x{value.dtype}>" + return repr(value) + + +def _fx_attributes(fx_node: Any) -> str: + """ + An FX node's configuration: the parts of its arguments that are not nodes. + + The counterpart of `_attr_digest` for a graph with no MLIR behind it. Without it + `aten.mean(x, dim=-1)` and `aten.mean(x, dim=0)` share an identity, as do two + `aten.to` calls to different dtypes. + + Args: + fx_node: The FX node to describe. + + Returns: + The node's constant arguments and keyword arguments, as a comparable string. + + """ + parts = [repr(_fx_constant(arg)) for arg in fx_node.args] + parts.extend( + f"{name}={_fx_constant(value)!r}" + for name, value in sorted(fx_node.kwargs.items()) + ) + return "|".join(parts) + + +def _fx_result_type(fx_node: Any) -> str: + """ + The shape and dtype an FX node produces, when export recorded them. + + Where a torch graph keeps what MLIR keeps in a result type. An FX graph is not + bipartite -- a node *is* its result -- so this goes in the node's own `ir_type`, + which `structural_labels` drops and `node_labels` keeps. That split is what lets a + node whose shape changed still pair, and then read as modified. + + Args: + fx_node: The FX node to describe. + + Returns: + The type as text, or an empty string when export recorded none. + + """ + value = fx_node.meta.get("val") + if isinstance(value, (list, tuple)): + return ",".join(_fx_result_type_of(item) for item in value) + return _fx_result_type_of(value) + + +def _fx_result_type_of(value: Any) -> str: + """ + One tensor's shape and dtype. + + Args: + value: A fake tensor from `node.meta["val"]`, or any other metadata value. + + Returns: + The type as text, or an empty string when it is not a tensor. + + """ + shape = getattr(value, "shape", None) + dtype = getattr(value, "dtype", None) + if shape is None or dtype is None: + return "" if value is None else type(value).__name__ + return f"tensor<{'x'.join(str(dimension) for dimension in shape)}x{dtype}>" + + class _TorchFXGraphBuilder: """Helper class to build NetworkX graph from PyTorch FX graphs.""" @@ -1564,20 +1681,45 @@ def build(self, fx_graph: Any) -> nx.DiGraph: # Add edges based on node inputs for node in fx_graph.nodes: if node in self.fx_node_to_id: - node_id = self.fx_node_to_id[node] - for i, arg in enumerate(node.args): - # Check if arg is an FX node (has 'op' attribute and is in our mapping) - if hasattr(arg, "op") and arg in self.fx_node_to_id: - arg_id = self.fx_node_to_id[arg] - self.graph.add_edge( - arg_id, - node_id, - edge_type="data_flow", - index=i, - ) + self._add_edges(node) return self.graph + def _add_edges(self, fx_node: Any) -> None: + """ + Add one edge per node this node consumes, positional args then keywords. + + The slot an edge occupies is `(edge_type, index)`, and operand order is + semantic, so positional inputs are numbered across the *flattened* arguments: + `cat([a, b])` gives `a` slot 0 and `b` slot 1, where before it gave neither an + edge. Keyword inputs are keyed by name instead of position, since that is what + identifies them -- reordering kwargs is not a change. + + Args: + fx_node: The FX node whose inputs are being wired up. + + """ + node_id = self.fx_node_to_id[fx_node] + + for index, argument in enumerate(_fx_nodes_in(fx_node.args)): + if argument in self.fx_node_to_id: + self.graph.add_edge( + self.fx_node_to_id[argument], + node_id, + edge_type="data_flow", + index=index, + ) + + for name, value in sorted(fx_node.kwargs.items()): + for index, argument in enumerate(_fx_nodes_in(value)): + if argument in self.fx_node_to_id: + self.graph.add_edge( + self.fx_node_to_id[argument], + node_id, + edge_type=f"kwarg:{name}", + index=index, + ) + def _get_next_id(self) -> int: """Generate a unique node ID.""" node_id = self.node_counter @@ -1593,13 +1735,17 @@ def _process_fx_node(self, fx_node: Any) -> None: op_type = str(fx_node.op) target = str(fx_node.target) if fx_node.target else "unknown" - # Add node to graph + # Add node to graph. `attributes` and `ir_type` are precomputed here rather + # than read from an `ir_object`, which an FX node does not have: see + # `graph_match._attr_digest` for the fallback they feed. self.graph.add_node( node_id, type="op", op_name=f"{op_type}:{target}", op_type=op_type, target=target, + attributes=_fx_attributes(fx_node), + ir_type=_fx_result_type(fx_node), torch_object=fx_node, ) @@ -1622,16 +1768,36 @@ def _build_torch_fx_graph(exported_program: torch.export.ExportedProgram) -> nx. def compute_exported_program_diff( source_program: torch.export.ExportedProgram, target_program: torch.export.ExportedProgram, + *, + ignore_attributes: Collection[str] = UNSTABLE_ATTRIBUTES, ) -> GraphDiff: """ Compute structural diff between two PyTorch ExportedPrograms. - Extracts the FX graphs, builds NetworkX graphs, and computes structural diff - using graph isomorphism. + Extracts the FX graphs, builds NetworkX graphs, and compares them with the same + `graph_match.align` the Core AI diff uses. + + A node's identity is its op and target, its constant arguments and keyword + arguments, and the shape and dtype export recorded for it. Its edges are every FX + node reachable inside its arguments, so an input passed in a list -- `aten.cat`, + `aten.stack` -- is followed like any other. + + Two differences from :func:`compute_coreai_program_diff` are worth knowing: + + * **An FX graph is not bipartite.** A node is its own result, so there are no value + nodes and `responsible_op` is the identity. Node counts are therefore much smaller + than the Core AI ones for the same model, and not comparable with them. + * **There is no `WeightPolicy`.** Parameters live in the module's state dict rather + than in the graph, so the graph carries no values to compare; a parameter appears + only as a `placeholder` named after it. A retrained model diffs as unchanged here. + Use :func:`compute_coreai_program_diff` with `DIGEST` to compare weights. Args: source_program: Source (reference/expected) ExportedProgram target_program: Target (actual/test) ExportedProgram + ignore_attributes: Attribute names left out of a node's identity. Accepted for + symmetry with the Core AI entry points; an FX node carries no named MLIR + attributes, so this reaches only the labelling of nodes that have them. Returns: GraphDiff object with source_graph and target_graph included @@ -1642,4 +1808,8 @@ def compute_exported_program_diff( target_graph = _build_torch_fx_graph(target_program) # Compute and return diff - return compute_graph_diff(source_graph, target_graph) + return compute_graph_diff( + source_graph, + target_graph, + ignore_attributes=ignore_attributes, + ) diff --git a/coreai_torch/debugging/graph_match.py b/coreai_torch/debugging/graph_match.py index 74cd800..da2c3ef 100644 --- a/coreai_torch/debugging/graph_match.py +++ b/coreai_torch/debugging/graph_match.py @@ -435,12 +435,18 @@ def _attr_digest( normalised away: both differ between two conversions of one unchanged model, so including either reports a change that did not happen. Read through the bindings, never by parsing printed IR. + + A graph with no MLIR behind it -- a torch FX graph -- has no attributes to read, so + a builder for one precomputes the digest and stores it under `attributes`. Without + that fallback every such node's digest was the empty string, and an op's whole + configuration (`dim`, `dtype`, a scalar operand) was absent from its identity. """ ir_object = attrs.get("ir_object") operation = getattr(ir_object, "operation", ir_object) attributes = getattr(operation, "attributes", None) if attributes is None: - return "" + precomputed = attrs.get("attributes") + return "" if precomputed is None else str(precomputed) parts: list[str] = [] for attribute in attributes: diff --git a/coreai_torch/debugging/histogram.py b/coreai_torch/debugging/histogram.py index 3a1839e..b9553ab 100644 --- a/coreai_torch/debugging/histogram.py +++ b/coreai_torch/debugging/histogram.py @@ -283,7 +283,12 @@ def _body_counts( counts[operation.name] += 1 continue callee = _invoke_callee(operation) - if callee is None: + # A callee this module defines no `coreai.graph` for has no body to + # describe, so it is counted as a plain invoke. Counting it as a call + # meant `_histogram` looked the symbol up in `bodies` and raised + # KeyError -- an invoke naming a `func.func`, or an external symbol, + # crashed the whole histogram rather than being reported as itself. + if callee is None or callee not in entry_points: counts[operation.name] += 1 continue calls[callee] += 1 @@ -303,6 +308,9 @@ def _histogram( Recurses without a depth or cycle guard: a `coreai.invoke` names a symbol the module defines, and the graphs form a DAG, so the walk terminates. + Every symbol in *calls* has an entry in *bodies*: `_body_counts` records a call + only for a callee `entry_points` holds, and *bodies* is keyed by exactly those. + Args: counts: The body's own operation counts. calls: How many invokes in this body name each callee. diff --git a/coreai_torch/debugging/search_strategy.py b/coreai_torch/debugging/search_strategy.py index fe3bfa6..c0d1da9 100644 --- a/coreai_torch/debugging/search_strategy.py +++ b/coreai_torch/debugging/search_strategy.py @@ -13,7 +13,7 @@ import logging from abc import ABC, abstractmethod from collections.abc import AsyncIterator, Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum, auto from typing import Generic, TypeVar @@ -39,6 +39,18 @@ class ValidationResult(Enum): PASS = auto() FAIL = auto() UNKNOWN = auto() + SKIPPED = auto() + """This node was never a candidate, so nothing was checked. + + A placeholder, a `get_attr`, the `output` node, or an operation an exclusion + policy removed. Distinct from UNKNOWN, which means a value was expected and + could not be obtained: a strategy narrows on UNKNOWN because unverified work + might hide the fault, and doing that for a node that computes nothing collapsed + the search on its first batch -- every FX graph starts with placeholders at + depth 0, so `LevelOrderStrategy` narrowed to depth 0 before checking anything. + + Recorded as a result rather than withheld: a strategy tracks which nodes are + still unchecked, so a node it is never told about is one it offers forever.""" @dataclass class CategorizedResults: @@ -47,6 +59,10 @@ class CategorizedResults: failed: list[ComputationGraph.Node] passed: list[ComputationGraph.Node] unknown: list[ComputationGraph.Node] + skipped: list[ComputationGraph.Node] = field(default_factory=list) + """Nodes that were never candidates; see `ValidationResult.SKIPPED`. Neither + evidence of a fault nor evidence against one, so no pass counts them and no + narrowing acts on them.""" def __aiter__(self) -> AsyncIterator[list[ComputationGraph.Node]]: return self @@ -357,11 +373,33 @@ def _initialize_search_scope(self) -> None: self._initialized = True + def _in_depth_range(self, node: ComputationGraph.Node) -> bool: + """ + Whether *node* still falls inside the depth range the search has narrowed to. + + Args: + node: Node to test. + + Returns: + True when the node is still a candidate. + + """ + if self._depth_range is None: + return True + min_depth, max_depth = self._depth_range + return min_depth <= node.depth < max_depth + def _group_unchecked_nodes_by_depth( self, ) -> dict[int, list[ComputationGraph.Node]]: """ - Group unchecked nodes by their depth level. + Group the still-eligible unchecked nodes by their depth level. + + Filtered by :attr:`_depth_range`, which is what makes narrowing mean anything. + Without the filter the range was computed, logged and never consulted except + for the `min_depth >= max_depth` exit, so every level was offered whatever the + results said and bisection reordered the graph instead of pruning it -- at one + model execution per side per batch, the most expensive way to check everything. Returns: Dictionary mapping depth -> list of unchecked nodes at that depth @@ -370,10 +408,11 @@ def _group_unchecked_nodes_by_depth( level_dict: dict[int, list[ComputationGraph.Node]] = {} for node in self._scope_nodes: - if node.op_id not in self._node_results: - if node.depth not in level_dict: - level_dict[node.depth] = [] - level_dict[node.depth].append(node) + if node.op_id in self._node_results: + continue + if not self._in_depth_range(node): + continue + level_dict.setdefault(node.depth, []).append(node) return level_dict @@ -422,18 +461,24 @@ def _get_next_batch_from_current_level( """ Get the next batch of nodes from the currently selected level. + Nodes narrowed out since the level was selected are stepped over rather than + yielded: a failure narrows the range mid-level, and the remainder of that level + was still being handed out because only level *selection* consulted the range. + Returns: - Next batch_size nodes from current position, or empty list if exhausted + Next batch of eligible nodes from the current position, or an empty list + when the level is exhausted. """ - if self._current_level_index >= len(self._current_level_nodes): - return [] - - # Get batch_size nodes from current position - batch = self._current_level_nodes[ - self._current_level_index : self._current_level_index + self.batch_size - ] - self._current_level_index += self.batch_size + batch: list[ComputationGraph.Node] = [] + while ( + self._current_level_index < len(self._current_level_nodes) + and len(batch) < self.batch_size + ): + node = self._current_level_nodes[self._current_level_index] + self._current_level_index += 1 + if self._in_depth_range(node): + batch.append(node) return batch @@ -544,12 +589,15 @@ def _categorize_validation_results( failed_nodes = [] passed_nodes = [] unknown_nodes = [] + skipped_nodes = [] for node, result in results: if result == SearchStrategy.ValidationResult.FAIL: failed_nodes.append(node) elif result == SearchStrategy.ValidationResult.PASS: passed_nodes.append(node) + elif result == SearchStrategy.ValidationResult.SKIPPED: + skipped_nodes.append(node) else: unknown_nodes.append(node) @@ -557,6 +605,7 @@ def _categorize_validation_results( failed=failed_nodes, passed=passed_nodes, unknown=unknown_nodes, + skipped=skipped_nodes, ) def _track_failed_parent_nodes_for_descent( @@ -656,6 +705,15 @@ def _narrow_search_range_on_failure( searching for failures. Don't narrow when everything passes, as we need to ensure all nodes eventually get checked. + SKIPPED results are ignored on both counts. They are neither evidence of a + fault nor evidence against one, so they must not pull the upper bound down -- + and they must not let the lower bound advance past depths that were never + actually verified. + + The range is enforced by :meth:`_group_unchecked_nodes_by_depth` and + :meth:`_get_next_batch_from_current_level`, so narrowing here genuinely stops + deeper nodes being offered. + Args: categorized: Categorized validation results from the last batch @@ -675,21 +733,25 @@ def _narrow_search_range_on_failure( # This focuses search on finding the root cause at or before the issue max_depth = min(max_depth, min_problematic_depth + 1) self._depth_range = (min_depth, max_depth) - # If all nodes passed, check if there are unchecked nodes at shallower depths - # Only narrow the lower bound if all shallower depths have been checked + # If all nodes passed, check if there are unchecked nodes at this depth or + # shallower. Only narrow the lower bound once they have all been checked. elif categorized.passed: # Find the maximum depth among passed nodes max_passed_depth = max(node.depth for node in categorized.passed) - # Check if there are any unchecked nodes at depths < max_passed_depth - unchecked_at_shallower_depths = any( - node.depth < max_passed_depth and node.op_id not in self._node_results + # `<=`, not `<`: the bound about to be set claims everything up to *and + # including* `max_passed_depth` has been checked, so the level itself has to + # be finished. A level larger than `batch_size` arrives in several batches, + # and testing only strictly shallower depths let the first batch advance the + # bound past its own unchecked siblings. Inert while the range was never + # consulted; now that it prunes, those siblings would simply never be + # checked, and a validator would report a clean result having skipped them. + unchecked_at_or_above = any( + node.depth <= max_passed_depth and node.op_id not in self._node_results for node in self._scope_nodes ) - # Only narrow if no unchecked nodes exist at shallower depths - # This ensures we eventually check all nodes when everything passes - if not unchecked_at_shallower_depths: + if not unchecked_at_or_above: # Narrow the lower bound to just after the passed depth # We know we've checked everything up to and including this depth min_depth = max(min_depth, max_passed_depth + 1) diff --git a/coreai_torch/debugging/validator.py b/coreai_torch/debugging/validator.py index ad66d29..5cc0677 100644 --- a/coreai_torch/debugging/validator.py +++ b/coreai_torch/debugging/validator.py @@ -289,9 +289,8 @@ async def check( # apart from the unknowns rather than counted as a failed retrieval. excluded.append(node) batch_results.append( - (node, SearchStrategy.ValidationResult.UNKNOWN), + (node, SearchStrategy.ValidationResult.SKIPPED), ) - unknown_count += 1 continue outputs = results.get(node.op_id) diff --git a/docs/api/debugging.md b/docs/api/debugging.md index 1ffc997..a826f22 100644 --- a/docs/api/debugging.md +++ b/docs/api/debugging.md @@ -207,9 +207,15 @@ result = await validator.check(check_large_values, inputs=example_input) Choose how to search through operations: ```python -from coreai_torch.debugging.search_strategy import LevelOrderStrategy +from coreai_torch.debugging.search_strategy import ( + ExhaustiveStrategy, + LevelOrderStrategy, +) + +# Exhaustive (default) - one batch, one model execution per side, reports every issue +strategy = ExhaustiveStrategy(graph) -# Binary search (default - fastest for finding first issue) +# Bisection - narrows by depth to the first failing level strategy = LevelOrderStrategy.bisection(graph, batch_size=10) # Top-down (systematic from inputs to outputs) @@ -219,6 +225,11 @@ strategy = LevelOrderStrategy.top_down(graph) strategy = LevelOrderStrategy.auto(graph) ``` +Every batch a strategy yields costs a full model execution on *both* sides, so +narrowing only pays when capturing a value is more expensive than a run — very large +intermediates, or an early exit that skips most of the graph. Otherwise the default +`ExhaustiveStrategy` is both cheaper and more complete. + ### Batch size ```python # Control batch size for memory efficiency diff --git a/tests/debugging/test_compute_plan.py b/tests/debugging/test_compute_plan.py index 05ae1d3..8b9c857 100644 --- a/tests/debugging/test_compute_plan.py +++ b/tests/debugging/test_compute_plan.py @@ -67,7 +67,7 @@ async def test_compute_plan_from_program( ) # An entry the planner made must name a device it chose. `<= set(ComputeDevice)` - # was the assertion here and cannot fail: `ComputeDevice.__missing__` maps any + # was the assertion here and cannot fail: `ComputeDevice._missing_` maps any # string at all to `UNKNOWN`, and `UNKNOWN` is itself a member, so a plan that # resolved nothing passed exactly as one that resolved everything. Measured on # this model: 31 entries all CPU under the bundled runtime, 10 all GPU under the @@ -184,3 +184,12 @@ async def test_compute_plan_annotate_source( assert any(device.value in output for device in ComputeDevice), ( "Expected at least one compute device label in the annotated output" ) + + +def test_known_residencies_still_parse() -> None: + """The names the enum does carry keep resolving, case- and alias-insensitively.""" + assert ComputeDevice.from_string("cpu") is ComputeDevice.CPU + assert ComputeDevice.from_string(" GPU ") is ComputeDevice.GPU + # "ANE" is the runtime's spelling of the Neural Engine. + assert ComputeDevice.from_string("ane") is ComputeDevice.NEURAL_ENGINE + assert ComputeDevice.from_string("NEURAL_ENGINE") is ComputeDevice.NEURAL_ENGINE diff --git a/tests/debugging/test_graph_diff.py b/tests/debugging/test_graph_diff.py index c06b697..812907a 100644 --- a/tests/debugging/test_graph_diff.py +++ b/tests/debugging/test_graph_diff.py @@ -14,6 +14,7 @@ from coreai_torch.converter import TorchConverter from coreai_torch.debugging.graph_diff import ( + _build_torch_fx_graph, compute_coreai_program_diff, compute_exported_program_diff, compute_per_graph_diff, @@ -753,3 +754,71 @@ async def test_per_graph_diff_applies_one_identity_to_every_graph() -> None: if diff is not None: assert diff.ignore_attributes == frozenset(no_values) assert diff.weights is WeightPolicy.DIGEST + + +class _ConcatModel(torch.nn.Module): + """A model whose inputs reach an operation inside a *list* argument.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Concatenate along the configured dimension.""" + return torch.cat([x, y], dim=self.dim).sum() + + +def _concat_program(dim: int) -> torch.export.ExportedProgram: + """Export `_ConcatModel` at *dim*.""" + args = (torch.randn(2, 4), torch.randn(2, 4)) + return torch.export.export(_ConcatModel(dim), args).run_decompositions() + + +def test_exported_program_follows_nodes_inside_list_arguments() -> None: + """An input passed in a list is still an operand. + + Only top-level `fx.Node` args were followed, so `aten.cat` -- which takes its inputs + as a list -- got no incoming edges at all, and the tensors feeding it looked unused. + Measured on this two-input model: 1 edge across 5 nodes, against 4 now. + """ + graph = _build_torch_fx_graph(_concat_program(0)) + + concat = [ + node + for node, data in graph.nodes(data=True) + if "cat" in data.get("op_name", "") + ] + assert len(concat) == 1, "Expected exactly one concatenation in the graph" + + operands = sorted( + (data["index"], producer) + for producer, _, data in graph.in_edges(concat[0], data=True) + ) + assert [index for index, _ in operands] == [0, 1], ( + f"Both list members should be wired as operands, got {operands}" + ) + + +def test_exported_program_diff_detects_a_changed_constant_argument() -> None: + """Two graphs differing only in a keyword argument are not the same graph. + + An FX node's label was its op and target alone -- `_attr_digest` reads an + `ir_object`, which an FX node does not have -- so its whole configuration was + outside its identity. `cat(dim=0)` and `cat(dim=1)` hashed identically and the diff + reported them isomorphic: a silent false negative, the worst answer a diff can give. + """ + diff = compute_exported_program_diff(_concat_program(0), _concat_program(1)) + + assert not diff.is_isomorphic, "A changed `dim` must not read as the same graph" + assert diff.summary.modified_node_count > 0, ( + "The concatenation is still the same operation, so it should be reported as " + "modified rather than removed and added" + ) + + +def test_exported_program_diff_still_matches_identical_programs() -> None: + """The stronger identity must not make two exports of one model differ.""" + diff = compute_exported_program_diff(_concat_program(0), _concat_program(0)) + + assert diff.is_isomorphic + assert diff.summary.modified_node_count == 0 diff --git a/tests/debugging/test_search_strategy.py b/tests/debugging/test_search_strategy.py index d1c795f..f2d1ccc 100644 --- a/tests/debugging/test_search_strategy.py +++ b/tests/debugging/test_search_strategy.py @@ -343,3 +343,112 @@ async def test_get_problematic_operations() -> None: assert failed_node.op_id in problematic_ids, ( f"Failed node {failed_node.op_id} not in problematic operations" ) + + +async def _drain( + strategy: LevelOrderStrategy, + result: SearchStrategy.ValidationResult, +) -> list[ComputationGraph.Node]: + """Run *strategy* to completion, reporting *result* for everything it yields.""" + seen: list[ComputationGraph.Node] = [] + while True: + try: + batch = await strategy.__anext__() + except StopAsyncIteration: + return seen + seen.extend(batch) + await strategy.update([(node, result) for node in batch]) + + +async def test_narrowing_actually_stops_deeper_nodes_being_offered() -> None: + """A narrowed range must prune, not merely be recorded. + + `_depth_range` was computed and logged but never consulted when choosing what to + yield -- only for the `min_depth >= max_depth` exit -- so every level was offered + whatever the results said. Bisection therefore visited the whole graph at one model + execution per side per batch: the most expensive way to check everything, and not + what it claims to do. + + The fault has to leave something *unverified* at its own depth for this to be + visible. When a level comes back wholly clean the lower bound marches up to meet the + narrowed upper bound and the search ends anyway, which is why a simpler version of + this test passed against the unenforced range. + """ + graph = create_hierarchical_graph_with_dependencies() + strategy = LevelOrderStrategy.top_down(graph) + + # Depth 0 is clean. + batch = await strategy.__anext__() + assert {node.depth for node in batch} == {0} + await strategy.update( + [(node, SearchStrategy.ValidationResult.PASS) for node in batch] + ) + + # Depth 1 holds the fault, and its sibling could not be verified -- so the lower + # bound cannot advance and only the upper bound moves. + batch = await strategy.__anext__() + assert {node.depth for node in batch} == {1} + assert len(batch) == 2 + await strategy.update( + [ + (batch[0], SearchStrategy.ValidationResult.FAIL), + (batch[1], SearchStrategy.ValidationResult.UNKNOWN), + ] + ) + + remaining = await _drain(strategy, SearchStrategy.ValidationResult.PASS) + assert all(node.depth <= 1 for node in remaining), ( + "After a failure at depth 1 the search still offered depths " + f"{sorted({node.depth for node in remaining})}" + ) + assert {node.op_id for node in strategy.get_problematic_operations()} == { + batch[0].op_id + } + + +async def test_a_level_is_finished_before_the_lower_bound_advances() -> None: + """A level larger than the batch must not be left half-checked. + + The bound says "everything up to and including this depth has been checked", but the + guard tested only strictly shallower depths, so the first batch of a split level + advanced it past its own unchecked siblings. Harmless while the range was never + consulted; once it prunes, those siblings are dropped and the validator reports a + clean result over a graph it did not finish. + """ + graph = create_hierarchical_graph_with_dependencies() + # Depth 0 holds two nodes, so a batch of one splits it. + strategy = LevelOrderStrategy.top_down(graph, batch_size=1) + + seen = await _drain(strategy, SearchStrategy.ValidationResult.PASS) + + checked = {node.op_id for node in seen} + top_level = {node.op_id for node in graph.get_nodes_in_scope((None, 0))} + assert top_level <= checked, ( + f"Never checked {sorted(top_level - checked)}, which a clean run must cover" + ) + + +async def test_skipped_nodes_do_not_narrow_the_search() -> None: + """A node that computes nothing is not evidence, and must not end the search. + + Placeholders sit at depth 0 of every FX graph and the validator reports them as + SKIPPED. Reported as UNKNOWN they narrowed the upper bound to depth 0 on the very + first batch, so a level-order search finished having verified nothing -- and said + so only as an empty failure list, which reads exactly like a clean model. + """ + graph = create_hierarchical_graph_with_dependencies() + strategy = LevelOrderStrategy.top_down(graph) + + batch = await strategy.__anext__() + assert {node.depth for node in batch} == {0} + await strategy.update( + [(node, SearchStrategy.ValidationResult.SKIPPED) for node in batch] + ) + + remaining = await _drain(strategy, SearchStrategy.ValidationResult.PASS) + depths = {node.depth for node in remaining} + assert depths == {1, 2, 3}, f"Search stopped early, only reaching depths {depths}" + assert not strategy.get_problematic_operations() + assert not strategy.get_unknown_operations(), ( + "A skipped node is not an unknown one: it was never a candidate" + )