PR types

New features

PR changes

Others

Description

Pcard-70448

Background
"Partial" is a data distribution type for tensor like "Replicated" and "Sharded".
A tensor is Partial means that: its shape is the same among ranks, but the local element value in each rank is only a partial of the global value, and reduced (sum/max/min/all/any) op over ranks is need to rebuild the global tensor from the locals.

Motivation
Before the PR, dist_tensor in auto parallel only has two type: "Replicated" and "Sharded". And it work quite well for most hybrid parallelism scenario. It is just a design choice to introduce a third distribution type "Partial".

1. "Partial" make it easier to explain why there need an Allreduce in Tensor Parallel (Vocab Parallel Embedding, Forward of Row Parallel Linear, Backward of Col Parallel Linear, etc ) for user and developer.
2. "Partial" make some Optimizations easier in Dynamic Graph mode of Auto Parallel. (Previously, Auto Parallel Only has Static Graph mode, and we could delay any Allreduce by Graph Transformation Pass. But we need "Partial" semantic for these optimizations in Dynamic Graph mode)
   

HOW
The Introducing of Partial is divided into two Stages.

First Stage(This PR):

• Introduce partial but not allow its propagation in computation graph.
• DistTensorAttribute will add new data member to maintain Partial related info.
• Op's inputs will always be non-partial.
• If an Op generate a Partial output, the partial will clean by Reshard(Redistribute) immediately.
• Meet the need of Motivation-I

Second Stage(Future PR):

• Allow partial to propagate in computation graph.
• User is allowed to mark tensor as Partial for Optimization.
• Op with linearity property will propagate the partial status from input to it output.
• Reshard(Partial --> Replicated) would only be called before the last non-linearity Op who consumes that tensor.
• Meet the need of Motivation-II