-
Notifications
You must be signed in to change notification settings - Fork 375
feat: improve engine caching and fix bugs #3932
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.
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
base: main
Are you sure you want to change the base?
Conversation
|
@cehongwang please take a pass so we have multiple eyes on this PR |
a54907e to
ea81677
Compare
| logger.warning( | ||
| "require_full_compilation arg is not applicable for torch.compile with backend='torch_tensorrt" | ||
| ) | ||
| if settings.strip_engine_weights: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When would a torch.compile use try to use strip weights?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
added the warning back. Not sure why strip_engine_weights arg doesn't work for torch.compile()
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It just doenst make sense. Torch compile is not serializable. So why would you ever want a callable that doesnt have the weights in it
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
My understanding is that for torch.compile, _pretraced_backend() calls compile_module(), which is the same as the AOT to build TRT engines. If we pass in settings.strip_engine_weights, doesn't it tell the builder to build a weight-stripped engine?
Description
As I requested, TensorRT 10.14 added an argument
trt.SerializationFlag.INCLUDE_REFITto allow refitted engines to keep refittable. That means engines can be refitted multiple times. Based on the capability, this PR enhances the existing engine caching and refitting features as follows:compilation_settings.strip_engine_weights. Then, when users pull out the cached engine, it will be automatically refitted and kept refittable.refit_module_weights(). e.g.:_conversion.py.Type of change
Checklist: