ColossalAI icon indicating copy to clipboard operation
ColossalAI copied to clipboard

[tensor] shape consistency generate transform path and communication cost

Open YuliangLiu0306 opened this issue 3 years ago • 0 comments

In #1418 PR, I add ShapeConsistencyManager into colossalai to support auto parallel strategy search and runtime sharding spec apply. This PR complete the auto parallel strategy search supporting part, mainly do following things:

  1. build a difference function for DimSpec and ShardingSpec based on 2D device mesh.
  2. Build a CommSpec class to record the communication operation during the transform, and use it to compute communication cost.
  3. Complete shape_consistency function development, which could generate transform path, communication action sequence and total transform cost for a given source spec and target spec pairs.
  4. Build a solved spec pair cache to reduce the runtime searching time.

Note: The apply function is not completed, it may use the CommSpec convert_spec_to_action function and some collective commucation implementation in colossalai.

YuliangLiu0306 avatar Aug 11 '22 03:08 YuliangLiu0306