Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: Remove pytorch overhead while finding fusions for fully convertible models #3311

Merged
merged 7 commits into from
Dec 16, 2024
Merged

fix: Remove pytorch overhead while finding fusions for fully convertible models #3311

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
1d269e2
feat: Reduce partitioning overhead
peri044 Nov 28, 2024
4e96a91
chore: updates
peri044 Dec 2, 2024
1c7a638
fix: skip fusions for fully convertible models
peri044 Dec 4, 2024
6e275ba
chore: remove flux example
peri044 Dec 5, 2024
3697b8a
Merge branch 'main' into partitioning_opt
peri044 Dec 12, 2024
a5e3b5a
Merge branch 'main' into partitioning_opt
peri044 Dec 12, 2024
f4c8bef
chore: minor fix
peri044 Dec 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion py/torch_tensorrt/dynamo/_compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -765,7 +765,9 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
min_block_size=settings.min_block_size,
torch_executed_ops=settings.torch_executed_ops,
require_full_compilation=settings.require_full_compilation,
skip_fusion=(num_supported_ops == total_ops),
)

except torch.fx.passes.splitter_base.FxNetSplitterInternalError:
logger.error(
"Partitioning failed on the subgraph with fast partition. See trace above. "
Expand All @@ -786,7 +788,7 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
require_full_compilation=settings.require_full_compilation,
)

dryrun_tracker.unsupported_ops = supported_ops.unsupported_operators
# dryrun_tracker.unsupported_ops = supported_ops.unsupported_operators
peri044 marked this conversation as resolved.
Show resolved Hide resolved

# The global partitioner leaves non-TRT nodes as-is
if not settings.use_fast_partitioner:
Expand Down
5 changes: 5 additions & 0 deletions py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@ def __init__(
min_block_size: int = MIN_BLOCK_SIZE,
require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
return_tuple: bool = False,
skip_fusion: bool = False,
):
"""
Preprocesses graph before splitting:
Expand All @@ -127,6 +128,7 @@ def __init__(
self.settings = _SplitterSettingBase(
min_acc_module_size=min_block_size,
allow_non_tensor=True,
skip_fusion=skip_fusion,
)
self.operator_support = operator_support

Expand Down Expand Up @@ -252,6 +254,7 @@ def partition(
min_block_size: int = MIN_BLOCK_SIZE,
torch_executed_ops: Collection[Target] = set(),
require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
skip_fusion: bool = False,
) -> Tuple[torch.fx.GraphModule, OpSupportTester]:
"""Partition an FX GraphModule with aten ops into TRT engines
Partitioning is based on converter operator support
Expand All @@ -262,6 +265,7 @@ def partition(
min_block_size: Minimum number of operators per TRT-Engine Block
torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage
require_full_compilation: Require that all computational operators be run in TRT
skip_fusion: Skip fusions found by FxNetAccFusionsFinder
Returns:
torch.fx.GraphModule, OpSupportTester
"""
Expand All @@ -277,6 +281,7 @@ def partition(
supported_ops,
min_block_size=min_block_size,
require_full_compilation=require_full_compilation,
skip_fusion=skip_fusion,
)

partitioned_graph = partitioner.partition_graph()
Expand Down
Loading