add lamb (PaddlePaddle#155)

graphcore · Sep 13, 2021 · cdc0287 · cdc0287
1 parent 79cccf7
commit cdc0287
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 19 deletions.
diff --git a/paddle/fluid/framework/ipu/common.h b/paddle/fluid/framework/ipu/common.h
@@ -24,8 +24,6 @@ static constexpr const char *sIpuIndexAttr = "ipu_index";
 static constexpr const char *sIpuStageAttr = "ipu_stage";
 static constexpr const char *sOpIdentifyIdAttr = "op_identify_id";
 
-using IdToInfo = std::pair<popart::TensorId, popart::TensorInfo>;
-
 }  // namespace ipu
 }  // namespace framework
 }  // namespace paddle
diff --git a/paddle/fluid/framework/ipu/ipu_executor.cc b/paddle/fluid/framework/ipu/ipu_executor.cc
@@ -150,9 +150,7 @@ void Executor::SetWeightsIO() {
   auto pre_post_fix = GetOptPrePostfix(opt_type);
   for (const auto &weight_id : weights_) {
     for (const auto &pair : pre_post_fix) {
-      // TODO(xiaobingw): add more optimizer support
-      // only support adam/sgd currently, skip opt state if not adam
-      if (opt_type != OptimizerType::Adam && opt_type != OptimizerType::SGD) {
+      if (!IsOptimizerSupported(opt_type)) {
         continue;
       }
 

diff --git a/paddle/fluid/framework/ipu/ipu_optimizer.cc b/paddle/fluid/framework/ipu/ipu_optimizer.cc
@@ -43,18 +43,23 @@ OptimizerType OptTypeStr2Enum(const std::string type) {
     return OptimizerType::SGD;
   } else if (type == "adam") {
     return OptimizerType::Adam;
+  } else if (type == "lamb") {
+    return OptimizerType::Lamb;
   } else {
     return OptimizerType::Undefined;
   }
 }
 
 std::string OptTypeEnum2Str(OptimizerType type) {
-  if (type == OptimizerType::SGD) {
-    return "sgd";
-  } else if (type == OptimizerType::Adam) {
-    return "adam";
-  } else {
-    return "undefined";
+  switch (type) {
+    case OptimizerType::SGD:
+      return "sgd";
+    case OptimizerType::Adam:
+      return "adam";
+    case OptimizerType::Lamb:
+      return "lamb";
+    default:
+      return "undefined";
   }
 }
 
@@ -86,25 +91,56 @@ std::unique_ptr<popart::Optimizer> GetPopartOptimizer(
         popart::DataType::FLOAT, popart::DataType::FLOAT,
         popart::DataType::FLOAT);
     return optimizer;
+  } else if (opt_type == OptimizerType::Lamb) {
+    auto optimizer = std::make_unique<popart::Adam>(
+        popart::OptimizerValue(opt_meta_info.GetLR(), false),
+        popart::OptimizerValue(opt_meta_info.GetAttr("weight_decay"), false),
+        popart::OptimizerValue(opt_meta_info.GetAttr("beta1"), false),
+        popart::OptimizerValue(opt_meta_info.GetAttr("beta2"), false),
+        popart::OptimizerValue(opt_meta_info.GetAttr("epsilon"), false),
+        popart::OptimizerValue(popart::Adam::getUnsetLossScaling()),
+        popart::AdamMode::Lamb, popart::WeightDecayMode::Decay,
+        popart::DataType::FLOAT, popart::DataType::FLOAT,
+        popart::DataType::FLOAT);
+    return optimizer;
   } else {
     PADDLE_THROW(platform::errors::Unimplemented(
         "Optimizer %s is not implemented now.", OptTypeEnum2Str(opt_type)));
   }
 }
 
+bool IsOptimizerSupported(OptimizerType type) {
+  switch (type) {
+    case OptimizerType::SGD:
+    case OptimizerType::Adam:
+    case OptimizerType::Lamb:
+      return true;
+    default:
+      return false;
+  }
+}
+
 std::vector<std::pair<std::string, std::string>> GetOptPrePostfix(
     OptimizerType opt_type) {
   // format: {popart_tensor_id, paddle_tensor_id}, ...
   std::vector<std::pair<std::string, std::string>> pre_post_fix;
 
-  pre_post_fix.push_back(std::make_pair("", ""));
-  if (opt_type == OptimizerType::SGD) {
-  } else if (opt_type == OptimizerType::Adam) {
-    pre_post_fix.push_back(std::make_pair("Accl1___", "_moment1_0"));
-    pre_post_fix.push_back(std::make_pair("Accl2___", "_moment2_0"));
-    pre_post_fix.push_back(std::make_pair("Step___", "_beta1_pow_acc_0"));
-  } else {
+  switch (opt_type) {
+    case OptimizerType::SGD:
+      pre_post_fix.push_back(std::make_pair("", ""));
+      break;
+    case OptimizerType::Adam:
+    case OptimizerType::Lamb:
+      pre_post_fix.push_back(std::make_pair("", ""));
+      pre_post_fix.push_back(std::make_pair("Accl1___", "_moment1_0"));
+      pre_post_fix.push_back(std::make_pair("Accl2___", "_moment2_0"));
+      pre_post_fix.push_back(std::make_pair("Step___", "_beta1_pow_acc_0"));
+      break;
+    default:
+      pre_post_fix.push_back(std::make_pair("", ""));
+      break;
   }
+
   return pre_post_fix;
 }
 

diff --git a/paddle/fluid/framework/ipu/ipu_optimizer.h b/paddle/fluid/framework/ipu/ipu_optimizer.h
@@ -24,7 +24,7 @@ namespace paddle {
 namespace framework {
 namespace ipu {
 
-enum class OptimizerType { SGD = 0, Adam, Undefined };
+enum class OptimizerType { SGD = 0, Adam, Lamb, Undefined };
 
 class OptmizerMetaInfo {
  public:
@@ -67,6 +67,8 @@ std::string OptTypeEnum2Str(OptimizerType type);
 std::unique_ptr<popart::Optimizer> GetPopartOptimizer(
     const OptmizerMetaInfo &info);
 
+bool IsOptimizerSupported(OptimizerType type);
+
 std::vector<std::pair<std::string, std::string>> GetOptPrePostfix(
     OptimizerType type);
 

diff --git a/python/paddle/fluid/tests/unittests/ipu/test_save_load.py b/python/paddle/fluid/tests/unittests/ipu/test_save_load.py
@@ -81,6 +81,9 @@ def _test_base(self, save_otherwise_load):
                         elif self.attrs['opt_type'] == 'adam':
                             adam = paddle.optimizer.Adam(learning_rate=1e-2)
                             adam.minimize(loss)
+                        elif self.attrs['opt_type'] == 'lamb':
+                            lamb = paddle.optimizer.Lamb(learning_rate=1e-2)
+                            lamb.minimize(loss)
                     fetch_list = [loss.name]
 
                 place = paddle.IPUPlace()
@@ -137,5 +140,14 @@ def set_attrs(self):
         self.attrs['opt_type'] = 'adam'
 
 
+class TestLamb(TestBase):
+    def set_attrs(self):
+        self.attrs = {}
+        self.attrs['steps'] = 100
+        self.attrs['save_at_step'] = 20
+        self.attrs['is_training'] = True
+        self.attrs['opt_type'] = 'lamb'
+
+
 if __name__ == "__main__":
     unittest.main()