Skip to content

Enable Dynamic Shapes For MD Ops #4402

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

Open
jloftin-nv wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Enable Dynamic Shapes For MD Ops #4402

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
24 changes: 18 additions & 6 deletions py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,9 @@
_LOGGER.info("Using native TensorRT DistCollective API for distributed operations")

@dynamo_tensorrt_converter(
tensorrt_fused_nccl_all_gather_op, requires_native_multidevice=True
tensorrt_fused_nccl_all_gather_op,
requires_native_multidevice=True,
supports_dynamic_shapes=True,
)
def fused_nccl_all_gather(
ctx: ConversionContext,
Expand All @@ -47,7 +49,9 @@ def fused_nccl_all_gather(
)

@dynamo_tensorrt_converter(
tensorrt_fused_nccl_reduce_scatter_op, requires_native_multidevice=True
tensorrt_fused_nccl_reduce_scatter_op,
requires_native_multidevice=True,
supports_dynamic_shapes=True,
)
def fused_nccl_reduce_scatter(
ctx: ConversionContext,
Expand All @@ -69,7 +73,9 @@ def fused_nccl_reduce_scatter(
)

@dynamo_tensorrt_converter(
tensorrt_fused_nccl_all_reduce_op, requires_native_multidevice=True
tensorrt_fused_nccl_all_reduce_op,
requires_native_multidevice=True,
supports_dynamic_shapes=True,
)
def fused_nccl_all_reduce(
ctx: ConversionContext,
Expand All @@ -90,7 +96,9 @@ def fused_nccl_all_reduce(
)

@dynamo_tensorrt_converter(
tensorrt_fused_nccl_all_to_all_op, requires_native_multidevice=True
tensorrt_fused_nccl_all_to_all_op,
requires_native_multidevice=True,
supports_dynamic_shapes=True,
)
def fused_nccl_all_to_all(
ctx: ConversionContext,
Expand All @@ -109,7 +117,9 @@ def fused_nccl_all_to_all(
)

@dynamo_tensorrt_converter(
tensorrt_fused_nccl_scatter_op, requires_native_multidevice=True
tensorrt_fused_nccl_scatter_op,
requires_native_multidevice=True,
supports_dynamic_shapes=True,
)
def fused_nccl_scatter(
ctx: ConversionContext,
Expand All @@ -125,7 +135,9 @@ def fused_nccl_scatter(
)

@dynamo_tensorrt_converter(
tensorrt_fused_nccl_gather_op, requires_native_multidevice=True
tensorrt_fused_nccl_gather_op,
requires_native_multidevice=True,
supports_dynamic_shapes=True,
)
def fused_nccl_gather(
ctx: ConversionContext,
Expand Down
Loading