Skip to content

Report debugging findings by the module and line they came from - #114

Merged
cymbalrush merged 2 commits into
apple:mainfrom
cymbalrush:feat/module-attribution
Sep 24, 2026
Merged

cymbalrush merged 2 commits into
apple:mainfrom
cymbalrush:feat/module-attribution

Conversation

@cymbalrush

Copy link
Copy Markdown
Contributor

graph_diff answers which node became which, plan_diff which op id moved device, intermediates which op id diverged. None of those is the question a reader asked. A node id or an op id names nothing they wrote, and a report over a real model is not read at all unless it is grouped -- so every tool here stopped one step short of the answer, and each stopped in its own way.

This adds the step, once, and points the three reports at it.

modules.py -- the rollup, shared rather than per-tool.

build_module_tree groups any per-operation finding into the module instance tree it came from, generalising the insert-or-create walk that was inline in BenchmarkResult.get_module_timings. attributions_by_op_id composes the two existing lookups, which deliberately disagree on how to filter source locations: provenance must not trim frames or a model of stock nn modules collapses to one path, while a source line must trim them or it names a line inside torch.

Instances, not types. Block$1 and Block$3 are different answers, and an operation name cannot tell them apart: a model with twelve identical blocks has twelve identical linear operations, and "the slow one is a linear" is not a finding.

changes.py -- compute_changes, built on that rollup.

Three things stand between a correspondence and "7 operations changed, 6 of them in Block$3, at model.py:41". Each is a measured defect in the naive version rather than a refinement:

  • A difference lands on a value node as readily as on an operation. Reported as-is, a rewiring inside a callee body says nothing was modified. Resolved through responsible_op.
  • Most differing pairs are consequences of one edit, not edits: operands are other operations, so inserting one changes every neighbour's slots. _is_independent_change classifies each disagreeing slot by what happened to its endpoints and keeps only the two transitions nothing else reports -- a rewiring and a substitution. Chawathe et al. (1996) minimality.
  • Which instance of an interchangeable operation is left unmapped is arbitrary. Unaided, adding a third block reports its extra constant against Block$1, and a reader concludes something changed in a block that did not. _relocate_to_new_modules breaks the tie on module scope, which the matcher never looked at. Additions and removals only: a modification names a pair and is interchangeable with nothing, so re-seating one end is simply false.

plan_diff and intermediates -- the same rollup, so "which layer" is one call whatever was measured.

`graph_diff` answers which node became which, `plan_diff` which op id moved
device, `intermediates` which op id diverged. None of those is the question a
reader asked. A node id or an op id names nothing they wrote, and a report over a
real model is not read at all unless it is grouped -- so every tool here stopped
one step short of the answer, and each stopped in its own way.

This adds the step, once, and points the three reports at it.

modules.py -- the rollup, shared rather than per-tool.

`build_module_tree` groups any per-operation finding into the module instance tree
it came from, generalising the insert-or-create walk that was inline in
`BenchmarkResult.get_module_timings`. `attributions_by_op_id` composes the two
existing lookups, which deliberately disagree on how to filter source locations:
provenance must not trim frames or a model of stock `nn` modules collapses to one
path, while a source line must trim them or it names a line inside torch.

Instances, not types. `Block$1` and `Block$3` are different answers, and an
operation name cannot tell them apart: a model with twelve identical blocks has
twelve identical `linear` operations, and "the slow one is a linear" is not a
finding.

changes.py -- `compute_changes`, built on that rollup.

Three things stand between a correspondence and "7 operations changed, 6 of them
in Block$3, at model.py:41". Each is a measured defect in the naive version rather
than a refinement:

  * A difference lands on a value node as readily as on an operation. Reported
    as-is, a rewiring inside a callee body says nothing was modified. Resolved
    through `responsible_op`.
  * Most differing pairs are consequences of one edit, not edits: operands are
    other operations, so inserting one changes every neighbour's slots.
    `_is_independent_change` classifies each disagreeing slot by what happened to
    its endpoints and keeps only the two transitions nothing else reports -- a
    rewiring and a substitution. Chawathe et al. (1996) minimality.
  * Which instance of an interchangeable operation is left unmapped is arbitrary.
    Unaided, adding a third block reports its extra constant against `Block$1`,
    and a reader concludes something changed in a block that did not.
    `_relocate_to_new_modules` breaks the tie on module scope, which the matcher
    never looked at. Additions and removals only: a modification names a pair and
    is interchangeable with nothing, so re-seating one end is simply false.

plan_diff and intermediates -- the same rollup, so "which layer" is one call
whatever was measured.

`ComputePlanDiff.by_module()` groups placement moves; a device is not a place in
the model, and "Block$3/MLP$1 moved to CPU" is actionable where "6 operations
moved to CPU" is not. `IntermediatesReport.by_module()` groups divergences;
"NormBlock$1/Linear$1 diverges" is a fact about the model where "mul_1 diverges"
is a fact about the graph. Module paths arrive as arguments -- per side for
`plan_diff`, since op id numberings overlap and one merged mapping would attribute
an operation of one program to a module of the other, and at all for
`compare_intermediates`, so it stays the pure function of captured values its
docstring promises.

Both rollups are deliberately partial and say so. A plan diff groups moves, not
`only_before`/`only_after`: an operation that arrived or left did not *move*, that
is `compute_changes`' question, and reporting it twice would give one fact two
homes. An intermediates report groups failures, not every comparison: a real model
is mostly PASS and EXCLUDED, which would bury the handful that matter.

No inference fallback, on measurement.

A dataflow guess at an unattributed operation's module was in the design until the
coverage was measured: 87-99% of operations carry a path across the debugging
models, and the shortfall is exactly two per model regardless of size --
`coreai.graph` and `coreai.output`, both structure. Commit 6cb0a4c already fixed
the recorder to follow operand edges, which is what the guess was for, so it would
have been dead code. `utils.STRUCTURAL_OPS` filters those two instead, rather than
growing an "unattributed" bucket whose entire contents are plumbing.

Reports lead with counts, then the module tree, then the operations: a caller's
context is the scarce resource, and a 2,500-operation model must not arrive as
2,500 rows to answer "did anything change, and where". `to_dict` throughout, so
the JSON a programmatic caller reads comes from the same function that renders the
table.

Tests are constructed pairs with the answer written down, not comparisons against
what the code currently does -- a diff scored against itself enshrines its own
defects as the reference. Fixtures come from `test_model.py`, including a new
depth-parameterised `BlockStack`: the relocation case needs two programs sharing
every module path but one, and two different classes share none, so it cannot tell
"named the instance that changed" from "named the first of its kind".

Not done here: `BenchmarkResult.get_module_timings` keeps its own walk. It is the
one rollup with no user-visible gain from sharing, and `test_module_timings`
cannot run in this environment -- the benchmark returns zero operation timings
under both `interpreter` and `cpu` -- so the refactor could not be verified.

Suite: 228 passed, 1 skipped. The three failures
(`test_intermediates_torch_vs_coreai[TwoNormModel]`, `test_module_timings`,
`test_annotate_dominant_source`) are pre-existing on main, verified against a
clean tree. The first now localises itself, which is the point: it reports
`NormBlock$1/RMSNorm$1/RMSNormImpl$1` rather than "op 15".
@cymbalrush
cymbalrush merged commit a36b92e into apple:main Sep 24, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants