Skip to content

fix: Add support for truncate_long_and_double in FX #1865

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

Jump to bottom
Closed
wants to merge 2 commits into from
Closed

fix: Add support for truncate_long_and_double in FX #1865

wants to merge 2 commits into from

Conversation

gs-olive
Copy link
Collaborator

@gs-olive gs-olive commented Apr 27, 2023 •
edited
Loading

Description

  • Add utility capabilities for accepting int64 and float64 inputs to TRTModules to support multiple use cases
  • Support cases include situations where internal tensors in split modules are int64 (generally used for indexing torch Tensors)
  • This also supports cases where the user wants to input long or double tensors as forward inputs
  • Add test cases to verify functionality and accuracy
  • Enable tests for TRTModuleNext, which are now fully supported on main

Fixes #1864
Addresses #1740

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)

Checklist:

  • [ x ] My code follows the style guidelines of this project (You can use the linters)
  • [ x ] I have performed a self-review of my own code
  • [ x ] I have commented my code, particularly in hard-to-understand areas and hacks
  • [ x ] I have made corresponding changes to the documentation
  • [ x ] I have added tests to verify my fix or my feature
  • [ x ] New and existing unit tests pass locally with my changes
  • [ x ] I have added the relevant labels to my PR in so that relevant reviewers are notified

@gs-olive gs-olive requested a review from frank-wei April 27, 2023 21:28
@gs-olive gs-olive self-assigned this Apr 27, 2023
gs-olive
gs-olive commented Apr 27, 2023
View reviewed changes
py/torch_tensorrt/fx/test/core/test_trt_module.py
ref_output,
rtol=1e-04,
atol=1e-04,
check_dtype=False,
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The output data type will be different, since TRT cannot output int64 types

@github-actions github-actions bot requested a review from wushirong April 27, 2023 22:55
@gs-olive gs-olive changed the title fix: Add support for torch.int64 inputs in FX fix: Add support for truncate_long_and_double in FX Apr 28, 2023
@gs-olive gs-olive force-pushed the fx_int64_fix branch 2 times, most recently from 65bc360 to ad8aecf Compare April 28, 2023 22:34
gs-olive
gs-olive commented Apr 28, 2023
View reviewed changes
py/torch_tensorrt/fx/utils.py
Comment on lines +46 to +70
elif dtype == torch.int64:
if truncate_long_and_double:
_LOGGER.warn(
"Detected Int64 Input, Casting to Int32 for TRT Engine Compatibility"
)
return trt.int32
else:
raise TypeError(
"Detected Int64 Input which is not supported by tensorrt, enable compilation"
+ "option truncate_long_and_double=True to cast input to Int32 for TRT Engine"
)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly to the TorchScript path, allow the truncate_long_and_double argument to automatically cast inputs as needed by TRT Engines, while informing the user. This is primarily helpful for intermediate inputs (not user-provided), which happen to be long-type tensors (such as indices for embeddings).

@narendasan
Copy link
Collaborator

@gs-olive is this PR still needed?

@gs-olive
Copy link
Collaborator Author

Yes, this PR is still needed to support T5 in the torch_tensorrt.dynamo.compile path, and is also needed to support T5 in the FX aten path (#1740)

@narendasan
Copy link
Collaborator

Can we create a seperate PR for dynamo so we can land the feature there at least?

@gs-olive gs-olive mentioned this pull request May 22, 2023
Closed
gs-olive added 2 commits May 30, 2023 12:40
@gs-olive
fix: Add support for torch.int64 inputs in FXTRT
5382916 
- Add utility capabilities for accepting `int64` inputs to TRTModules to
support multiple use cases
- Support cases include situations where internal tensors in split
modules are `int64` (generally used for indexing torch Tensors)
- This also supports cases where the user wants to input `long` tensors
as `forward` inputs
- Add test cases to verify functionality and accuracy
- Enable tests for `TRTModuleNext`, which are now fully supported on
`main`
@gs-olive
fix: Add support for truncate_long_and_double in FX
6e53805 
- Add support and testing for `double` type inputs
@narendasan
Copy link
Collaborator

@gs-olive can you create separate PRs for each backend? Will be easier to merge then

@gs-olive gs-olive added the WIP Work is in progress, pull request should not be merged yet label May 30, 2023
@github-actions github-actions bot requested a review from yinghai May 30, 2023 21:57
@gs-olive gs-olive changed the title fix: Add support for truncate_long_and_double in FX fix: Add support for truncate_long_and_double in Dynamo May 31, 2023
@gs-olive gs-olive changed the title fix: Add support for truncate_long_and_double in Dynamo fix: Add support for truncate_long in Dynamo May 31, 2023
@gs-olive gs-olive changed the title fix: Add support for truncate_long in Dynamo fix: Add support for truncate_long_and_double in FX May 31, 2023
@gs-olive
Copy link
Collaborator Author

Closed in favor of the more robust #2021 (no need to manually downcast, have the FX graph/Dynamo utilities automatically handle this for us).

@gs-olive gs-olive closed this Jun 14, 2023
@gs-olive gs-olive deleted the fx_int64_fix branch June 14, 2023 19:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@borisfom borisfom borisfom approved these changes

@frank-wei frank-wei Awaiting requested review from frank-wei

@wushirong wushirong Awaiting requested review from wushirong

@yinghai yinghai Awaiting requested review from yinghai

Assignees

@gs-olive gs-olive

Labels
cla signed component: api [Python] Issues re: Python API component: fx fx WIP Work is in progress, pull request should not be merged yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

🐛 [Bug] Error when compiling aten model with intermediate int64 tensors
4 participants
@gs-olive @narendasan @borisfom @facebook-github-bot