ColossalAI
ColossalAI copied to clipboard
[tensor] shape consistency generate transform path and communication cost
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:
- build a difference function for DimSpec and ShardingSpec based on 2D device mesh.
- Build a CommSpec class to record the communication operation during the transform, and use it to compute communication cost.
- 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.
- 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.