[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.
please add docstring for it
There was a problem hiding this comment.
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.
it is a little strange that both root_graph and graph are used, they are the same graph?
| @@ -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.
better to add docstring for this class
There was a problem hiding this comment.
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.
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.
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. |
|
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_wrapperdecorator outside the full PyTorch model.config.execution_engine = 'py'toRetiariiExeConfig.