Skip to content

Commit be879f4

Browse files
committed
Skip packing GraphModule.recompile() codegen bookkeeping as module state
torch.fx.GraphModule.recompile() writes `_code`/`_lineno_map` (and `_in_spec`/`_out_spec` for pytree codegen) straight into the module's __dict__. Thunder's attribute-modification tracking mistook these codegen-cache internals for real module state and tried to pack them into the epilogue trace, asserting when the value was a plain str/dict instead of a Proxy. Fixes test_higher_order_inplace_alias_update failing intermittently in CI depending on test order/dynamo cache state.
1 parent 5dd55d3 commit be879f4

1 file changed

Lines changed: 6 additions & 1 deletion

File tree

thunder/core/jit_ext.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2063,7 +2063,12 @@ def process_recorded_modifications(ctx, epilogue_trace):
20632063
):
20642064
name = k
20652065
setattr_obj_provenance = modified_object.provenance.inputs[0]
2066-
if hasattr(setattr_obj_provenance, "proxy"):
2066+
# GraphModule.recompile() codegen bookkeeping, not real module state.
2067+
if name in ("_code", "_lineno_map", "_in_spec", "_out_spec") and isinstance(
2068+
umodified_object.get("_graph"), torch.fx.Graph
2069+
):
2070+
pass
2071+
elif hasattr(setattr_obj_provenance, "proxy"):
20672072
assert isinstance(
20682073
value.value, (Proxy, int, float, tuple, NoneType, thunder.devices.Device)
20692074
), (

0 commit comments

Comments
 (0)