fix model run error when use auto parallel and recompute(use_reentrant=false)

[jeff41404]

PR Category

Auto Parallel

PR Types

Bug fixes

Description

pcard-84677
model will run error and core dump when meets 2 conditions:

class Large_model(...... #sub-class of paddle.nn.Layer):
    def __init__(
        self,
        ......
    ):
        super().__init__()
        ......
        self.layers = nn.LayerList()
        for i in range(num_layers):
            # condition 1 : use auto parallel shard parameters, e.g. shard parameters in TransformerBlock
            self.layers.append(TransformerBlock(......))

    def forward(self, x, label):
        ......
        for i, layer in enumerate(self.layers):
            # condition 2 : use recompute and set use_reentrant=False
            x = paddle.distributed.fleet.utils.recompute(layer, x, use_reentrant=False)
        ......
        return loss

model = Large_model(......)
for step, inputs in enumerate(dataloader):
    # inputs contain data and label
    loss = model(inputs)
    loss.backward() # error and core dump, this PR will fix this issue

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[zhiqiu]

LGTM

[JiabinYang]

LGTM for Tensor construction