Skip to content

Commit 5424a51

Browse files
houj04houj04
committed
update for KL2
1 parent 3abf26b commit 5424a51

6 files changed

Lines changed: 90 additions & 27 deletions

File tree

‎paddle/phi/backends/xpu/xpu2_op_list.cc‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,10 @@ XPUOpMap& get_kl2_ops() {
3636
XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
3737
{"accuracy", XPUKernelSet({phi::DataType::FLOAT32})},
3838
{"adadelta", XPUKernelSet({phi::DataType::FLOAT32})},
39-
{"adamw", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
39+
{"adamw",
40+
XPUKernelSet({phi::DataType::FLOAT32,
41+
phi::DataType::FLOAT16,
42+
phi::DataType::BFLOAT16})},
4043
{"adam", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
4144
{"adam_dense_param_sparse_grad",
4245
XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
@@ -723,7 +726,8 @@ XPUOpMap& get_kl2_ops() {
723726
phi::DataType::INT32,
724727
phi::DataType::INT64,
725728
phi::DataType::FLOAT16})},
726-
{"reduce_mean_grad", XPUKernelSet({phi::DataType::FLOAT32})},
729+
{"reduce_mean_grad",
730+
XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
727731
{"reduce_mean",
728732
XPUKernelSet({phi::DataType::FLOAT32,
729733
phi::DataType::FLOAT16,

‎paddle/phi/backends/xpu/xpu3_op_list.cc‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -721,7 +721,8 @@ XPUOpMap& get_kl3_ops() {
721721
phi::DataType::BFLOAT16,
722722
phi::DataType::INT32,
723723
phi::DataType::INT64})},
724-
{"reduce_mean_grad", XPUKernelSet({phi::DataType::FLOAT32})},
724+
{"reduce_mean_grad",
725+
XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
725726
{"reduce_mean",
726727
XPUKernelSet({phi::DataType::FLOAT32,
727728
phi::DataType::FLOAT16,

‎paddle/phi/kernels/xpu/adamw_kernel.cc‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -97,8 +97,7 @@ void AdamwDenseKernelKL3(const Context& dev_ctx,
9797
skip_update_ = skip_update_vec[0];
9898
}
9999

100-
// skip_update=true, just copy input to output, and TensorCopy will call
101-
// mutable_data
100+
// skip_update=true, just copy input to output
102101
if (skip_update_) {
103102
VLOG(4) << "Adamw skip update";
104103
phi::Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
@@ -495,23 +494,27 @@ void AdamwDenseKernel(const Context& dev_ctx,
495494
const float* master_param_in_data = master_param->data<float>();
496495
float* master_param_out_data =
497496
dev_ctx.template Alloc<float>(master_param_outs);
498-
// convert grad to float
499-
float* grad_fp32 = RAII_GUARD.alloc_l3_or_gm<float>(grad.numel());
500-
PADDLE_ENFORCE_XDNN_NOT_NULL(grad_fp32);
501-
// int cast(Context* ctx, const TX* x, TY* y, int64_t len);
502-
int r = xpu::cast<XPUType, float>(
503-
dev_ctx.x_context(),
504-
reinterpret_cast<const XPUType*>(grad.template data<T>()),
505-
grad_fp32,
506-
grad.numel());
507-
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
497+
// convert grad to float if necessary
498+
float* grad_fp32 = nullptr;
499+
const auto grad_type = grad.dtype();
500+
if (grad_type != phi::DataType::FLOAT32) {
501+
grad_fp32 = RAII_GUARD.alloc_l3_or_gm<float>(grad.numel());
502+
PADDLE_ENFORCE_XDNN_NOT_NULL(grad_fp32);
503+
// int cast(Context* ctx, const TX* x, TY* y, int64_t len);
504+
int r = xpu::cast<XPUType, float>(
505+
dev_ctx.x_context(),
506+
reinterpret_cast<const XPUType*>(grad.template data<T>()),
507+
grad_fp32,
508+
grad.numel());
509+
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
510+
}
508511
// int adamw(Context* ctx, const T* g, const float* mom1, const float* mom2,
509512
// const T* param, const float* beta1_pow, const float* beta2_pow, const
510513
// float* lr, float* moment1_out, float* moment2_out, T* param_out, float
511514
// beta1, float beta2, float epsilon, float coeff, int64_t n);
512515
r = xpu::adamw<float>(
513516
dev_ctx.x_context(),
514-
grad_fp32,
517+
(grad_type == phi::DataType::FLOAT32) ? grad.data<float>() : grad_fp32,
515518
moment_in_fp16 ? moment1_input_for_xdnn
516519
: moment1.template data<float>(),
517520
moment_in_fp16 ? moment2_input_for_xdnn

‎paddle/phi/kernels/xpu/reduce_mean_grad_kernel.cc‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,5 +84,9 @@ void ReduceMeanGradKernel(const Context& dev_ctx,
8484

8585
} // namespace phi
8686

87-
PD_REGISTER_KERNEL(
88-
mean_grad, XPU, ALL_LAYOUT, phi::ReduceMeanGradKernel, float) {}
87+
PD_REGISTER_KERNEL(mean_grad,
88+
XPU,
89+
ALL_LAYOUT,
90+
phi::ReduceMeanGradKernel,
91+
float,
92+
phi::dtype::float16) {}

‎test/xpu/test_adamw_op_xpu.py‎

Lines changed: 59 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626

2727
import paddle
2828
from paddle import base
29-
from paddle.base import core
3029

3130

3231
def adamw_step(inputs, attributes):
@@ -60,8 +59,8 @@ def adamw_step(inputs, attributes):
6059

6160
moment1_out = beta1 * moment1 + (1 - beta1) * grad
6261
moment2_out = beta2 * moment2 + (1 - beta2) * np.square(grad)
63-
lr_t = lr * np.sqrt(1 - beta2_pow) / (1 - beta1_pow)
64-
param_out = param - lr_t * (moment1_out / (np.sqrt(moment2_out) + epsilon))
62+
denom = (np.sqrt(moment2_out) / np.sqrt(1.0 - beta2_pow)) + epsilon
63+
param_out = param + ((moment1_out / denom) * (-(lr / (1.0 - beta1_pow))))
6564
return param_out, moment1_out, moment2_out
6665

6766

@@ -677,7 +676,7 @@ def _test_adamw_op_dygraph_place_amp_with_maingrad(
677676
param = paddle.randn(shape).astype(paddle.bfloat16)
678677
master_weight = param.astype(paddle.float32)
679678
grad = paddle.randn(shape).astype(paddle.bfloat16)
680-
main_grad = grad.astype(paddle.bfloat16)
679+
main_grad = grad.astype(paddle.float32)
681680
moment1 = paddle.randn(shape).astype(paddle.float32)
682681
moment2 = paddle.randn(shape).astype(paddle.float32).abs()
683682
lr = paddle.zeros([1]).astype(paddle.float32)
@@ -696,7 +695,7 @@ def _test_adamw_op_dygraph_place_amp_with_maingrad(
696695
# reference code
697696
_, _, _, _, _, _ = paddle._C_ops.adamw_(
698697
ref_param,
699-
main_grad.astype(paddle.float32),
698+
main_grad,
700699
lr,
701700
ref_moment_1,
702701
ref_moment_2,
@@ -780,9 +779,6 @@ def _get_places(self):
780779
return places
781780

782781
def test_main(self):
783-
xpu_version = core.get_xpu_device_version(0)
784-
if xpu_version != core.XPUVersion.XPU3:
785-
return
786782
for _ in range(1):
787783
shape = paddle.randint(1, 1024, [2])
788784
for place in self._get_places():
@@ -793,6 +789,61 @@ def test_main(self):
793789
)
794790

795791

792+
class TestAdamWOpMultiPrecison(unittest.TestCase):
793+
def _test_adamw_op_dygraph_place_amp(self, place, use_amp=False):
794+
paddle.disable_static()
795+
paddle.seed(10)
796+
paddle.set_device(place)
797+
798+
input = paddle.randn((5, 5))
799+
800+
model = paddle.nn.Linear(5, 5)
801+
802+
optimizer = paddle.optimizer.AdamW(
803+
parameters=[
804+
{
805+
'params': model.parameters(),
806+
'weight_decay': 0.001,
807+
'beta1': 0.1,
808+
'beta2': 0.99,
809+
}
810+
],
811+
multi_precision=use_amp,
812+
)
813+
814+
for idx in range(2):
815+
if place == 'xpu' and use_amp:
816+
model = paddle.amp.decorate(models=model, level='O2')
817+
scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
818+
819+
if place == 'xpu' and use_amp:
820+
with paddle.amp.auto_cast(level='O2'):
821+
output = model(input)
822+
loss = paddle.mean(output)
823+
scaled = scaler.scale(loss)
824+
scaled.backward()
825+
scaler.step(optimizer)
826+
optimizer.clear_grad()
827+
else:
828+
output = model(input)
829+
loss = paddle.mean(output)
830+
loss.backward()
831+
optimizer.step()
832+
optimizer.clear_grad()
833+
834+
def _get_places(self):
835+
places = ['cpu']
836+
if paddle.is_compiled_with_xpu():
837+
places.append('xpu')
838+
return places
839+
840+
def test_main(self):
841+
for place in self._get_places():
842+
use_amp_list = [True, False]
843+
for use_amp in use_amp_list:
844+
self._test_adamw_op_dygraph_place_amp(place, use_amp)
845+
846+
796847
support_types = get_xpu_op_support_types('adamw')
797848
for stype in support_types:
798849
create_test_class(globals(), XPUTestAdamwOp1, stype)

‎test/xpu/test_flash_attention_op_xpu.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ def setUp(self):
7979
def test_all(self):
8080
self.run_case(dtype="float32", tolerance=5e-4, tolerance_dv=5e-4)
8181
self.run_case(dtype="float16", tolerance=5e-4, tolerance_dv=1e-3)
82-
self.run_case(dtype="bfloat16", tolerance=5e-3, tolerance_dv=1e-2)
82+
self.run_case(dtype="bfloat16", tolerance=6e-3, tolerance_dv=1e-2)
8383

8484
def run_case(self, dtype, tolerance, tolerance_dv):
8585
# TODO(houj04) remove debug codes after correctness check

0 commit comments

Comments
 (0)