[autoparallel] refactored shape consistency to remove redundancy #1591
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.
This PR refactored the usage of shape consistency manager. The details are described below.
Problem
Previously, each handler and strategy constructor requires a shape consistency manager object as input for initialization. However, this manager is only used in calculating the resharding cost and nowhere else. Meanwhile, calculating resharding cost is implemented twice in the op handler and strategy constructor. After moving this to utils, the handler and constructor do not need to keep the manager as a member any more.
Fix
The shape consistency manager is now a singleton class where there will only be one instance in the global namespace. The usage of resharding is now purely via function call (i.e.
generate_resharding_cost
andset_shape_consistency_manager_options
, even though the latter is not implemented yet). The handler and constructor only need to call these functions now.Side note
generate_sharding_spec
.handle_backward
. This is to configure whether backward stats should be computed, can be useful for inference as backward is not needed.