feat: improve engine caching and fix bugs #3932
Conversation
github-actions
bot
added
component: tests
| """ | ||
|
|
||
| def _insert_engine_to_cache( | ||
| hash_val: str, interpreter_result: TRTInterpreterResult |
There was a problem hiding this comment.
Is there a reason that the function needs to be in the interpret functions scope?
There was a problem hiding this comment.
Not a specific reason, but I just don't know when the engine_cache will be used other than in the function interpret_module_to_result(). To make it safe and self-contained, I picked the smallest scope. Is there any other cases that might use engine_cache?
There was a problem hiding this comment.
Well it will get redefined each time interpret is called if its in scope, not sure its necessary unless there was some reason like it uses context from interpret module
| logger.info(f"Engine was successfully inserted into cache for hash: {hash_val}") | ||
|
|
||
| @needs_refit # type: ignore[misc] | ||
| def _pull_cached_engine(hash_val: str) -> Optional[SerializedInterpreterResult]: |
| """ | ||
|
|
||
| def _insert_engine_to_cache( | ||
| hash_val: str, interpreter_result: TRTInterpreterResult |
There was a problem hiding this comment.
Is there a reason that the function needs to be in the interpret functions scope?
| ) | ||
| logger.info(f"Engine was successfully inserted into cache for hash: {hash_val}") | ||
|
|
||
| @needs_refit # type: ignore[misc] |
There was a problem hiding this comment.
Should the insert and extract both be needs refit?
Also shouldnt this gracefully pass through vs the typically unimplemented error?
There was a problem hiding this comment.
Should the insert and extract both be needs refit?
insert seems not to involve any refitting stuff. It supports a scenario that users insert engines on machine A that doesn't support refit but pull engines on machine B that supports refit. Please correct me if wrong.
Also shouldnt this gracefully pass through vs the typically unimplemented error?
Not sure if I understand your question correctly. The reason why we need refit in pull is that we save weight-stripped engine in this implementation, which needs to be refitted to get correct weights before using.
There was a problem hiding this comment.
Should the insert and extract both be needs refit?
If we support saving and pulling identical immutable weights engines then we dont need to guard at all, but if we are just saving immutable weights engines but never reusing them I dont see the point of supporting just insert
Also shouldnt this gracefully pass through vs the typically unimplemented error?
The needs_X decorator will throw unimplemented if its hit with the feature not enabled. But in the case of refit if its not enabled then we should just bypass the cache
There was a problem hiding this comment.
Ah I got you. fixed the two issues
|
@cehongwang please take a pass so we have multiple eyes on this PR |
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: