Modify Parallelization Strategy to Make it More General #1988
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
…parallelization_strategy
optimum/fx/parallelization/api.py
Outdated
| """ | ||
| API for automatic model parallelism through Pytorch FX. | ||
|
|
||
| Args: | ||
| model (Union[torch.nn.Module, str]): | ||
| Model to parallelize, could either be a module or a model id on the Huggingface Hub. | ||
| model (str): |
There was a problem hiding this comment.
Suggested change
| model (str): | |
| model (`str`): |
optimum/fx/parallelization/decomp.py
Outdated
Comment on lines
63
to
67
| class DecompTracer(GraphAppendingTracer): | ||
| def __init__(self, graph: Graph): | ||
| super().__init__(graph) | ||
| self.tensor_tracker = WeakTensorKeyDictionary() | ||
| self.symnode_tracker = _SymNodeDict() |
There was a problem hiding this comment.
Maybe a docstring explaining what it does.
optimum/fx/parallelization/decomp.py
Outdated
Comment on lines
74
to
79
| that certain primitive layers(like `nn.Linear`, `nn.Embedding`, and activation layers) are preserved because we have specific | ||
| heuristic based parallelization strategy for them so that we can conveniently replace them into their parallelized counterparts | ||
| in the orignal graph module. | ||
|
|
||
| Note that the traced graph is a low-level equivalent representation of the original graph module, and is only used for | ||
| parallel axis propagation and analysis, the original graph module is still used for real execution. |
There was a problem hiding this comment.
Can you group notes as follows:
Notes:
1. Certain primitive layers ....
2. The traced graph is a low-level equivalent...
optimum/fx/parallelization/decomp.py
Outdated
| leaf_function_targets: List[Callable] = [F.scaled_dot_product_attention], | ||
| ) -> Callable: | ||
| """ | ||
| API to decompose and funcitonalize a high-level graph module. |
There was a problem hiding this comment.
Suggested change
| API to decompose and funcitonalize a high-level graph module. | |
| API to decompose and functionalize a high-level graph module. |
optimum/fx/parallelization/decomp.py
Outdated
Comment on lines
196
to
200
| graph_module (GraphModule): | ||
| The high-level graph module to be decomposed and functionalized. | ||
| decomposition_table (Dict[torch._ops.OperatorBase, Callable], defaults to `core_aten_decompostions()`): | ||
| The lookup table which maps high-level torch op to their equivalent low-level implementation. | ||
| leaf_function_targets (List[Callable], defaults to `[F.scaled_dot_product_attention]`): |
There was a problem hiding this comment.
Suggested change
| graph_module (GraphModule): | |
| The high-level graph module to be decomposed and functionalized. | |
| decomposition_table (Dict[torch._ops.OperatorBase, Callable], defaults to `core_aten_decompostions()`): | |
| The lookup table which maps high-level torch op to their equivalent low-level implementation. | |
| leaf_function_targets (List[Callable], defaults to `[F.scaled_dot_product_attention]`): | |
| graph_module (`GraphModule`): | |
| The high-level graph module to be decomposed and functionalized. | |
| decomposition_table (`Dict[torch._ops.OperatorBase, Callable]`, defaults to `core_aten_decompostions()`): | |
| The lookup table which maps high-level torch op to their equivalent low-level implementation. | |
| leaf_function_targets (`List[Callable]`, defaults to `[F.scaled_dot_product_attention]`): |
Comment on lines
25
to
41
| class Registry: | ||
| def __init__(self) -> None: | ||
| self.mapping = {} | ||
|
|
||
| def register(self, op_types): | ||
| def wrapper(cls): | ||
| if isinstance(op_types, (list, tuple)): | ||
| for op_type in op_types: | ||
| self.mapping[op_type] = cls | ||
| else: | ||
| self.mapping[op_types] = cls | ||
| return cls | ||
|
|
||
| return wrapper | ||
|
|
||
| def is_supported(self, op_type) -> bool: | ||
| return op_type in self.mapping |
There was a problem hiding this comment.
What is this registry used for?
There was a problem hiding this comment.
This is for registration of parallel axis policy of different aten ops
There was a problem hiding this comment.
Can you add a docstring to explain that please?
Comment on lines
127
to
130
| def propagate(self) -> bool: | ||
| arg = self.node.all_input_nodes[0] | ||
| axis = self.extract_axis(arg) | ||
| return [axis] |
There was a problem hiding this comment.
nit it's not returning a bool.
If I understand properly it returns the axis that is supposed to be parallel?
There was a problem hiding this comment.
Yes, you are right, it looks up the axis of all the inputs and try inferencing the axis of the output, and will return empty list if no valid axis on which the output can be parallelized
Comment on lines
25
to
41
| class Registry: | ||
| def __init__(self) -> None: | ||
| self.mapping = {} | ||
|
|
||
| def register(self, op_types): | ||
| def wrapper(cls): | ||
| if isinstance(op_types, (list, tuple)): | ||
| for op_type in op_types: | ||
| self.mapping[op_type] = cls | ||
| else: | ||
| self.mapping[op_types] = cls | ||
| return cls | ||
|
|
||
| return wrapper | ||
|
|
||
| def is_supported(self, op_type) -> bool: | ||
| return op_type in self.mapping |
There was a problem hiding this comment.
Can you add a docstring to explain that please?
Comment on lines
+199
to
+201
| def search(idx: int): | ||
| if idx == len(nodes): | ||
| return True |
There was a problem hiding this comment.
Suggested change
| def search(idx: int): | |
| if idx == len(nodes): | |
| return True | |
| def search(idx: int) -> bool: | |
| return idx == len(nodes) |
There was a problem hiding this comment.
the search here actually is a backtracking search function entailing more logic following, so we can only return True if we have reached the very last op
…parallelization_strategy
michaelbenayoun
approved these changes
Sep 2, 2024
Comment on lines
27
to
28
| Registry class handles registration of parallel axis propagation handlers of different aten ops, to support a new | ||
| aten op, you need to register the corresponding handler class by decorating it with `register` function. |
There was a problem hiding this comment.
nit:
Suggested change
| Registry class handles registration of parallel axis propagation handlers of different aten ops, to support a new | |
| aten op, you need to register the corresponding handler class by decorating it with `register` function. | |
| Registry class handles registration of parallel axis propagation handlers of different aten ops. | |
| To support a new aten op, you need to register the corresponding handler class by decorating | |
| it with the `register` function. |
…parallelization_strategy
|
merge this for irrelevant failures |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
As per title, this PR tries a more general approach rather than relying purely on human heuristics, basically it uses the following steps to search a possible parallelization strategy for a transformer model
And for the API design, we disable the support of passing custom modules and only focus on models in transformers because supporting custom models is not the priority for now.