support Huggingface's adamw (#3318)

* add weight decay mode to support both pytorch and huggingface's adamw

orttraining/orttraining/core/graph/gradient_schema_defs.cc

@@ -544,6 +544,13 @@ void RegisterGradientSchemas() {
           "Compute unbiased 1st and 2nd momentums.",
           AttributeProto::INT,
           static_cast<int64_t>(1))
+      .Attr(
+          "weight_decay_mode",
+          "Modes for applying weight decay, "
+          "0 means applying decay before weight update, "
+          "1 means applying decay after weight update.",
+          AttributeProto::INT,
+          static_cast<int64_t>(0))
       .TypeConstraint(
           "T1",
           {"tensor(float16)", "tensor(float)", "tensor(double)"},

orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h

@@ -11,7 +11,12 @@ namespace training {
 class AdamOptimizerBuilder final : public OptimizerBuilder {
  public:
   AdamOptimizerBuilder() : OptimizerBuilder("AdamOptimizer",
-                                            {"alpha", "beta", "lambda", "epsilon", "do_bias_correction"}) {}
+                                            {"alpha",
+                                             "beta",
+                                             "lambda",
+                                             "epsilon",
+                                             "do_bias_correction",
+                                             "weight_decay_mode"}) {}
 
   virtual Status Build(
       const std::vector<ArgDef>& weight_argdefs,

orttraining/orttraining/models/bert/main.cc

@@ -141,6 +141,11 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
         "Default is false, which means no bias correction. "
         "Use true to enable bias correction.",
         cxxopts::value<bool>()->default_value("false"))
+      ("weight_decay_mode",
+        "Chooses the weight decay mode for Adam optimizer "
+        "Default is 0, which does weight decay before updating weight. "
+        "Use 1 to do weight decay after updating weight.",
+        cxxopts::value<int64_t>()->default_value("0"))
       ("ratio_min", "Lamb min ratio parameter", cxxopts::value<float>()->default_value("0.05"))
       ("ratio_max", "Lamb max ratio parameter", cxxopts::value<float>()->default_value("5.0"))
       ("cuda_mem_limit_in_gb", "Max cuda memory ort can use, in GB", cxxopts::value<float>()->default_value("-1.0"))
@@ -317,10 +322,12 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
     float beta = flags["beta"].as<float>();
     float lambda = flags["lambda"].as<float>();
     float epsilon = flags["epsilon"].as<float>();
+    int64_t weight_decay_mode = flags["weight_decay_mode"].as<int64_t>();
     float ratio_min = flags["ratio_min"].as<float>();
     float ratio_max = flags["ratio_max"].as<float>();
     ORT_RETURN_IF_NOT(alpha >= 0.f && alpha <= 1.f, "alpha is not in valid range [0.0, 1.0]");
     ORT_RETURN_IF_NOT(beta >= 0.f && beta <= 1.f, "alpha is not in valid range [0.0, 1.0]");
+    ORT_RETURN_IF_NOT(weight_decay_mode == 0 || weight_decay_mode == 1, "Only 0 and 1 are supported for weight decay mode.");
     ORT_RETURN_IF_NOT(epsilon >= 0.f, "epsilon should be non-negative.");
     ORT_RETURN_IF_NOT(epsilon >= 0.f, "epsilon should be non-negative.");
     ORT_RETURN_IF_NOT(ratio_min >= 0.f, "ratio_min should be non-negative.");
@@ -350,7 +357,8 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet
     // Optimizer's int attributes.
     params.optimizer_int_attributes = [=](const std::string& /*weight*/) {
       return std::unordered_map<std::string, int64_t>{
-          {"do_bias_correction", do_bias_correction ? static_cast<int64_t>(1) : static_cast<int64_t>(0)}
+          {"do_bias_correction", do_bias_correction ? static_cast<int64_t>(1) : static_cast<int64_t>(0)},
+          {"weight_decay_mode", weight_decay_mode}
       };
     };
 

orttraining/orttraining/training_ops/cpu/optimizer/optimizers.cc

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#include "optimizers.h"
+#include "orttraining/training_ops/cpu/optimizer/optimizers.h"
 
 #include "core/framework/op_kernel.h"
 #include "core/providers/common.h"
@@ -70,20 +70,53 @@ Status AdamOptimizer<T>::Compute(OpKernelContext* ctx) const {
   const float beta_correction = do_bias_correction_ ?
     compute_bias_correction_coefficient(beta_, step) : 1.f;
 
-  // Compute weight update.
-  const auto& denom = (MakeEigenArrayMap<T>(NM2) / beta_correction).sqrt() + epsilon_;
-  const auto& update = ( (MakeEigenArrayMap<T>(NM1) / alpha_correction) / denom) + (lambda_ * MakeEigenArrayMap<T>(W));
-  const auto& delta = -eta * update;
+  // Currently two modes of Adamw are supported:
+  // Mode 0: Pytorch https://pytorch.org/docs/stable/_modules/torch/optim/adamw.html#AdamW,
+  //         bias correction is applied on m and v individually,
+  //         weight decay is applied before weight is updated.
+  // Mode 1: Huggingface https://huggingface.co/transformers/_modules/transformers/optimization.html#AdamW.,
+  //         bias correction is applied on learning rate,
+  //         weight decay is applied after weight is updated.
+  if(weight_decay_mode_ == 0) {
+    // Compute weight update.
+    const auto& denom = (MakeEigenArrayMap<T>(NM2) / beta_correction).sqrt() + epsilon_;
+    const auto& update = ( (MakeEigenArrayMap<T>(NM1) / alpha_correction) / denom) + (lambda_ * MakeEigenArrayMap<T>(W));
+    const auto& delta = -eta * update;
 
-  // Weight, gradient, and step update.
-  if (NG != nullptr) {
-    MakeEigenArrayMap<T>(*NG) = delta;
+    // Weight, gradient, and step update.
+    if (NG != nullptr) {
+      MakeEigenArrayMap<T>(*NG) = delta;
+    }
+    if (NW != nullptr) {
+      MakeEigenArrayMap<T>(*NW) = MakeEigenArrayMap<T>(W) + delta;
+    }
   }
-  if (NW != nullptr) {
-    MakeEigenArrayMap<T>(*NW) = MakeEigenArrayMap<T>(W) + delta;
+  else if (weight_decay_mode_ == 1) {
+    const auto& denom = MakeEigenArrayMap<T>(NM2).sqrt() + epsilon_;
+    const auto& step_size = eta * std::sqrt(beta_correction) / alpha_correction;
+
+    // Huggingface updates weights in the following logic:
+    // param' = param - step_size * m1o / denom
+    // param_out = param' - original_lr * lambda * param'
+    // then param_out = param - step_size * m1o / denom - original_lr * lambda * (param - step_size * m1o / denom)
+    // so delta = -step_size * m1o / denom - original_lr * lambda * (param - step_size * m1o / denom)
+    const auto& delta = -step_size * MakeEigenArrayMap<T>(NM1) / denom
+                        - eta * lambda_ * (MakeEigenArrayMap<T>(W) - step_size * MakeEigenArrayMap<T>(NM1) / denom);
+
+    // Weight, gradient, and step update.
+    if (NG != nullptr) {
+      MakeEigenArrayMap<T>(*NG) = delta;
+    }
+    if (NW != nullptr) {
+      MakeEigenArrayMap<T>(*NW) = MakeEigenArrayMap<T>(W) + delta;
+    }
   }
+  else {
+    // Shouldn't reach here
+    ORT_THROW("Unsupported Adamw optimizer mode.");
+  }
+
   *NS.template MutableData<int64_t>() = step + 1;
-
   return Status::OK();
 }
 

orttraining/orttraining/training_ops/cpu/optimizer/optimizers.h

@@ -3,7 +3,7 @@
 
 #pragma once
 
-#include "common.h"
+#include "orttraining/training_ops/cpu/optimizer/common.h"
 #include "core/common/common.h"
 #include "core/framework/op_kernel.h"
 
@@ -30,11 +30,12 @@ class AdamOptimizer final : public OpKernel {
     ORT_ENFORCE(beta_ >= 0);
     ORT_ENFORCE(lambda_ >= 0);
     ORT_ENFORCE(epsilon_ >= 0);
-
     int64_t tmp_flag = static_cast<int64_t>(0);
     ORT_ENFORCE(info.GetAttr<int64_t>("do_bias_correction", &tmp_flag).IsOK(), "Missing/Invalid do_bias_correction");
     ORT_ENFORCE(tmp_flag == 0 || tmp_flag == 1, "do_bias_correction must be either 0 or 1.");
     do_bias_correction_ = tmp_flag != 0 ? true : false;
+    info.GetAttrOrDefault("weight_decay_mode", &weight_decay_mode_, static_cast<int64_t>(0));
+    ORT_ENFORCE(weight_decay_mode_ == 0 || weight_decay_mode_ == 1, "Only 0 and 1 are supported for weight decay mode.");
   }
 
   Status Compute(OpKernelContext* context) const override;
@@ -45,6 +46,7 @@ class AdamOptimizer final : public OpKernel {
   float lambda_;
   float epsilon_;
   bool do_bias_correction_;
+  int64_t weight_decay_mode_;
 };
 }  // namespace contrib
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cuda/optimizer/adam.cc

@@ -4,8 +4,8 @@
 #include "core/providers/cuda/cuda_allocator.h"
 #include "core/providers/cuda/reduction/reduction_functions.h"
 #include "core/providers/cuda/math/binary_elementwise_ops.h"
-#include "common.h"
-#include "adam.h"
+#include "orttraining/training_ops/cuda/optimizer/common.h"
+#include "orttraining/training_ops/cuda/optimizer/adam.h"
 
 namespace onnxruntime {
 namespace cuda {
@@ -136,6 +136,7 @@ Status AdamOptimizer<T1, T2, T3, T4, T_GRAD, T_GRAD_NORM>::ComputeInternal(OpKer
       ToCudaType<T4>::FromFloat(lambda_),
       ToCudaType<T4>::FromFloat(epsilon_),
       do_bias_correction_,
+      weight_decay_mode_,
       reinterpret_cast<CudaT4*>(NM1.template MutableData<T4>()),
       reinterpret_cast<CudaT4*>(NM2.template MutableData<T4>()),
       NW != nullptr ? reinterpret_cast<CudaT3*>(NW->template MutableData<T3>()) : nullptr,

orttraining/orttraining/training_ops/cuda/optimizer/adam.cu

@@ -4,13 +4,13 @@
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "orttraining/training_ops/cuda/optimizer/common.cuh"
-#include "adam.h"
-#include "common.h"
+#include "orttraining/training_ops/cuda/optimizer/adam.h"
+#include "orttraining/training_ops/cuda/optimizer/common.h"
 
 namespace onnxruntime {
 namespace cuda {
 template <typename T1, typename T3, typename T4, typename T_GRAD, typename T_GRAD_NORM>
-__global__ void _AdamOptimizer(
+__global__ void _AdamOptimizer_mode0(
     const T1* eta,
     const T3* weights,
     const T_GRAD* grads,
@@ -49,6 +49,7 @@ __global__ void _AdamOptimizer(
   // Compute weight update.
   const T4 denom = _Sqrt(m2o_corrected) + epsilon;
   const T4 update = (m1o_corrected / denom) + (lambda * T4(weights[id]));
+
   const T4 delta = -T4(*eta) * update;
 
   // Compute the new gradient.
@@ -69,6 +70,71 @@ __global__ void _AdamOptimizer(
   moment_2_out[id] = m2o;
 }
 
+template <typename T1, typename T3, typename T4, typename T_GRAD, typename T_GRAD_NORM>
+__global__ void _AdamOptimizer_mode1(
+    const T1* eta,
+    const T3* weights,
+    const T_GRAD* grads,
+    const T4* moment_1,
+    const T4* moment_2,
+    const T3* loss_scale,
+    const T_GRAD_NORM* grad_norm,
+    const T4 alpha,
+    const T4 beta,
+    const T4 lambda,
+    const T4 epsilon,
+    const T4 alpha_correction,
+    const T4 beta_correction,
+    T4* moment_1_out,
+    T4* moment_2_out,
+    T3* weights_out,
+    T_GRAD* grads_out,
+    half* fp16_weights_out,
+    CUDA_LONG N) {
+  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
+  const T4 actual_scale = _ComputeGradScale<T3, T_GRAD_NORM, T4>(loss_scale, grad_norm);
+
+  // Gradient scaling/clipping.
+  const T4 g = T4(grads[id]) / actual_scale;
+  // A shared constant.
+  const T4 one = T4(1.0f);
+
+  // Compute exponentially-averaged historical gradient.
+  const T4 m1o = alpha * moment_1[id] + (one - alpha) * g;
+
+  // Compute exponentially-averaged historical squared gradient.
+  const T4 m2o = beta * moment_2[id] + (one - beta) * g * g;
+
+  const T4 denom = _Sqrt(m2o) + epsilon;
+
+  // Apply bias correction terms on learning rate
+  const T4 step_size = T4(*eta) * _Sqrt(beta_correction) / alpha_correction;
+
+  // Huggingface updates weights in the following logic:
+  // param' = param - step_size * m1o / denom
+  // param_out = param' - original_lr * lambda * param'
+  // then param_out = param - step_size * m1o / denom - original_lr * lambda * (param - step_size * m1o / denom)
+  // so delta = -step_size * m1o / denom - original_lr * lambda * (param - step_size * m1o / denom)
+  const T4 delta = -step_size * m1o / denom - T4(*eta) * lambda * (T4(weights[id]) - step_size * m1o / denom);
+
+  // Compute the new gradient.
+  if (grads_out) {
+    grads_out[id] = T_GRAD(delta);
+  }
+  
+  // Compute the new weight.
+  if (weights_out) {
+    weights_out[id] = weights[id] + T3(delta);
+
+    if (fp16_weights_out) {
+      fp16_weights_out[id] = static_cast<half>(weights_out[id]);
+    }
+  }
+
+  moment_1_out[id] = m1o;
+  moment_2_out[id] = m2o;
+}
+
 template <typename T1, typename T2, typename T3, typename T4, typename T_GRAD, typename T_GRAD_NORM>
 void AdamOptimizerImpl(
     const T1* eta,
@@ -84,6 +150,7 @@ void AdamOptimizerImpl(
     const T4 lambda,
     const T4 epsilon,
     const bool do_bias_correction,
+    const int64_t weight_decay_mode,
     T4* moment_1_out,
     T4* moment_2_out,
     T3* weights_out,
@@ -97,7 +164,16 @@ void AdamOptimizerImpl(
     onnxruntime::contrib::compute_bias_correction_coefficient(alpha, update_count) : T4(1.f);
   const T4 beta_correction = do_bias_correction ?
     onnxruntime::contrib::compute_bias_correction_coefficient(beta, update_count) : T4(1.f);
-  _AdamOptimizer<T1, T3, T4, T_GRAD, T_GRAD_NORM><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+  
+  // Currently two modes of Adamw are supported:
+  // Mode 0: Pytorch https://pytorch.org/docs/stable/_modules/torch/optim/adamw.html#AdamW,
+  //         bias correction is applied on m and v individually,
+  //         weight decay is applied before weight is updated.
+  // Mode 1: Huggingface https://huggingface.co/transformers/_modules/transformers/optimization.html#AdamW.,
+  //         bias correction is applied on learning rate,
+  //         weight decay is applied after weight is updated.
+  if (weight_decay_mode == 0) {
+    _AdamOptimizer_mode0<T1, T3, T4, T_GRAD, T_GRAD_NORM><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
       eta,
       weights,
       grads,
@@ -117,6 +193,33 @@ void AdamOptimizerImpl(
       grads_out,
       fp16_weights_out,
       N);
+  }
+  else if (weight_decay_mode == 1) {
+    _AdamOptimizer_mode1<T1, T3, T4, T_GRAD, T_GRAD_NORM><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+      eta,
+      weights,
+      grads,
+      moment_1,
+      moment_2,
+      loss_scale,
+      grad_norm,
+      alpha,
+      beta,
+      lambda,
+      epsilon,
+      alpha_correction,
+      beta_correction,
+      moment_1_out,
+      moment_2_out,
+      weights_out,
+      grads_out,
+      fp16_weights_out,
+      N);
+  }
+  else {
+    // Shouldn't reach here
+    ORT_THROW("Unsupported Adamw optimizer mode.");
+  }
 }
 
 #define SPECIALIZED_AdamOptimizerImpl(T1, T2, T3, T4, T_GRAD, T_GRAD_NORM) \
@@ -134,6 +237,7 @@ void AdamOptimizerImpl(
       const T4 lambda,                                                     \
       const T4 epsilon,                                                    \
       const bool do_bias_correction,                                       \
+      const int64_t weight_decay_mode,                                     \
       T4* moment_1_out,                                                    \
       T4* moment_2_out,                                                    \
       T3* weights_out,                                                     \

orttraining/orttraining/training_ops/cuda/optimizer/adam.h

@@ -24,6 +24,7 @@ void AdamOptimizerImpl(
     const T4 lambda,
     const T4 epsilon,
     const bool do_bias_correction,
+    const int64_t weight_decay_mode,
     T4* moment_1_out,
     T4* moment_2_out,
     T3* weights_out,
@@ -44,6 +45,7 @@ class AdamOptimizer final : public CudaKernel {
     ORT_ENFORCE(info.GetAttr<int64_t>("do_bias_correction", &tmp_flag).IsOK(), "Missing/Invalid do_bias_correction");
     ORT_ENFORCE(tmp_flag == 0 || tmp_flag == 1, "do_bias_correction must be either 0 or 1.");
     do_bias_correction_ = tmp_flag != 0 ? true : false;
+    info.GetAttrOrDefault("weight_decay_mode", &weight_decay_mode_, static_cast<int64_t>(0));
   }
 
   Status ComputeInternal(OpKernelContext* context) const override;
@@ -53,7 +55,8 @@ class AdamOptimizer final : public CudaKernel {
   float beta_;
   float lambda_;
   float epsilon_;
-  bool do_bias_correction_; 
+  bool do_bias_correction_;
+  int64_t weight_decay_mode_;
 };
 
 }  // namespace cuda

orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc

@@ -5,8 +5,8 @@
 #include "core/providers/cuda/cuda_allocator.h"
 #include "core/providers/cuda/reduction/reduction_functions.h"
 #include "core/providers/cuda/math/binary_elementwise_ops.h"
-#include "common.h"
-#include "lamb.h"
+#include "orttraining/training_ops/cuda/optimizer/common.h"
+#include "orttraining/training_ops/cuda/optimizer/lamb.h"
 
 namespace onnxruntime {
 namespace cuda {

orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu

@@ -6,7 +6,7 @@
 #include "core/providers/cuda/atomic/common.cuh"
 #include "orttraining/training_ops/cuda/math/isfinite.cuh"
 #include "orttraining/training_ops/cuda/optimizer/common.cuh"
-#include "lamb.h"
+#include "orttraining/training_ops/cuda/optimizer/lamb.h"
 namespace onnxruntime {
 namespace cuda {
 template <typename T1, typename T2, typename T3>