PR types

New features

PR changes

Others

Description

Pcard-73145

支持流水线并行前向的非计算rank计算裁剪。动半架构中，每一张卡的组网相同，流水线并行需要在PHI API这一层，根据设备rank区分当前Layer是否需要计算。不需要计算的Layer会跳过。

在PHI API层，DistTensor会进行切分信息推导（InferSPMD）、Global Shape推导（InferMeta）、Reshard Input、Launch Kernel等操作。由于Reshard模块的设计，ChooseProperReshardFunction需要用到input的dist_attr（具体需要用到process_mesh、dims_mapping及partial_status信息），仍需执行InferSPMD来推导正确的dist_attr。InferSPMD需要Input的真实shape，这依赖InferMeta的shape推导（后续需要推全InferMeta正确性，当前matmul和elementwise的shape推导绝大情况是正确的，必须要launch kernel才能得到正确shape的op需要特殊处理）。目前PHI API实现非计算rank裁剪，只执行InferSPMD、InferMeta、创建空DistTensor操作。

1. InferMeta：如上所述，目前有两种情况需要处理：Op的InferMeta无法推导正确shape（正确的shape推导可能放在kernel里），必须launch kernel才能得到正确shape（如bincount算子）。前者需要推全算子InferMeta函数的实现正确性，后者出现情况较少，一种兜底方案是：这类算子通过手动插入paddle.distributed.reshard，或者在PHI API层面对这类算子生成特殊的reshard逻辑，首先从正确设备获得Input的dist_attr，然后退化成所有设备均执行这类算子，得到的Output释放显存，仅保留其shape信息。其它算子保持非计算rank裁剪的逻辑。
2. 创建空Tensor：当前动半的流水线并行实现不考虑性能，只需要保证能跑起来，不爆显存。由于dist_attr和global_shape信息需要在同设备的不同层间流动，需要创建一个uninitialized的DistTensor。它的显存为0，包含正确的dist_attr和global_shape信息。
3. Tensor的dtype和layout：由于PHI API设计时默认Input Tensor都是defined的，持有dtype、layout和place信息。对于reshard，它执行send/recv操作时，需要用到paddle.distributed.reshard的input的dtype来选择正确的kernel，layout也需要不同层间的正确传递（构建反向图的时候也有多处用到这三个属性）。目前的做法是，DistTensor在InferMeta时，会设置DistTensor.value的dtype和layout。这里可能会重复设置（local DenseTensor也会进行InferMeta）。
4. Tensor的place：place的处理会复杂一些，因为它是取的DistTensor.value()->holder_->place()。对于没有分配显存的DistTensor，place显然是无法获得的，目前处理方法是，PHI API对于不需要计算的算子，构建一个空的DistTensor，它的value会用空显存的holder初始化。其他情况下，如果DistTensor.value没有分配holder，可以获取默认的place。

示例代码如下：

import numpy as np
import paddle
import paddle.distributed as dist
from paddle import nn
class PPDemoNet(nn.Layer):
    def __init__(self, np_w0, np_w1, mesh0, mesh1):
        super().__init__()
        self.replicate_dist_attr0 = dist.DistAttr(
            mesh=dist.ProcessMesh([0], dim_names=["x"]),
            sharding_specs=[None, None]
        )
        self.replicate_dist_attr1 = dist.DistAttr(
            mesh=dist.ProcessMesh([1], dim_names=["x"]),
            sharding_specs=[None, None]
        )
        np_w0 = np.random.random([784, 784]).astype('float32')
        self.w0 = dist.shard_tensor(
            self.create_parameter(
                shape=[784, 784],
                attr=paddle.framework.ParamAttr(
                    name="pp_demo_weight_0",
                    initializer=paddle.nn.initializer.Assign(np_w0),
                ),
            ),
            dist_attr=self.replicate_dist_attr0,
        )
        np_w1 = np.random.random([784, 10]).astype('float32')
        self.w1 = dist.shard_tensor(
            self.create_parameter(
                shape=[784, 10],
                attr=paddle.framework.ParamAttr(
                    name="pp_nemo_weight_1",
                    initializer=paddle.nn.initializer.Assign(np_w1),
                ),
            ),
            dist_attr=self.replicate_dist_attr1,
        )
    def forward(self, x):
        y = paddle.matmul(x, self.w0)
        y = dist.reshard(y, dist_attr=self.replicate_dist_attr1)
        z = paddle.matmul(y, self.w1)
        return z