feat: Support weight-stripped engine and REFIT_IDENTICAL flag #3167
Conversation
github-actions
bot
added
component: tests
| name: str = "", | ||
| settings: CompilationSettings = CompilationSettings(), # Assumes engine was built with default compilation settings if object not passed | ||
| weight_name_map: Optional[dict[Any, Any]] = None, | ||
| graph_module: torch.fx.GraphModule = None, |
There was a problem hiding this comment.
@narendasan I tried to do refitting for C++ runtime like for Python runtime but didn't work. Any suggestions? should I do in C++ or Python?
There was a problem hiding this comment.
Doesnt refit already work on both apis?
There was a problem hiding this comment.
Also why do we need the graph module in this module?
There was a problem hiding this comment.
-
In this PR I moved the refitting part into TRTModule, so only works for Python runtime.
-
graph module is used for refitting
| @@ -619,27 +609,32 @@ def run( | |||
| builder_config, self.compilation_settings.timing_cache_path | |||
| ) | |||
|
|
|||
| serialized_engine = self.builder.build_serialized_network( | |||
| # if strip_engine_weights is true, the serialized engine need to be refitted before using | |||
| maybe_unrefitted_serialized_engine = self.builder.build_serialized_network( | |||
There was a problem hiding this comment.
Why is this maybe unrefitted engine?
There was a problem hiding this comment.
please see the design in the comment below. If compilation_settings.strip_engine_weights is true, it needs to be refitted, else it doesn't. so it's maybe
| ), "weight-stripped engines must be refittable, please set make_refittable=True" | ||
|
|
||
| # Refit the weights | ||
| refitter = trt.Refitter(self.engine, TRT_LOGGER) |
There was a problem hiding this comment.
Can you use this function?
TensorRT/py/torch_tensorrt/dynamo/_refit.py
Line 138 in fa02fd3
There was a problem hiding this comment.
The function requires input_list which is not provided in the caller.
| @@ -121,6 +124,52 @@ def setup_engine(self) -> None: | |||
| self.engine = runtime.deserialize_cuda_engine(self.serialized_engine) | |||
| self.context = self.engine.create_execution_context() | |||
|
|
|||
| if self.settings.strip_engine_weights: | |||
There was a problem hiding this comment.
We likely shouldnt be doing the refit in these modules
I think for weight stripping there are 3 workflows.
- a user just wants a weight stripped engine. They should use
convert_exported_program_to_trt_enginewith settingsstrip_weights. The choice ofmake_refittablecan be used to decide betweenkREFITandkREFIT_IDENTICAL(though it might not be entirely clear so we might want to think about that setting). - We want to utilize weight stripping to have a lighter weight cache. Here this choice is opaque to the user. The user choice of
make_refittablecontrols if we usekREFITorkREFIT_IDENTICAL. But once the engine is loaded or we pull from cache we immediately refit (prior to passing the engine to the TRTModule). Same as we do today - The user wants a stripped weights compiled program (im not sure why or if this is a real usecase). Here, this is basically the same as lazy engine loading. We would require that users need to run through
refit_engine_weightsbefore executing.
There was a problem hiding this comment.
Got it. The very beginning idea/design is commented below. I'll move the refitting part back to TRTInterpreter.run()
The choice of
make_refittablecan be used to decide betweenkREFITandkREFIT_IDENTICAL
Do you mean we use make_refittable to control both kREFIT and kREFIT_IDENTICAL?
|
@narendasan Ok, at first the overall design was like: In TRTInterpreter.run(): if compilation_settings.strip_engine_weights is True:
if engine_cache not hit:
1. build a weight-stripped engine
2. save the weight-stripped engine if engine_cache is set
3. return the weight-stripped engine (not yet refit)
else:
load and return the weight-stripped engine (not yet refit)
else:
if engine_cache not hit:
1. build a weight-included engine
2. save the weight-included engine if engine_cache is set
3. return the weight-included engine (don't need to refit)
else:
load and return the weight-included engine (not yet refit)Then, in TRTModule, refit if necessary before inference. |
Description
Fixes #3146
Type of change
Checklist: