Repository navigation
Conversation
|
|
|
The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:
Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows. |
004eb2e to
6072f70
Compare
6072f70 to
64ca0f9
Compare
|
@claude review these changes |
|
Claude finished @jansel's task in 2m 16s —— View job The fix is correct for the paths it touches, but it is partial: the Full reviewPR Review: #198124SummaryThe PR installs Infrastructure
Performance
Testing
Code Quality
RecommendationNeeds Discussion. The key question for dynamo maintainers is whether to accept per-method |
Review follow-up for pytorch#198124. forward. UnspecializedNNModuleVariable.call_function short-circuits _call_impl to forward and built it with a plain UserFunctionVariable, explicitly to skip the forward.__code__ guard; the rationale was the ~20% guard overhead from pytorch#124975. That left `model(x)` reusing a stale graph after `Model.forward.__code__ = ...` or `Model.forward = other`, while a helper that forward calls was guarded, which is the inconsistency the review pointed out. The short-circuit now uses create_with_source on the type(mod).__dict__ source when the target is forward. This is not the pytorch#124975 trade-off: that change put an ID_MATCH on every inlined function, whereas this is one CLOSURE_MATCH per (class, method), deduped by the get_source_by_walking_mro cache, so a model built from N copies of a block pays once per block class, not per instance. Instance-level patching is still handled by guard_to_detect_forward_monkeypatching, whose comment now says where class-level patching is caught. classmethod. Main reworked ClassMethodVariable while this was in review: from_descriptor now builds the wrapped function with VariableTracker.build on the __func__ source, which installs the code guard itself. The merge takes main's version, and test_classmethod_code_mutation_recompiles passes on main without this PR, so that half of the original change is gone. Other sites. object_protocol.py's FunctionType branch sources the function as obj.<name>.__func__, a GetAttr chain off the instance, which is the tag-unsafe shape this change already leaves unguarded on the TypeSource fallback, so it is left alone. The dunder helpers that use AttrSource(TypeSource(...)) are the same case. Cleanup: the resolve_type_attr comment is down to two lines, and the TypeSource fallback source is built inline in the else branch, where it is the only use. Known limit: with the opt-in use_recursive_dict_tags_for_guards, a tag-safe root skips its subtree while its dict tag is unchanged, so an in-place __code__ swap is still missed there. That applies to every CLOSURE_MATCH under such a root, predates this change, and the flag is off by default. Tests: - test_nn_module_method_mutation_recompiles (4): forward vs. a helper it calls, __code__ replacement vs. rebinding, through `mod(x)`. - test_nn_module_tag_safe_with_method_code_guard: a module reaching both guarded paths stays a tag-safe root. Test Plan: Run against the 2.15.0.dev20261003 CPU nightly (which matches origin/main for every file touched here) with this branch's torch/_dynamo files laid over it: ``` python test/dynamo/test_recompiles.py python test/dynamo/test_guard_manager.py ``` All pass. Without the change, the four existing class-function mutation tests and the four new nn.Module ones fail with a stale result. ``` pytest test/dynamo/test_modules.py test/dynamo/test_functions.py test/dynamo/test_guard_exclusion.py test/dynamo/test_autograd_function.py ``` 943 passed, 3 failed: test_torch_function_metadata_attrs_constant, test_inject_module_parameters and test_save_and_load_inductor. All three fail the same way on the unpatched nightly (Windows environment), so none is from this change. This change was authored with the assistance of an AI coding agent.
|
@williamwen42 @jansel I have addressed Claude's reviews. Please review when you have the chance, thanks! |
Fixes pytorch#197860. Replacing a method's code object after compilation (`Model.m.__code__ = other.__code__`) changed eager execution but left `torch.compile` running the graph traced from the old body, with no recompile, warning or graph break. `CLOSURE_MATCH`, the guard that pins a function's `__code__`, is installed by `VariableBuilder` or by `UserFunctionVariable.create_with_source`. A function found on the type MRO takes neither route: `UserDefinedObjectVariable.resolve_type_attr` and `ClassMethodVariable.tp_descr_get_impl` both build the bound method's `im_func` with the plain `UserFunctionVariable(...)` constructor, so no code guard is installed. (`StaticMethodVariable.tp_descr_get_impl` goes through `VariableTracker.build` and is guarded, which is why only the other two are affected.) Both sites now use `create_with_source` when the function's source is rooted at the class `__dict__`: the MRO-walk source in `resolve_type_attr`, and the descriptor's own `__func__` in `ClassMethodVariable`. The fallback sources (`TypeSource(obj).name`, `owner.name.__func__`) stay unguarded. A guard on them would hang a `GetAttrGuardAccessor` off the type's guard node, which the tag-safety pass rejects, so every `nn.Module` reaching that path would lose recursive dict-tag optimization (an earlier revision regressed `test_nn_module_tag_overridden_getattr_safe` this way). The class-`__dict__` chain is tag-safe: `DictGetItemGuardAccessor`, `CodeGuardAccessor` and the `__func__` accessor are all accepted, and code objects are immutable. This also fixes related cases that were broken in the same way: rebinding the class attribute (`Model.m = other`), both mutations on an inherited method, and classmethod code mutation via an instance or the class. Cost: three guard-tree nodes per distinct (class, method) pair, deduped by the existing cache in `get_source_by_walking_mro`. On an earlier revision with an equivalent guard, a 20-block `nn.Module` with four methods per block went from 220 to 232 guard nodes and guard eval from 199us to 215us. ``` python test/dynamo/test_recompiles.py ``` Nine new tests in `test/dynamo/test_recompiles.py`: - `test_class_function_attr_mutation_recompiles` (4): `__code__` replacement vs. rebinding, on own class vs. inherited (`type.__mro__[1].__dict__` source). - `test_classmethod_code_mutation_recompiles` (4): instance vs. class access, own class vs. inherited. - `test_unmutated_method_does_not_recompile`: control; two instances of an unmutated class share one cache entry. Local verification is limited. The patched code is newer than any torch build available to me (Windows, 2026-08-12 nightly), and the change can't be backported onto that wheel because the refactor it builds on spans modules. On that nightly, the eight mutation tests fail and the control passes without the fix, and an equivalent fix on the pre-refactor code made all nine pass without changing the failure sets of test_autograd_function, test_guard_manager, test_guard_exclusion, test_modules, test_subclasses or test_functions. The rebased change here is compiled and ruff-clean but not executed locally; CI is the check. This change was authored with the assistance of an AI coding agent.
Review follow-up for pytorch#198124. forward. UnspecializedNNModuleVariable.call_function short-circuits _call_impl to forward and built it with a plain UserFunctionVariable, explicitly to skip the forward.__code__ guard; the rationale was the ~20% guard overhead from pytorch#124975. That left `model(x)` reusing a stale graph after `Model.forward.__code__ = ...` or `Model.forward = other`, while a helper that forward calls was guarded, which is the inconsistency the review pointed out. The short-circuit now uses create_with_source on the type(mod).__dict__ source when the target is forward. This is not the pytorch#124975 trade-off: that change put an ID_MATCH on every inlined function, whereas this is one CLOSURE_MATCH per (class, method), deduped by the get_source_by_walking_mro cache, so a model built from N copies of a block pays once per block class, not per instance. Instance-level patching is still handled by guard_to_detect_forward_monkeypatching, whose comment now says where class-level patching is caught. classmethod. Main reworked ClassMethodVariable while this was in review: from_descriptor now builds the wrapped function with VariableTracker.build on the __func__ source, which installs the code guard itself. The merge takes main's version, and test_classmethod_code_mutation_recompiles passes on main without this PR, so that half of the original change is gone. Other sites. object_protocol.py's FunctionType branch sources the function as obj.<name>.__func__, a GetAttr chain off the instance, which is the tag-unsafe shape this change already leaves unguarded on the TypeSource fallback, so it is left alone. The dunder helpers that use AttrSource(TypeSource(...)) are the same case. Cleanup: the resolve_type_attr comment is down to two lines, and the TypeSource fallback source is built inline in the else branch, where it is the only use. Known limit: with the opt-in use_recursive_dict_tags_for_guards, a tag-safe root skips its subtree while its dict tag is unchanged, so an in-place __code__ swap is still missed there. That applies to every CLOSURE_MATCH under such a root, predates this change, and the flag is off by default. Tests: - test_nn_module_method_mutation_recompiles (4): forward vs. a helper it calls, __code__ replacement vs. rebinding, through `mod(x)`. - test_nn_module_tag_safe_with_method_code_guard: a module reaching both guarded paths stays a tag-safe root. Test Plan: Run against the 2.15.0.dev20261003 CPU nightly (which matches origin/main for every file touched here) with this branch's torch/_dynamo files laid over it: ``` python test/dynamo/test_recompiles.py python test/dynamo/test_guard_manager.py ``` All pass. Without the change, the four existing class-function mutation tests and the four new nn.Module ones fail with a stale result. ``` pytest test/dynamo/test_modules.py test/dynamo/test_functions.py test/dynamo/test_guard_exclusion.py test/dynamo/test_autograd_function.py ``` 943 passed, 3 failed: test_torch_function_metadata_attrs_constant, test_inject_module_parameters and test_save_and_load_inductor. All three fail the same way on the unpatched nightly (Windows environment), so none is from this change. This change was authored with the assistance of an AI coding agent.
Review follow-up for pytorch#198124. CI on aef3192 failed test_aot_compile_module (PackageError: CLOSURE_MATCH guard cannot be serialized) and, on six shards, test_precompile_session_captures_variants_and_summarizes (one risky dropped guard). The code guard this PR adds on type(obj).__dict__[name] is local-rooted, so skip_guard_on_globals_unsafe keeps it, and at that merge base serialize_guards refused every CLOSURE_MATCH; precompile's default filter dropped it and, with the slot holding different functions across variants, flagged the drop as risky. The automated review pointed out the same exposure for any method call on a user object through resolve_type_attr. Main has since fixed this at the shared point. pytorch#198810 rewrites a CLOSURE_MATCH on a plain function into FUNCTION_CODE_MATCH (code object plus the globals' module, compared by value) when guards are saved, and is_portable_function_guard lets precompile's default filter keep it. That covers the forward guard and the resolve_type_attr guard alike, so this branch is rebased onto main (taking main's ClassMethodVariable in the functions.py conflict, as the earlier merge did) and needs no source change. A method whose globals no importable module owns (e.g. a class built by exec) is still refused, which is main's policy for every CLOSURE_MATCH. New test, test_aot_compile_user_method_code_guard: aot_compile a function calling a method on a user object with skip_guard_on_globals_unsafe, save, load, check the result, then replace the method's __code__ and check that the loaded artifact's guard fails. Without this PR the last check fails (nothing guards the code); at the old merge base the compile itself raises the PackageError. Test Plan: Linux (WSL Ubuntu 22.04), Python 3.12, 2.16.0.dev20261016 CPU nightly (built from 99096b2, 41 commits behind origin/main), with every torch/*.py that differs between that commit and this branch laid over it. Reproduction on the old base: 2.15.0.dev20261004 nightly (e8f41c0 + 3 non-dynamo commits) with aef3192's nn_module.py and user_defined.py: both CI tests fail as reported (PackageError; risky_dropped_guards 1 != 0), and pass without the two files. A user-object method aot_compile with skip_guard_on_globals_unsafe raises the same PackageError with them and compiles without them. After rebase: ``` pytest test/dynamo/test_recompiles.py test/dynamo/test_guard_manager.py test/dynamo/test_aot_compile.py test/dynamo/test_precompile_package.py test/dynamo/test_modules.py test/dynamo/test_functions.py test/dynamo/test_guard_exclusion.py test/dynamo/test_autograd_function.py test/dynamo/test_guard_serialization.py test/dynamo/test_package.py ``` 1889 passed, 1 failed: test_nested_function_preserves_a_guarded_defaults_tuple, which fails the same way with origin/main's torch files, so it is not from this change. With origin/main's torch files the same run also fails the eight mutation tests in test_recompiles.py and the new aot test. This change was authored with the assistance of an AI coding agent.
aef3192 to
0ed727f
Compare
There was a problem hiding this comment.
Automated review of 0ed727f: nothing blocking; ready for human review.
The PR adds a code guard on class-level methods resolved through the cls.dict MRO source, and on the short-circuited nn.Module forward, so patching a class method now triggers a recompile. The source changes look correct and tag-safe, and regression tests that fail without the fix are included. The only maintainer comment was a bot-review trigger, which the author followed up on. Two minor issues remain.
⚪ 2 non-blocking findings
pytorch/test/dynamo/test_recompiles.py
Line 1040 in 0ed727f
This test does not exercise the change. Classmethods go through ClassMethodVariable.from_descriptor -> VariableTracker.build on func, which already installed CLOSURE_MATCH before this PR, so the test passes without the fix. Either drop it or say it only guards existing behavior.
pytorch/torch/_dynamo/variables/nn_module.py
Line 189 in 0ed727f
The comment claims class-level forward patching is now caught. That holds only for UnspecializedNNModuleVariable. The specialized NNModuleVariable.var_getattr path also calls this helper but still builds forward with a plain UserFunctionVariable and no code guard.
This review is AI-generated and advisory. Each new automated review replaces the previous one.
Review follow-up for pytorch#198124 (automated review of 0ed727f, two minor findings). test_classmethod_code_mutation_recompiles did not exercise this change. Classmethods go through ClassMethodVariable, which on main already guards __func__ with CLOSURE_MATCH; all four variants pass with origin/main's torch files, while the other eight mutation tests fail there. Main already tests classmethod __func__.__code__ mutation in test_autograd_function.py, so the test is removed rather than kept as a pin on existing behavior. The comment in guard_to_detect_forward_monkeypatching said class-level patching of forward is caught. That only holds for UnspecializedNNModuleVariable. The specialized NNModuleVariable path (e.g. a module with torchdynamo_force_dynamic = False) calls the same helper but builds forward without a code guard: replacing Mod.forward.__code__ or rebinding Mod.forward after compiling returns the stale result with no recompile, on main and on this branch alike. The comment now says so. That path is only reached for modules Dynamo still specializes (strict export, torchdynamo_force_dynamic = False); guarding it would need its own source choice and tests, so it is left as a known gap instead of widening this PR. Test Plan: Linux (WSL Ubuntu 22.04), Python 3.12, 2.16.0.dev20261016 CPU nightly (built from 99096b2), with every torch/*.py that differs between that commit and the tested tree laid over it. ``` pytest test/dynamo/test_recompiles.py -k "classmethod_code_mutation or class_function_attr_mutation or nn_module_method_mutation" ``` With origin/main's torch files (at 0ed727f): 8 failed, 4 passed; the 4 passing are the classmethod variants. Specialized-path probe (Mod with torchdynamo_force_dynamic = False, forward = two, compiled via lambda x: mod(x), then Mod.forward.__code__ replaced or Mod.forward rebound): compiled output stays x * 2 with frame_count 1 on both origin/main and this branch, eager gives x * 9. ``` ruff check torch/_dynamo/variables/nn_module.py test/dynamo/test_recompiles.py ruff format --check torch/_dynamo/variables/nn_module.py test/dynamo/test_recompiles.py pytest test/dynamo/test_recompiles.py test/dynamo/test_guard_manager.py test/dynamo/test_aot_compile.py test/dynamo/test_precompile_package.py test/dynamo/test_modules.py ``` ruff: clean. pytest: 779 passed, 9 skipped, 28 subtests passed. This change was authored with the assistance of an AI coding agent.
The previous commit installed CLOSURE_MATCH on every MRO-sourced method, including methods of objects that already have a source. Upstream sources those through the class MRO too but does not guard them (pytorch#197860, which pytorch#198124 addresses), and precompile now relies on that: since pytorch#200823, test_package.py's unnamed-scope tests call an exec()-defined method of a module-level object, and the new guard on it cannot be made portable, so saving the package raised "CLOSURE_MATCH guard cannot be serialized". Keep the guard for the sources this PR introduces - methods of objects constructed during tracing, which are reachable only through cls_source - and give objects that have a source back exactly the function source they had before.

Fixes #197860. Replacing a class method's code (
Model.m.__code__ = ...) or rebinding it (Model.m = other) after compilation lefttorch.compilerunning the stale graph. Same fornn.Module.forward.Fix
resolve_type_attrand theUnspecializedNNModuleVariableforwardshort-circuit built the method with a plainUserFunctionVariable, so noCLOSURE_MATCHwas installed. Both now usecreate_with_sourceon the class__dict__source.TypeSourcefallback andobj.<name>.__func__sources, since a GetAttr accessor there breaks the module's tag-safe root.ClassMethodVariablerefactor already guards these; the merge takes main's version.Tests
test_recompiles.py: class-function mutation (4),nn.Moduleforward/helper mutation (4), no-recompile control.test_guard_manager.py: module stays a tag-safe root with the new guards.test_aot_compile.py: the method code guard survivesaot_compilesave/load and rejects a swapped__code__.Verified locally against the
2.15.0.dev20261003+cpunightly (main + this PR vs. plain main) ontest_recompiles,test_guard_manager,test_modules,test_functions,test_guard_exclusion,test_autograd_function,test_misc,test_repros,test_unspec,test_model_output: no new failures; the 8 mutation tests go from failing to passing.This change was authored with the assistance of an AI coding agent.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @chauhang @amjames @jataylo @azahed98