-
Notifications
You must be signed in to change notification settings - Fork 1.8k
[Retiarii] Pure-python execution engine #3605
Conversation
# The following are written for pure-python mode | ||
|
||
|
||
class ManyChooseManyMutator(Mutator): |
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.
please add docstring for 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.
This is an internal API. Added a short description.
return model, None | ||
|
||
mutators = [] | ||
for nodes in _group_by_label_and_type(graph.hidden_nodes): |
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 is a little strange that both root_graph and graph are used, they are the same graph?
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.
I think so.
@@ -61,3 +65,35 @@ def get_module_name(cls_or_func): | |||
def get_importable_name(cls, relocate_module=False): | |||
module_name = get_module_name(cls) if relocate_module else cls.__module__ | |||
return module_name + '.' + cls.__name__ | |||
|
|||
|
|||
class ContextStack: |
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.
better to add docstring for this class
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.
for example why we need this context manager?
@@ -149,3 +151,15 @@ def basic_unit(cls): | |||
import torch.nn as nn | |||
assert issubclass(cls, nn.Module), 'When using @basic_unit, the class must be a subclass of nn.Module.' | |||
return serialize_cls(cls) | |||
|
|||
|
|||
def model_wrapper(cls): |
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 is not clear how to use this wrapper. what is the difference between model_wrapper
and serializer
?
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.
I think it's documented in the docstring.
@ultmaster looks good, but where is document for this execution engine? |
I think it's written in the PR description. The original plan is to fully test it and make this execution engine default. I didn't expect that this PR takes so longer to review. Anyway, now we are running out of time, and I think the safest plan is to document it and make this execution engine optional if user need it. |
Pure-python execution engine (experimental) | ||
------------------------------------------- | ||
|
||
If you are experiencing issues with TorchScript, or the generated model code by Retiarii, there is another execution engine called Pure-python execution engine which doesn't need the code-graph conversion. This should generally not affect models and strategies in most cases, but customized mutation might not be supported. |
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.
终于可以不用torchscript了,大佬牛逼
nice work! |
This PR adds an execution engine that can directly execute graph without touching any graph converter and codegen stuff. It is expected to become the default execution engine in version 2.3.
I've also implemented a more useful mutation history tracking tool. This should make the upcoming visualization thing possible.
I tried to refact/damage as little code as possible. Some of the refactoring effort might come in future releases.
I think a lot of tests are needed before it is ready for review.Depending on #3607.
User experience
Currently I don't have the confidence to make this execution engine default. Two modifications are needed to (manually) enable the new engine.
@model_wrapper
decorator outside the full PyTorch model.config.execution_engine = 'py'
toRetiariiExeConfig
.