[fx] Clear fake_mode on a view's base in GraphPickler (#194771) (#194771)
Summary: `_TensorPickleData.__init__` clears `fake_mode` on the top-level `MetaTensorDesc` but not on `metadata.base`. For a **view** tensor the base descriptor carries its own live `fake_mode`, so pickling reaches `FakeTensorMode` -> `FakeTensorConverter` -> `MetaConverter.storage_memo` and dies on a memoized storage. The asymmetry is the bug: `unpickle` already restores `fake_mode` on *both* the descriptor and its base, so the dump side clearing only one level was an oversight rather than a deliberate choice. This clears the base too, gated on the same `is_view and base is not None` condition `unpickle` uses, so a non-view carrying a base can never be dumped cleared and left unrestored. Found via a graph whose `aten.permute` nodes survive into the collected FX graph; `GraphPickler.debug_dumps` reported the leaf as: ``` ...nodes[5].meta['val'].metadata.base.fake_mode .fake_tensor_converter.meta_converter.storage_memo._remove ``` The concrete symptom was `RuntimeError: Cannot access data pointer of Tensor (e.g. FakeTensor, FunctionalTensor)`, though the exact failing leaf varies with what the mode's converter has memoized. This change was authored with assistance from Claude Code. The root cause analysis, fix, and tests were reviewed by the author. Test Plan: Added `TestViewTensorPickle` to `test/fx/test_graph_pickler.py` with two tests: `test_base_fake_mode_is_cleared` pins the descriptor invariant directly, and `test_view_val_roundtrips` pins that a view `FakeTensor` in `node.meta['val']` survives `dumps`/`loads` with its view-ness, shape, stride and dtype intact. The round-trip test builds its fake tensor via `fake_mode.from_tensor` so the mode's `storage_memo` is populated, and asserts that precondition -- without it the leak has nothing to drag in and the test passes vacuously. Both tests fail without the fix and pass with it. Note `test/fx/test_graph_pickler.py` is excluded from the fbcode build at `caffe2/test/BUCK:203`, so they were verified by temporarily hosting the two test classes in `caffe2/test/inductor:compiled_fx_graph_serialization`: ```bash buck2 run fbcode//caffe2/test/inductor:compiled_fx_graph_serialization ``` Without the fix: `FAILED (failures=1, errors=1)` -- `test_base_fake_mode_is_cleared` fails on `FakeTensorMode ... is not None`, `test_view_val_roundtrips` errors inside `GraphPickler.dumps`. With the fix: `Ran 5 tests ... OK`. Existing GraphPickler coverage, unaffected: ```bash buck2 run fbcode//caffe2/test/inductor:compiled_fx_graph_serialization # 3/3 OK buck2 run fbcode//mtia/accuracy/training/tests:test_graph_serialize # 9/9 OK ``` `caffe2/test/dynamo:test_dynamo` fails at collection with `ModuleNotFoundError: test_functions` in `test_ctx_manager.py` and `test_nested_graph_breaks_wrapped.py`. That reproduces without this change and is unrelated to it. Differential Revision: D117389490 Pull Request resolved: https://github.com/pytorch/pytorch/pull/194771 Approved by: https://github.com/bobrenjc93 Co-authored-by: Michael Roberts <mbr@meta.com>
M
Michael Roberts committed
550ed938e9c427834ec1b9a3e88506e8acc095a1
Parent: d397b88
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 8/26/2026, 4:05:32 AM