[autoparallel] refactor the runtime apply pass and add docstring to passes #1757
Conversation
| def runtime_preparation_pass(gm: torch.fx.GraphModule, solution: List[int], device_mesh: DeviceMesh): | ||
| gm, sharding_spec_convert_dict, origin_node_sharding_spec_dict, comm_actions_dict = _solution_annotatation_pass( | ||
| gm, solution) | ||
| # TODO: the pass below should be uncommented after the implementation of implicit_comm_action_apply_pass completed. | ||
| # gm = implicit_comm_action_apply_pass(gm) | ||
| gm = _module_params_sharding_pass(gm, device_mesh) |
There was a problem hiding this comment.
The function names seem a bit weird to me, it is kind of saying that there are passes within passes as runtime_preparation_pass calls _solution_annotatation_pass and _module_params_sharding_pass. Would it better if we treat _module_params_sharding_pass and _solution_annotatation_pass as the helper functions instead of passes.
| from copy import deepcopy | ||
| from typing import Dict, List | ||
|
|
||
| import torch | ||
| from torch.fx.node import Node | ||
|
|
||
| from colossalai.auto_parallel.tensor_shard.sharding_strategy import ( | ||
| CommAction, | ||
| CommType, | ||
| OperationData, | ||
| OperationDataType, | ||
| ) | ||
| from colossalai.device.device_mesh import DeviceMesh | ||
| from colossalai.tensor.comm_spec import CommSpec | ||
| from colossalai.tensor.shape_consistency import ShapeConsistencyManager | ||
|
|
||
| shape_consistency_manager = ShapeConsistencyManager() | ||
|
|
||
|
|
||
| def runtime_apply(node: Node, origin_dict: Dict, input_dict: Dict, node_index: int, user_node_index: int): |
There was a problem hiding this comment.
I think we should keep the passes in the auto_parallel module instead of the fx module. The relationship should be that auto_parallel module depends on the fx module. The current relationship is that auto parallel solver relies on fx tracing, and then fx passes relies on the auto parallel sharding strategy, which leads to high coupling between these modules.
What does this PR do