Access fused_adamw operator

DeepLink-org · Aug 29, 2024 · e70717e · e70717e
1 parent 67467e9
commit e70717e
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 3 deletions.
diff --git a/diopi_test/python/conformance/customized_test.py b/diopi_test/python/conformance/customized_test.py
@@ -177,6 +177,7 @@ def fused_adamw(
             weight_decay=weight_decay,
             eps=eps,
             maximize=maximize,
+            fused=True,
         )
         return params, exp_avgs, exp_avg_sqs, max_exp_avg_sqs
 

diff --git a/diopi_test/python/conformance/diopi_functions.py b/diopi_test/python/conformance/diopi_functions.py
@@ -3179,7 +3179,7 @@ def fused_adamw(
     for state_step in state_steps:
         c_state_steps.append(TensorP(state_step))
 
-    func = check_function("diopiFused_AdamW")
+    func = check_function("diopiFusedAdamW")
     ret = func(
         params[0].context(),
         list(c_params),

diff --git a/impl/torch/functions/functions.cpp b/impl/torch/functions/functions.cpp
@@ -2815,7 +2815,7 @@ diopiError_t diopiMeshGrid(diopiContextHandle_t ctx, diopiTensorHandle_t* outs,
     return diopiSuccess;
 }
 
-diopiError_t diopiFused_AdamW(diopiContextHandle_t ctx, diopiTensorHandle_t* params, diopiConstTensorHandle_t* grads, diopiTensorHandle_t* exp_avgs,
+diopiError_t diopiFusedAdamW(diopiContextHandle_t ctx, diopiTensorHandle_t* params, diopiConstTensorHandle_t* grads, diopiTensorHandle_t* exp_avgs,
                              diopiTensorHandle_t* exp_avg_sqs, diopiTensorHandle_t* max_exp_avg_sqs, diopiConstTensorHandle_t* state_steps, int64_t nums,
                              float lr, float beta1, float beta2, float eps, float weight_decay, bool amsgrad, bool maximize) {
     impl::aten::setCurStream(ctx);

diff --git a/proto/include/diopi/functions.h b/proto/include/diopi/functions.h
@@ -2469,7 +2469,7 @@ DIOPI_API diopiError_t diopiReciprocalInp(diopiContextHandle_t ctx, diopiTensorH
  * @param[in] amsgrad whether to use the AMSGrad variant of this algorithm from the paper `On the Convergence of Adam and Beyond`.
  * @param[in] maximize maximize the objective with respect to the params, instead of minimizing.
  */
-DIOPI_API diopiError_t diopiFused_AdamW(diopiContextHandle_t ctx, diopiTensorHandle_t* params, diopiConstTensorHandle_t* grads, diopiTensorHandle_t* exp_avgs,
+DIOPI_API diopiError_t diopiFusedAdamW(diopiContextHandle_t ctx, diopiTensorHandle_t* params, diopiConstTensorHandle_t* grads, diopiTensorHandle_t* exp_avgs,
                                        diopiTensorHandle_t* exp_avg_sqs, diopiTensorHandle_t* max_exp_avg_sqs, diopiConstTensorHandle_t* state_steps,
                                        int64_t nums, float lr, float beta1, float beta2, float eps, float weight_decay, bool amsgrad, bool maximize);