From 3a4ac110f8526ae45cd4b442f316c489aafd1912 Mon Sep 17 00:00:00 2001 From: Po-Chen Wu Date: Wed, 30 Sep 2026 15:17:47 -0700 Subject: [PATCH] Fix constant_prop_pass fqn collision with state_dict (#23266) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/23266 `replace_with_constant_node` names propagated constants `_prop_tensor_constant{len(exported_program.constants)}`, then checks for a collision. Both the check and the max-suffix rescan look only at `exported_program.constants`. A program re-entering this pass can already carry a `_prop_tensor_constant*` in `exported_program.state_dict` from an earlier run. That name is invisible to both, so the pass reuses it and writes its tensor into `constants` alongside the stale `state_dict` entry. Emission then resolves the placeholder against `state_dict` before `constants` (`exir/emit/_emitter.py:2044`), picks up the stale tensor, and fails the size check at `_emitter.py:513`: InternalError: Tensor spec has buffer of size 48, but expected nbytes of 4 The spec is correct at 4 bytes; the tensor bound to it is the wrong one. Note the existing comment already reads "If prop_constant_tensor_fqn already exists in the state dict" -- the implementation never matched that intent. Widen the collision check and the suffix rescan to span `state_dict` as well as `constants`. Observed on a Turing/Coleman hand-tracking model compiled with `fx_options.allow_quant_dequant_on_tce=True`, which folds a scalar `1` from a `1 - valid_multi_view` expression whose name was already taken in `state_dict`. Instrumenting the pass showed the collision directly: fqn=_prop_tensor_constant20 new_nbytes=4 already_in_constants=False already_in_state_dict=True With the fix the model compiles to a .pte (3 TCE delegates, 84.19% delegation); without it emission aborts. Differential Revision: D122298673 --- exir/passes/constant_prop_pass.py | 10 ++++++++-- exir/tests/test_passes.py | 26 ++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index 19c25a838af..ea600e1c222 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -213,11 +213,17 @@ def replace_with_constant_node( # If prop_constant_tensor_fqn already exists in the state dict, we need # to create a new name. Find the largest suffix of "_prop_tensor_constant", # and increment it by 1 to form the new name. - if prop_constant_tensor_fqn in exported_program.constants: + # A program re-entering this pass can already carry a + # `_prop_tensor_constant*` in state_dict from an earlier run. Emission + # resolves placeholders against state_dict before constants, so reusing + # such a name lets the stale entry shadow the tensor written here. Both + # the check and the suffix scan therefore span state_dict too. + taken = exported_program.constants.keys() | exported_program.state_dict.keys() + if prop_constant_tensor_fqn in taken: suffix = 1 + max( ( int(name[len(prefix) :]) - for name in exported_program.constants.keys() + for name in taken if name.startswith(prefix) and name[len(prefix) :].isdigit() ), default=-1, diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 0f0caa3736b..0ebd71c1690 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -1690,6 +1690,32 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: new_ep.graph_module.code ) + def test_constant_prop_pass_avoids_state_dict_name_collision(self) -> None: + class Add(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 3 + + edge = to_edge( + export(Add(), (torch.ones(1),), strict=True), + compile_config=EdgeCompileConfig(_skip_dim_order=False), + ) + edge = edge.transform([ScalarToTensorPass(), RemoveMixedTypeOperators()]) + exported_program = lift_constant_tensor_pass(edge.exported_program()) + + # A program re-entering this pass can already carry a + # `_prop_tensor_constant*` in state_dict from an earlier run. + # Emission resolves placeholders against state_dict before + # constants, so reusing the name lets the stale entry shadow the + # propagated tensor and the two disagree on size. + stale = torch.zeros(12) + exported_program.state_dict["_prop_tensor_constant0"] = stale + + new_ep = constant_prop_pass(exported_program) + + for name in new_ep.constants: + self.assertNotIn(name, new_ep.state_dict) + self.assertIs(new_ep.state_dict["_prop_tensor_constant0"], stale) + def test_pass_no_user_inputs(self) -> None: class NoUserInputs(torch.nn.Module): def __init__(self):