ColossalAI icon indicating copy to clipboard operation
ColossalAI copied to clipboard

[Auto Parallel]: Speed up intra-op plan generation by 44%

Open stephankoe opened this issue 1 year ago • 0 comments
trafficstars

📌 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.

stephankoe avatar Mar 12 '24 07:03 stephankoe