fix(dynamo): correct legacy exporter (retrace=False) submodule inlining for hybrid graphs - #4446
fix(dynamo): correct legacy exporter (retrace=False) submodule inlining for hybrid graphs#4446Conarnar wants to merge 1 commit into
Conversation
|
Reviewed this one carefully because #4440 lands right on top of it: the composable The bug is real, and it is worse than mis-wiringI built a small parent graph whose submodule input placeholder name collides with an Your two unit tests are genuine regression tests, not decoration. I ran them both ways:
That is exactly the property that matters, and it is nice that they need neither a GPU On the approachWiring positionally from Letting Questions on the two compat fixesThese are separate from the inlining fix and I would like to understand them a bit 1. 2. The I tried to reach that path and could not, which is reassuring but leaves the question So the fallback did not fire for either, and in both cases the exported program returned That leaves two possibilities, and it would help to know which you saw. If the fallback Minor: the comment says "torch>=2.13" while the Composition with the ExecuTorch workFor anyone tracking how these fit together:
Thanks for tracking this down. The name-collision path has probably been quietly wrong |
d504244 to
60e39a4
Compare
|
Thank you for your suggestions. I have added clarifications in the comments.
Verified that |
…ng for hybrid graphs
torch_tensorrt.save(retrace=False) uses the legacy dynamo exporter, which inlines the
partitioned _run_on_gpu (non-TensorRT) submodules back into the graph before building an
ExportedProgram. For a hybrid graph interleaving TensorRT engines with a CUDA/pytorch
delegated op, inline_torch_modules wired each submodule's inputs by MATCHING placeholder
names to graph nodes (get_duplicate_nodes). Name matching binds an input to a same-named
but unrelated node on a collision (e.g. a submodule input placeholder name-matching a
different engine's getitem), which:
- rewires a consumer to the wrong producer and orphans the real one; the orphan is then
pruned by dead-code elimination, leaving a delegate short an output at runtime (an
aliased engine reports "expected N args, got N-1"); and
- for a submodule mixing graph-input and computed-intermediate inputs, leaks the
computed intermediates as spurious graph placeholders (misclassified USER_INPUTs).
Wire submodule inputs POSITIONALLY from the call_module args (gm_node.args, which is
authoritative) instead of by name: let graph_copy create a fresh placeholder for each
submodule input, then rewire each to submodule_inputs[i] by position and erase it. Drop
get_duplicate_nodes (now unused).
Also fix two torch-version-compat gaps this path hits on recent torch:
- lift(): pass an explicit persistent= flag on BUFFER InputSpecs (required since 2.3).
- create_trt_exp_program(): an inlined GraphModule may carry a plain fx.CodeGen (no
pytree_info); fall back to specs rebuilt from the example inputs + graph outputs.
With these, retrace=False export of a hybrid TensorRT+CUDA program is bit-identical to
retrace=True (validated on a 2-layer int4 MoE decode: per-step argmax + logits match).
Tests: tests/py/dynamo/models/test_exporter_inlining.py -- positional input wiring under a
name collision, and multi-output preservation (GPU-free fx unit tests).
60e39a4 to
464d756
Compare
torch_tensorrt.save(retrace=False) uses the legacy dynamo exporter, which inlines the partitioned _run_on_gpu (non-TensorRT) submodules back into the graph before building an ExportedProgram. For a hybrid graph interleaving TensorRT engines with a CUDA/pytorch delegated op, inline_torch_modules wired each submodule's inputs by MATCHING placeholder names to graph nodes (get_duplicate_nodes). Name matching binds an input to a same-named but unrelated node on a collision (e.g. a submodule input placeholder name-matching a different engine's getitem), which:
Wire submodule inputs POSITIONALLY from the call_module args (gm_node.args, which is authoritative) instead of by name: let graph_copy create a fresh placeholder for each submodule input, then rewire each to submodule_inputs[i] by position and erase it. Drop get_duplicate_nodes (now unused).
Also fix two torch-version-compat gaps this path hits on recent torch:
With these, retrace=False export of a hybrid TensorRT+CUDA program is bit-identical to retrace=True (validated on a 2-layer int4 MoE decode: per-step argmax + logits match).
Tests: tests/py/dynamo/models/test_exporter_inlining.py -- positional input wiring under a name collision, and multi-output preservation (GPU-free fx unit tests).