fix pylayer problem with amp

[zhiqiu]

PR types

Bug fixes

PR changes

APIs

Describe

When using PyLayer with amp, the forward part of PyLayer may be invoked with amp enabled, while the backward part is not, thus resulting in errors.

related #39881

import paddle
from paddle.autograd import PyLayer

class MyMM(PyLayer):
    @staticmethod
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return a.mm(b)
    
    @staticmethod
    def backward(ctx, grad):
        a, b = ctx.saved_tensor()  
        # NOTE(zhiqiu): a and b is float32 now, while grad is fp16 when forward runs with auto_cast()
        # thus, the mm operation raise errors because of the dtype of inputs are inconsistent
        return grad.mm(b.t()), a.t().mm(grad)

x = paddle.rand([10, 10])
y = paddle.rand([10, 10])
x.stop_gradient=False
y.stop_gradient=False

with paddle.amp.auto_cast():
   res = MyMM.apply(x, y)
   loss = paddle.mean(res)

loss.backward()

• before
  run failed.
  
• after
  run ok

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[chenwhql]

LGTM

[zhangbo9674]

LGTM

[zhangbo9674]

LGTM