From a8565dc0695de6becd4ce9e1ac55d9e2ff23da11 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Fri, 28 Aug 2026 12:58:21 -0700 Subject: [PATCH] fix: order the placeholders first before computing serde range constraints torch_tensorrt.save() raises IndexError: list index out of range whenever the partitioner has sent ops back to PyTorch. It is not architecture-specific -- any caller that hits fallback reaches it on any GPU: trt_gm = torchtrt.dynamo.compile( exp_program, inputs=[...], min_block_size=1, torch_executed_ops={"torch.ops.aten.linear.default"}, ) torchtrt.save(trt_gm, path, inputs=[Input(shape=(4, 10))], retrace=False) IndexError: list index out of range torch_tensorrt/dynamo/_exporter.py:491 range_constraints = make_constraints(...) torch/_export/non_strict_utils.py:935 shape_spec = flat_dynamic_shapes[...] save(retrace=True, use_legacy_exporter=True) with dynamic shapes fails the same way; both options land in create_trt_exp_program(). make_constraints() walks enumerate(gm.graph.nodes) and indexes flat_dynamic_shapes with the *node* index rather than the placeholder ordinal, so it silently requires the placeholders to be the leading nodes of the graph. Every graph torch.export produces satisfies that. A graph whose ops fell back to PyTorch does not: unlifting reinstates the module parameters as get_attr nodes ahead of the user inputs, so a lone input at node index 2 reads flat_dynamic_shapes[2] on a one-element list. Measured on an L40S with the same model compiled twice: converted fallback [0] placeholder x [0] get_attr linear_weight [1] get_attr ..._engine [1] get_attr linear_bias [2] call_function execute_engine [2] placeholder x [3] call_function getitem [3] call_function aten.linear.default [4] output [4] output save -> OK save -> IndexError Both have one placeholder, dynamic_shapes {'x': {}} and a one-element flat_dynamic_shapes; only the placeholder's node position differs. That is also why the failure surfaces as a bare IndexError: make_constraints()' own length check, len(flat_dynamic_shapes) == num_placeholders - num_lifted_inputs, passes. create_trt_exp_program() now normalizes its graph to the placeholders-first shape torch.export emits before computing range_constraints. Placeholders take no arguments, so hoisting them cannot break the topological order, and preserving their relative order leaves the forward signature, in_spec and the InputSpec ordering unchanged; the helper returns immediately when the placeholders already lead, so a fully converted graph is untouched. Fixing this in torch instead -- indexing by placeholder ordinal -- would be the better repair, but torch._export.non_strict_utils is a private API and the ordering contract, though unstated, is the caller's to meet. Not fixed in transform(): the non-legacy exporter retraces the inlined module through torch.export.export() and is unaffected, and transform() is also reachable from refit and executorch paths that should not change. The regression test needs neither a GPU nor a TensorRT build: an exported module reproduces the get_attr-before-placeholder node order exactly, so create_trt_exp_program() can be driven directly with a Dim spec. The one-word "mis-wired" -> "miswired" edit in the same file is unrelated: it is a pre-existing typo the repo's own typos hook rejects, which blocks any commit touching this file. Testing (T4 ipp1-2023 and L40S a1u1g-mil-0572, driver 595.58.03, identical stacks, -n 1, full models/ runs on both arms, test_hf_gqa_model.py ignored): * The standalone reproducer above prints "save -> OK" for both the converted and the fallback row on both arms; before, the fallback row raised IndexError on both. * models/ on the T4: 266 collected, 222 passed / 19 failed / 25 skipped -> 228 passed / 13 failed / 25 skipped. Two of the six recovered tests are this fix (test_save_load_input_objects_retrace_false and test_save_load_retrace_true_legacy_true_dynamic); neutralising the fix at this tip leaves the other four passing, so they belong to commits that landed after the baseline was taken. * models/ on the L40S: 266 collected, no pre-existing test changes status. The only difference on either arm is the new test appearing as a pass. * The 3 view_as_real refit failures and the 10 remaining SM 7.5 failures are untouched, and no test regressed on either arm. Co-Authored-By: Claude Opus 5 --- py/torch_tensorrt/dynamo/_exporter.py | 37 ++++++++++++ .../dynamo/models/test_exporter_inlining.py | 60 ++++++++++++++++++- 2 files changed, 96 insertions(+), 1 deletion(-) diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index c6c992d720..34c5b5b838 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -386,6 +386,41 @@ def copy_submodule_attributes( _assign_attr(value, gm, key, _AttrKind.BUFFER) +def _order_placeholders_first(gm: torch.fx.GraphModule) -> None: + """Hoist every placeholder to the front of the node list, keeping their + relative order. + + ``make_constraints()`` indexes ``flat_dynamic_shapes`` by *node* position while + walking ``enumerate(gm.graph.nodes)``, so it silently requires placeholders to lead + the graph. For ``nn.Linear(10, 5)`` with one dynamic input:: + + converted (placeholders lead) fully fallen back to PyTorch + [0] placeholder x [0] get_attr linear_weight + [1] get_attr engine_0 [1] get_attr linear_bias + [2] call_function ... [2] placeholder x + [3] output [3] call_function linear + [4] output + + Unlifting puts the parameters ahead of the user input, so on the right the lone + input sits at node index 2 and the lookup becomes ``flat_dynamic_shapes[2]`` on a + one-element list: ``save()`` dies with a bare ``IndexError`` while its own length + check passes. Placeholders take no arguments, so hoisting cannot break topological + order, and preserving their relative order leaves the forward signature, + ``input_nodes``, ``in_spec`` and InputSpec ordering unchanged. + """ + nodes = list(gm.graph.nodes) + placeholders = [node for node in nodes if node.op == "placeholder"] + if not placeholders or nodes[: len(placeholders)] == placeholders: + return + + for node in reversed(placeholders): + first = next(iter(gm.graph.nodes)) + if node is not first: + first.prepend(node) + + gm.graph.lint() + + def create_trt_exp_program( gm: torch.fx.GraphModule, *, @@ -397,6 +432,8 @@ def create_trt_exp_program( and constructs an Exported Program object with the new IO node names and state_dict """ + _order_placeholders_first(gm) + input_nodes = [node for node in gm.graph.nodes if node.op == "placeholder"] output_nodes = [node for node in gm.graph.nodes if node.op == "output"] assert output_nodes diff --git a/tests/py/dynamo/models/test_exporter_inlining.py b/tests/py/dynamo/models/test_exporter_inlining.py index 2fa2f0470e..9dad5de16b 100644 --- a/tests/py/dynamo/models/test_exporter_inlining.py +++ b/tests/py/dynamo/models/test_exporter_inlining.py @@ -71,7 +71,7 @@ def test_inline_torch_modules_wires_inputs_by_position(): @pytest.mark.unit def test_inline_torch_modules_preserves_all_submodule_outputs(): """A multi-output _run_on_gpu submodule must keep every output wired to its - consumer after inlining. Regression: a mis-wired input orphaned one submodule + consumer after inlining. Regression: a miswired input orphaned one submodule output, which dead-code elimination then pruned, leaving a downstream consumer (or, in the hybrid case, a TensorRT engine) short an output at runtime. @@ -231,6 +231,64 @@ def forward(self, a, b): assert torch.allclose(out, torch.tensor(7.0)) +@pytest.mark.unit +def test_create_trt_exp_program_handles_get_attrs_before_placeholders(): + """create_trt_exp_program must build range_constraints on a graph whose + get_attr nodes precede its user input. + + Regression: when the partitioner sends ops back to PyTorch the parameters + come back as get_attr nodes AHEAD of the user inputs. torch's + make_constraints() indexes flat_dynamic_shapes by NODE position, not by + placeholder ordinal, so the lone input sitting at node index 2 made it read + flat_dynamic_shapes[2] on a one-element list and torch_tensorrt.save() died + with a bare "IndexError: list index out of range". Its own length check + passes, so nothing pointed at the real problem. An exported module reproduces + the node order exactly. + """ + + class Linear(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + return self.linear(x) + + batch = torch.export.Dim("batch", min=2, max=8) + gm = torch.export.export( + Linear().eval(), (torch.randn(4, 10),), dynamic_shapes={"x": {0: batch}} + ).module() + # A compiled TRT GraphModule carries no _guards_fn; drop it so the graph + # matches what create_trt_exp_program is handed in the normal flow. + for node in list(gm.graph.nodes): + if node.op == "call_module" and node.target == "_guards_fn": + gm.graph.erase_node(node) + break + gm.graph.lint() + gm.recompile() + + nodes = list(gm.graph.nodes) + assert nodes[0].op == "get_attr" + assert [n.op for n in nodes].index("placeholder") > 0 + + ep = create_trt_exp_program( + gm, arg_inputs=(torch.randn(4, 10),), dynamic_shapes={"x": {0: batch}} + ) + + # The user-specified Dim bound survived, so path 1 (make_constraints) ran. + assert [str(vr) for vr in ep.range_constraints.values()] == ["VR[2, 8]"] + assert [s.kind.name for s in ep.graph_signature.input_specs] == [ + "PARAMETER", + "PARAMETER", + "USER_INPUT", + ] + + # And the program still runs at another batch size within the range. + out = ep.module()(torch.randn(3, 10)) + out = out[0] if isinstance(out, (tuple, list)) else out + assert out.shape == (3, 5) + + @pytest.mark.unit def test_lift_sets_persistent_true_on_buffer_spec(): """lift() must mark a lifted BUFFER InputSpec as persistent.