ColossalAI
ColossalAI copied to clipboard
[Auto Parallel]: Speed up intra-op plan generation by 44%
📌 Checklist before creating the PR
- [X] I have created an issue for this PR for traceability
- [ ] The title follows the standard format:
[doc/gemini/tensor/...]: A concise description - [ ] I have added relevant tags if possible for us to better distinguish different PRs
🚨 Issue number
resolves #5436
📝 What does this PR do?
As stated in issue #5436, the generation of a DimSpec object is comparatively costly because a dictionary is created every time and a deep copy of two strings is being made. The large volume of DimSpec objects created leads to several seconds spent in total by creating DimSpec objects when generating intra-op plans with ColossalAuto.
This pull requests addresses this inefficiency by converting this dict to a class attribute of the DimSpec class, so it's shared among all its instances. This dict is initialized lazily the first time the property difference_dict is used. This is possible because the contents of difference_dict are not modified by other portions of the code.
Additionally, the methods build_difference_2d_dict and convert_str_to_shard_list are made class/static methods and private because they don't need access to the instance properties.
Effect
These changes reduce the end-to-end wall-clock time to build the strategy constructor by 44% while running the script examples/tutorial/auto_parallel/auto_parallel_with_resnet.py. The time required to compute a solution on a laptop with Intel Core i7 7700HQ was on average 56.79 s before the change. After applying above-mentioned changes, the wall-clock time reduced to about 31.93 s.
💥 Checklist before requesting a review
- [x] I have linked my PR to an issue (instruction)
- [x] My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
- [x] I have performed a self-review of my code
- [x] I have added thorough tests.
- [x] I have added docstrings for all the functions/methods I implemented
⭐️ Do you enjoy contributing to Colossal-AI?
- [x] 🌝 Yes, I do.
- [ ] 🌚 No, I don't.