Moe support #3811
Moe support #3811
Conversation
github-actions
bot
added
component: tests
cehongwang
force-pushed
the
moe-support
branch
from
22edfe6 to
fa53986
Compare
cehongwang
force-pushed
the
moe-support
branch
from
fa53986 to
6ea89ae
Compare
There was a problem hiding this comment.
Could you add detailed comments demonstrating what it does now that you've gone through the entire converter ? that would be helpful
| model = model.to(torch.float32) | ||
|
|
||
| return model | ||
| return model.cuda() |
There was a problem hiding this comment.
model.cuda() is already done during initialization
| # TODO: @Evan is waiting for TRT's feature to cache the weight-stripped engine | ||
| # if not self.compilation_settings.strip_engine_weights: | ||
| # # set EXCLUDE_WEIGHTS flag to strip weights | ||
| # runtime = trt.Runtime(TRT_LOGGER) |
There was a problem hiding this comment.
why are these being deleted?
There was a problem hiding this comment.
Don't need to deserialize the engine anymore. The engine is already live
| assert isinstance( | ||
| serialized_engine, bytes | ||
| ), "Serialized engine must be a bytes object" | ||
| self.engine = serialized_engine |
There was a problem hiding this comment.
deserialize the engine here
|
|
||
| return TRTInterpreterResult( | ||
| engine_str, | ||
| cuda_engine, |
There was a problem hiding this comment.
How will this work with the deferred engine setup?
There was a problem hiding this comment.
Generally I think this is not clearly implemented. If we decided that PyRuntime will use initialized engines, then you should quickly reduce other cases to an initialized engine.
| assert isinstance( | ||
| serialized_engine, bytes | ||
| ), "Serialized engine must be a bytes object" | ||
| self.engine = serialized_engine |
There was a problem hiding this comment.
Also do not overload the meaning of self.engine. Either it should be the serialized engine or the live engine not both
|
|
||
| def setup_engine(self) -> None: | ||
|
|
||
| if isinstance(self.engine, trt.ICudaEngine): |
There was a problem hiding this comment.
We need to think about how this will interact with the greater system like the lazy_engine_init option.
|
|
||
| if isinstance(self.engine, trt.ICudaEngine): | ||
| pass | ||
| elif isinstance(self.engine, bytes): |
There was a problem hiding this comment.
we should not check by type, have two members and a state flag
|
|
||
| def _on_state_dict(self, state_dict: Dict[str, Any], prefix: str, _: Any) -> None: | ||
| state_dict[prefix + "engine"] = self.serialized_engine | ||
| state_dict[prefix + "engine"] = self.engine |
There was a problem hiding this comment.
Same thing here, this should just be serialized engine
| error_msgs: Any, | ||
| ) -> None: | ||
| self.serialized_engine = state_dict[prefix + "engine"] | ||
| self.engine = state_dict[prefix + "engine"] |
|
|
||
| def __init__( | ||
| self, | ||
| cuda_engine: Optional[trt.ICudaEngine | bytes] = None, |
There was a problem hiding this comment.
Should not be a union. Its either you give a cuda_engine or a serialized engine
cehongwang
force-pushed
the
moe-support
branch
from
9e7ca5d to
c286767
Compare
Description
Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.
Fixes # (issue)
Type of change
Please delete options that are not relevant and/or add your own.
Checklist: