add adamw for deeplink ext

csrc/extensions.cpp

@@ -39,6 +39,15 @@ at::IntArrayRef optionalIntArrayToIntArrayRefOrDefault(
 
 }  // namespace
 
+auto extAdamW(at::Tensor& param, at::Tensor& exp_avg, at::Tensor& exp_avg_sq,
+              at::Tensor& max_exp_avg_sq, at::Tensor& grad, float lr,
+              float beta1, float beta2, float epsilon, float weight_decay,
+              int64_t step, bool amsgrad) {
+  // the diopiAdamW func has no "maximize" param
+  callDiopi(diopiAdamW, param, grad, exp_avg, exp_avg_sq, max_exp_avg_sq, lr,
+            beta1, beta2, epsilon, weight_decay, step, amsgrad);
+}
+
 auto extRmsNorm(const at::Tensor& input,
                 const OptionalIntArray& normalized_shape,
                 const at::Tensor& weight, const at::Tensor& bias, double eps) {
@@ -249,6 +258,9 @@ void extRotaryEmb(at::Tensor& q, const at::Tensor& cos, const at::Tensor& sin) {
 //   否则不注册, 等到 python 层处理.
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  if (&diopiAdamW != nullptr) {
+    m.def("adamw", &extAdamW, "deeplink ext_adamw");
+  }
   if (&diopiRMSNorm != nullptr) {  // Check if weak symbol defined
     m.def("rms_norm", &extRmsNorm, "deeplink ext_rms_norm");
     m.def("rms_norm_lightllm", &extRmsNormLightllm,

deeplink_ext/ascend_speed/__init__.py

@@ -0,0 +1,3 @@
+from .adamw import adamw
+
+all = ["adamw"]

deeplink_ext/ascend_speed/adamw.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2024, DeepLink.
+
+from typing import Optional, Union, List
+import torch
+import deeplink_ext.cpp_extensions as ext
+
+assert hasattr(ext, "adamw")
+
+
+def adamw(
+    params: List[torch.Tensor],
+    grads: List[torch.Tensor],
+    exp_avgs: List[torch.Tensor],
+    exp_avg_sqs: List[torch.Tensor],
+    max_exp_avg_sqs: List[torch.Tensor],
+    state_steps: List[int],
+    *,
+    amsgrad: bool,
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    maximize: bool,
+    norm_coeff_scale: float
+):
+    r"""Functional API that performs AdamW algorithm computation.
+    See :class:`~torch.optim.AdamW` for details.
+    """
+
+    assert maximize == False, "ascend diopiAdamW only support False 'maximize'."
+    assert amsgrad == False, "ascend diopiAdamW only support False 'amsgrad'."
+
+    for i, param in enumerate(params):
+        if norm_coeff_scale is not None:
+            grad = grads[i].float() * norm_coeff_scale
+        else:
+            grad = grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        step = state_steps[i]
+        if not max_exp_avg_sqs:
+            max_exp_avg_sq = torch.Tensor().cuda()
+        else:
+            max_exp_avg_sq = max_exp_avg_sqs[i]
+        ext.adamw(
+            param,
+            exp_avg,
+            exp_avg_sq,
+            max_exp_avg_sq,
+            grad,
+            lr,
+            beta1,
+            beta2,
+            eps,
+            weight_decay,
+            step,
+            amsgrad,
+        )
+    return params, exp_avgs, exp_avg_sqs