fix: resolve ExecuTorch TRT target_device per partition (coalesced multi-engine)#4350
fix: resolve ExecuTorch TRT target_device per partition (coalesced multi-engine)#4350shoumikhin wants to merge 2 commits into
Conversation
github-actions
Bot
added
component: tests
shoumikhin
force-pushed
the
fix/executorch-per-partition-device
branch
2 times, most recently
from
380e295 to
6971383
Compare
|
The existing |
|
Good catch, fixed in the latest commit. Updated |
shoumikhin
force-pushed
the
fix/executorch-per-partition-device
branch
from
fd3c628 to
5933f20
Compare
TensorRTPartitioner resolved target_device once for the whole exported program via _get_engine_info_from_edge_program(), which requires exactly one engine node. A coalesced graph (TensorRT + CUDA delegates) has multiple TRT engines, so that call raised and every TRT partition fell back to cuda:0 with a spurious "expects exactly 1 engine node per partition, found N" warning; multi-GPU graphs also could not be labeled per partition. Extract _get_engine_info_for_node() (single-node engine-info extraction) from _get_engine_info_from_edge_program() and resolve target_device per partition from that partition's own engine node. Single-GPU behavior is unchanged (still cuda:0) minus the warning; multi-engine/multi-GPU graphs now label each delegate correctly.
The partitioner now derives target_device from each partition's own engine node (coalesced multi-engine), so the test must place the engine node in the partition and assert per-partition devices. Adds a coalesced multi-engine case and a >1-engine-per-partition fallback case.
shoumikhin
force-pushed
the
fix/executorch-per-partition-device
branch
from
5933f20 to
a5d4dbb
Compare
|
|
||
| partition_tags: Dict[str, DelegationSpec] = {} | ||
| for partition in partition_list: | ||
| tag = f"tensorrt_{partition.id}" |
There was a problem hiding this comment.
@narendasan Should the naming be consistent with dynamo's format like "run_on_acc" and "run_on_gpu"?
There was a problem hiding this comment.
Dont think it matters as much for the executorch partitioner. For the other partitioners we need to match on tag
3 participants
What's broken
When coalescing TensorRT with other delegates into one
.pte, the graph has multiple TensorRT engines.TensorRTPartitionerresolvedtarget_deviceonce for the whole program via_get_engine_info_from_edge_program(), which requires exactly one engine node. With more than one engine it raised, so every TensorRT partition fell back tocuda:0with a spurious warning:On a single GPU this is just noise, but a multi-GPU graph cannot label each delegate with its own device.
The fix
Extract
_get_engine_info_for_node()(single-node engine-info extraction) out of_get_engine_info_from_edge_program()(which keeps its one-engine contract used bypreprocess()), then resolvetarget_deviceper partition from that partition's own engine node.cuda:0), minus the spurious warning.cuda:0rather than aborting the export.Test
tests/py/dynamo/executorch/test_partitioner_target_device.pyandtests/py/dynamo/executorch/test_api.py: default GPU, non-zero GPU, a coalesced multi-engine graph resolving partition 0 →cuda:0and partition 1 →cuda:1, the >1-engine-in-one-partition fallback, a malformed-partition fallback, and an explicittarget_devicepin taking precedence..ptestill contains bothTensorRTBackendandAnotherBackend, and the "found 2 engines" warning no longer fires.Type of change
Checklist