【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad #68079

megemini · 2024-09-08T09:25:58Z

PR Category

User Experience

PR Types

New features

Description

【Hackathon 7th No.12】Adam、AdamW 优化器支持 amsgrad

关联：

本地对比 pytorch 的结果，两者一致：

比对代码

import numpy as np

import torch
import paddle


def func(t, x):
    if t % 101 == 1:
        return 1010 * x
    else:
        return -10 * x


np.random.seed(2024)
data = np.array(0).astype("float64")
epoch = 500
lr = 0.1

for amsgrad in [True, False]:
    for opt_name, opt_torch, opt_paddle in [
        ["Adam", torch.optim.Adam, paddle.optimizer.Adam],
        ["AdamW", torch.optim.AdamW, paddle.optimizer.AdamW],
    ]:
        for torch_device, paddle_device in [["cpu", "cpu"], ["cuda", "gpu"]]:
            print(f"------ optimizer is : {opt_name} ; compare : {paddle_device}------")
            print(f"------ pytorch ------")
            x = torch.tensor(data, device=torch.device(torch_device))
            x.requires_grad = True

            optimizer = opt_torch([x], lr=lr, amsgrad=amsgrad)
            for i in range(epoch):
                y = func(i, x)
                optimizer.zero_grad()
                y.backward()
                optimizer.step()

            if torch_device == "cuda":
                x_torch = x.cpu().detach().numpy()
                y_torch = y.cpu().detach().numpy()
            else:
                x_torch = x.detach().numpy()
                y_torch = y.detach().numpy()

            print(f"------ paddle ------")
            paddle.set_device(paddle_device)
            x = paddle.to_tensor(data)
            x.stop_gradient = False

            optimizer = opt_paddle(parameters=[x], learning_rate=lr, amsgrad=amsgrad)
            for i in range(epoch):
                y = func(i, x)
                optimizer.clear_grad()
                y.backward()
                optimizer.step()

            x_paddle = x.numpy()
            y_paddle = y.numpy()

            np.testing.assert_allclose(x_torch, x_paddle, atol=1e-06, rtol=1e-06)
            print(x_torch, x_paddle)
            print(y_torch, y_paddle)
            print(f"------- compare finish ---------")

输出结果：

------ optimizer is : Adam ; compare : cpu------
------ pytorch ------
------ paddle ------
0.382819332566745 0.3828193325667452
-3.7319234136114865 -3.7319234136114887
------- compare finish ---------
------ optimizer is : Adam ; compare : gpu------
------ pytorch ------
------ paddle ------
0.3828193325667449 0.38281933256674533
-3.7319234136114856 -3.73192341361149
------- compare finish ---------
------ optimizer is : AdamW ; compare : cpu------
------ pytorch ------
------ paddle ------
0.38940724227589385 0.389407242265435
-3.801604114817793 -3.8016041146280424
------- compare finish ---------
------ optimizer is : AdamW ; compare : gpu------
------ pytorch ------
------ paddle ------
0.38940724227589385 0.3894072422654346
-3.801604114817793 -3.801604114628038
------- compare finish ---------
------ optimizer is : Adam ; compare : cpu------
------ pytorch ------
------ paddle ------
0.47233193956960806 0.47233193956960845
-4.62253146676283 -4.622531466762833
------- compare finish ---------
------ optimizer is : Adam ; compare : gpu------
------ pytorch ------
------ paddle ------
0.472331939569608 0.4723319395696082
-4.62253146676283 -4.6225314667628306
------- compare finish ---------
------ optimizer is : AdamW ; compare : cpu------
------ pytorch ------
------ paddle ------
0.462192080569021 0.46219208087997216
-4.525658535292251 -4.525658538303618
------- compare finish ---------
------ optimizer is : AdamW ; compare : gpu------
------ pytorch ------
------ paddle ------
0.46219208056902106 0.46219208087997266
-4.525658535292251 -4.525658538303623
------- compare finish ---------

Update 20240908

已在本地完成：
- test_adam_op.py
- test_adamw_op.py
- test_merged_adam_op.py
- test_fused_adam_op.py
相关测试。
需要在 CI 环境中验证分布式的测试项目
需要在 CI 环境中验证其他测试项目

另外，xpu 的 amsgrad 变体，由于 xpu 底层接口暂不支持，因此，此处只修改了相关的输入输出参数列表。

… hack7_amsgrad

paddle-bot · 2024-09-08T09:26:03Z

你的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.

HydrogenSulfate

添加的ams_grad是否会影响原有的代码执行逻辑和存储空间占用情况？PR的代码起来无论是否开启ams_grad，都会比原先没有amsgrad的代码多申请一段mom2_max的空间，以及有一些多余的变量产生。

HydrogenSulfate · 2024-09-09T11:08:12Z

paddle/phi/kernels/funcs/adam_functors.h


  inline HOSTDEVICE void operator()(size_t i) const {
    // Merge all memory access together.
    T g = grad_[i];
    T mom1 = moment1_[i];
    T mom2 = moment2_[i];
+    T mom2_max = moment2_max_[i];


这个是必须要记录的吗？

HydrogenSulfate · 2024-09-09T11:08:55Z

paddle/phi/kernels/funcs/adam_functors.h

+    T mom2_max_;
+    if (amsgrad_) {
+      mom2_max_ = std::max(mom2, mom2_max);
+      p -= lr * (mom1 / (sqrt(mom2_max_) + epsilon_ * sqrt(1 - beta2_pow)));
+    } else {
+      mom2_max_ = mom2_max;
+      p -= lr * (mom1 / (sqrt(mom2) + epsilon_ * sqrt(1 - beta2_pow)));
+    }

    // Write back to global memory
    moment1_out_[i] = mom1;
    moment2_out_[i] = mom2;
+    moment2_max_out_[i] = mom2_max_;


同理，如果amsgrad没有开启，建议不要添加任何多余的变量和相关计算逻辑，保持原样即可

HydrogenSulfate · 2024-09-09T11:10:31Z

paddle/phi/kernels/funcs/adam_functors.h

+    Eigen::Map<Eigen::Array<T, 1, Eigen::Dynamic>> moment2_max_out{
+        moment2_max_out_, static_cast<Eigen::Index>(numel)};


同上，如果没有开启amsgrad，是否会有mom2_max相关的冗余运算和存储占用？

HydrogenSulfate · 2024-09-09T11:10:50Z

paddle/phi/kernels/funcs/adam_functors.h


  inline HOSTDEVICE void adam_update(size_t i, T g) const {
    // The following code is the same as dense
    T mom1 = moment1_[i];
    T mom2 = moment2_[i];
+    T mom2_max = moment2_max_[i];


HydrogenSulfate · 2024-09-09T11:12:13Z

paddle/phi/kernels/funcs/jit/refer/refer.h

@@ -14,6 +14,7 @@

 #pragma once

+#include <stdio.h>


这个头文件是什么有代码依赖吗？

调试之后忘记删掉了，抱歉～

HydrogenSulfate · 2024-09-09T11:17:40Z

python/paddle/optimizer/adam.py

@@ -117,6 +117,7 @@ class Adam(Optimizer):
            The default value is False.
        multi_precision (bool, optional): Whether to use multi-precision during weight updating. Default is false.
        use_multi_tensor (bool, optional): Whether to use multi-tensor strategy to update all parameters at once . Default is false.
+        amsgrad (bool, optional): Whether to use the AMSGrad of this algorithm. Default is false.


参数说明请参考：https://pytorch.org/docs/stable/_modules/torch/optim/adam.html#Adam写一下

HydrogenSulfate · 2024-09-09T11:18:20Z

python/paddle/optimizer/adamw.py

@@ -104,6 +104,7 @@ class AdamW(Optimizer):
            different semantics with the original Adam algorithm and may lead to different result.
            The default value is False.
        multi_precision (bool, optional): Whether to use multi-precision during weight updating. Default is false.
+        amsgrad (bool, optional): Whether to use the AMSGrad of this algorithm. Default is false.


… hack7_amsgrad

megemini · 2024-09-09T13:39:16Z

添加的ams_grad是否会影响原有的代码执行逻辑和存储空间占用情况？PR的代码起来无论是否开启ams_grad，都会比原先没有amsgrad的代码多申请一段mom2_max的空间，以及有一些多余的变量产生。

这个之前考虑过，主要是因为，目前涉及到 amsgrad 的地方太多了，所以优化相关的事情想先往后放一下～

那我现在改一下试试吧～

HydrogenSulfate · 2024-09-09T14:23:13Z

添加的ams_grad是否会影响原有的代码执行逻辑和存储空间占用情况？PR的代码起来无论是否开启ams_grad，都会比原先没有amsgrad的代码多申请一段mom2_max的空间，以及有一些多余的变量产生。

这个之前考虑过，主要是因为，目前涉及到 amsgrad 的地方太多了，所以优化相关的事情想先往后放一下～

那我现在改一下试试吧～

这一点影响是比较大的。因为一般情况下优化器是逐元素跟踪参数状态，所以优化器每一个统计量需要记录的数量都等于模型参数数量，adam(w)这种带动量的优化器则更会多。因此模型训练过程中显存占比前三就是中间状态、优化器参数、模型参数，如果没有优化，很可能原先在16G上能训的下的CV、NLP模型就会OOM了，更不用说B级别参数量的大模型
代码本身的计算逻辑应该没太大问题，目前没有优化的版本可以用于快速验证正确性，但最终版本一定要考虑到这种基本但必要的优化

HydrogenSulfate · 2024-09-09T14:31:15Z

另外可以在修改完成后，用ResNet50或者其他模型，以fake data为输入做一个对比，确认下amsgrad关闭时，显存无变化，开启时显存增加量与参数量基本相同。

… hack7_amsgrad

HydrogenSulfate · 2024-11-26T06:23:01Z

@megemini hello大佬，我们内部测试了最新的这个PR，应该是没问题了，还麻烦解决一下冲突

… hack7_amsgrad

megemini · 2024-11-27T04:25:42Z

@megemini hello大佬，我们内部测试了最新的这个PR，应该是没问题了，还麻烦解决一下冲突

非常感谢！！！听说过程非常曲折 😂😂😂 感谢～～～

冲突已经解决～ PR-CI-NPU-910B-Paddle 这个 CI 的错误，看上去是 npu 那边没有正确的处理 paddle::optional ，传了个空指针过去？～

… hack7_amsgrad

HydrogenSulfate

LGTM

HydrogenSulfate · 2024-12-03T08:23:20Z

@megemini 能麻烦再合一下develop分支吗，windows和Hygon-DCU-Test这两个挂了，应该不是PR的原因

luotao1 · 2024-12-03T08:55:20Z

windows和Hygon-DCU-Test这两个挂了，应该不是PR的原因

@HydrogenSulfate DCU已豁免，是单侧随机挂。windows我重跑了，根据其他开发者反馈，最近windows流水线网络不好。

… hack7_amsgrad

XiaoguangHu01

LGTM

HydrogenSulfate

LGTM

megemini added 12 commits August 29, 2024 18:39

[init] amsgrad

b45f2c4

[update] refer.h

640be9b

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

2028825

… hack7_amsgrad

[Add] amsgrad gpu

caf919a

[Add] amsgrad for adamw and fused

aa289ad

[Fix] adamw gpu kernel

106f817

[Update] fused adam kernel for gpu

fddb46a

[Update] xpu adam/adamw param list

d206442

[Update] tests for amsgrad

8cc9b5b

[Fix] moment2 max out settting values without amsgrad

eb5de54

[Update] unittest passed for adam and adamw

7aa9d60

[Update] unittest passed for merged and fused amda

96216e4

megemini requested review from LiYuRio, ForFishes and zhiqiu as code owners September 8, 2024 09:25

paddle-bot bot added the contributor External developers label Sep 8, 2024

luotao1 added the PaddlePaddle Hackathon label Sep 9, 2024

luotao1 assigned luotao1 and HydrogenSulfate Sep 9, 2024

PaddlePaddle locked and limited conversation to collaborators Sep 9, 2024

PaddlePaddle unlocked this conversation Sep 9, 2024

HydrogenSulfate reviewed Sep 9, 2024

View reviewed changes

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

7398a2f

… hack7_amsgrad

megemini added 3 commits September 10, 2024 23:45

[Update] make moment2_max optional

98abe71

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

e159b70

… hack7_amsgrad

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

4564d32

… hack7_amsgrad

megemini added 2 commits November 8, 2024 13:36

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

f17d737

… hack7_amsgrad

[Update] codestyle

af27337

megemini dismissed stale reviews from phlrain, heavengate, and zyfncg via af27337 November 8, 2024 05:39

[Update] npu test rtol adamw

d7bb19a

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

c8f27cd

… hack7_amsgrad

megemini added 3 commits November 27, 2024 15:26

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

45337e6

… hack7_amsgrad

[Update] xpu amsgrad raise errors

ee8d94f

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

4b8c9dc

… hack7_amsgrad

megemini mentioned this pull request Nov 27, 2024

【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad PaddlePaddle/PaddleCustomDevice#1484

Merged

megemini added 2 commits November 27, 2024 20:26

[Fix] not test xpu amsgrad

40ca555

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

0c15f55

… hack7_amsgrad

zyfncg approved these changes Dec 2, 2024

View reviewed changes

phlrain self-requested a review December 2, 2024 07:34

phlrain approved these changes Dec 2, 2024

View reviewed changes

heavengate approved these changes Dec 2, 2024

View reviewed changes

HydrogenSulfate approved these changes Dec 3, 2024

View reviewed changes

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

83c433b

… hack7_amsgrad

XiaoguangHu01 approved these changes Dec 4, 2024

View reviewed changes

luotao1 merged commit d774b83 into PaddlePaddle:develop Dec 4, 2024
28 checks passed

HydrogenSulfate reviewed Dec 4, 2024

View reviewed changes

This was referenced Dec 4, 2024

[Fea] Support amsgrad in Adam/AdamW PaddlePaddle/PaddleScience#1033

Merged

adapt code to amsgrad supported adamw PaddlePaddle/PaddleNLP#9568

Merged

luotao1 changed the title ~~【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad -part~~ 【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad Dec 4, 2024

megemini mentioned this pull request Jan 14, 2025

2024下半年飞桨开源之星评选-信息征集 PaddlePaddle/community#1043

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad #68079

【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad #68079

megemini commented Sep 8, 2024 •

edited

Loading

paddle-bot bot commented Sep 8, 2024

HydrogenSulfate left a comment

HydrogenSulfate Sep 9, 2024

HydrogenSulfate Sep 9, 2024

HydrogenSulfate Sep 9, 2024

HydrogenSulfate Sep 9, 2024

HydrogenSulfate Sep 9, 2024

megemini Sep 9, 2024

HydrogenSulfate Sep 9, 2024

HydrogenSulfate Sep 9, 2024

megemini commented Sep 9, 2024

HydrogenSulfate commented Sep 9, 2024 •

edited

Loading

HydrogenSulfate commented Sep 9, 2024

HydrogenSulfate commented Nov 26, 2024

megemini commented Nov 27, 2024

HydrogenSulfate left a comment

HydrogenSulfate commented Dec 3, 2024

luotao1 commented Dec 3, 2024

XiaoguangHu01 left a comment

HydrogenSulfate left a comment

		Eigen::Map<Eigen::Array<T, 1, Eigen::Dynamic>> moment2_max_out{
		moment2_max_out_, static_cast<Eigen::Index>(numel)};

【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad #68079

【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad #68079

Conversation

megemini commented Sep 8, 2024 • edited Loading

PR Category

PR Types

Description

Update 20240908

paddle-bot bot commented Sep 8, 2024

HydrogenSulfate left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

megemini commented Sep 9, 2024

HydrogenSulfate commented Sep 9, 2024 • edited Loading

HydrogenSulfate commented Sep 9, 2024

HydrogenSulfate commented Nov 26, 2024

megemini commented Nov 27, 2024

HydrogenSulfate left a comment

Choose a reason for hiding this comment

HydrogenSulfate commented Dec 3, 2024

luotao1 commented Dec 3, 2024

XiaoguangHu01 left a comment

Choose a reason for hiding this comment

HydrogenSulfate left a comment

Choose a reason for hiding this comment

megemini commented Sep 8, 2024 •

edited

Loading

HydrogenSulfate commented Sep 9, 2024 •

edited

Loading