PR types

Function optimization

PR changes

Others

Description

Pcard-70448

Refactor the Op internal dims mapping update Logic:

• Enable New SPMD Rule in static mode Completion, by FLAG: FLAGS_infer_spmd_enable=1 (default=1)
• Deprecate the original update_dims_mapping and is_xxx_compatible methods used in Completion (matmul, elementwise, reduction, layer_norm are adapted).
• Re-mapping DistOp: New SPMD Rule may result an parallel pattern which has NOT corresponding DistOperatorImpl for it. We re-mapping the op to be replicated when it happen.

the new logic would has two main branch:
left: orignal logic for ops that has not new spmd rule
right: new logic for ops that has new spmd rule implemented.

How we plan to refactor static mode Semi Auto Parallel:

Currently:

there are 3 problems:

1. Limited Coverage: each DistOperatorImpl is hard code for a special parallel pattern of the operator (both spmd_rule and parallel_transform), and only limited parallel pattern are supported.
2. Code Redundancy: There are plenty of redundant logic between DistOperatorImpls of same op (in spmd_rule and parallel_transform)
3. Not Compatible with Semi Auto Parallel in Dynamic mode.

After Refactor Completion (This and next PRs):

• Each Operator will only have one spmd_rule which will be general for all parallel patterns.
• Resulted parallel pattern from SpmdRule is mapped to DistOperatorImpl. If there is not corresponding impl, it will be reverted to replicated.

After Refactor DistOperatorImpl (Future PRs):

• LocalReshard module will integraled to unify all DistOperatorImpls of same operator and support transformation need by any parallel patterns.