diff --git a/orttraining/orttraining/core/graph/gradient_schema_defs.cc b/orttraining/orttraining/core/graph/gradient_schema_defs.cc index 14804a07ff..8fde46e304 100644 --- a/orttraining/orttraining/core/graph/gradient_schema_defs.cc +++ b/orttraining/orttraining/core/graph/gradient_schema_defs.cc @@ -544,6 +544,13 @@ void RegisterGradientSchemas() { "Compute unbiased 1st and 2nd momentums.", AttributeProto::INT, static_cast(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(0)) .TypeConstraint( "T1", {"tensor(float16)", "tensor(float)", "tensor(double)"}, diff --git a/orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h b/orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h index 0e22a3c5cd..60077dc5ff 100644 --- a/orttraining/orttraining/core/graph/optimizer/adam_optimizer_builder.h +++ b/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& weight_argdefs, diff --git a/orttraining/orttraining/models/bert/main.cc b/orttraining/orttraining/models/bert/main.cc index 191d8b3460..c88417caf4 100644 --- a/orttraining/orttraining/models/bert/main.cc +++ b/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()->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()->default_value("0")) ("ratio_min", "Lamb min ratio parameter", cxxopts::value()->default_value("0.05")) ("ratio_max", "Lamb max ratio parameter", cxxopts::value()->default_value("5.0")) ("cuda_mem_limit_in_gb", "Max cuda memory ort can use, in GB", cxxopts::value()->default_value("-1.0")) @@ -317,10 +322,12 @@ Status ParseArguments(int argc, char* argv[], BertParameters& params, OrtParamet float beta = flags["beta"].as(); float lambda = flags["lambda"].as(); float epsilon = flags["epsilon"].as(); + int64_t weight_decay_mode = flags["weight_decay_mode"].as(); float ratio_min = flags["ratio_min"].as(); float ratio_max = flags["ratio_max"].as(); 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{ - {"do_bias_correction", do_bias_correction ? static_cast(1) : static_cast(0)} + {"do_bias_correction", do_bias_correction ? static_cast(1) : static_cast(0)}, + {"weight_decay_mode", weight_decay_mode} }; }; diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc index 70a1b74b5d..2e1cca4c8c 100644 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -1036,25 +1036,6 @@ TEST(GradientCheckerTest, SoftMaxGrad) { } } -TEST(OptimizerTest, SGDOptimizerTest) { - OpTester test("SGDOptimizer", 9, onnxruntime::kOnnxDomain); - test.AddInput("ETA", {}, {0.5f}); - test.AddInput("W", {3}, {1, 2, 3}); - test.AddInput("G", {3}, {4, 5, 6}); - test.AddOutput("W_New", {3}, {-1.f, -0.5f, 0.f}); - test.Run(); -} - -TEST(OptimizerTest, SGDOptimizerTest_Gradient) { - OpTester test("SGDOptimizer", 9, onnxruntime::kOnnxDomain); - test.AddInput("ETA", {}, {0.5f}); - test.AddInput("W", {3}, {1, 2, 3}); - test.AddInput("G", {3}, {4, 5, 6}); - test.AddMissingOptionalOutput(); - test.AddOutput("G_New", {3}, {-2.f, -2.5f, -3.f}); - test.Run(); -} - void TestSoftmaxCrossEntropyGrad(const TensorShape& input_shape, const std::string& reduction) { float max_error; GradientChecker gradient_checker; @@ -1129,123 +1110,6 @@ TEST(GradientCheckerTest, FastGeluGrad) { UnaryOpGradientTest("FastGelu", kMSDomain, 1); } -struct AdamOptimizerInputOutput { - AdamOptimizerInputOutput() { - eta_half.resize(eta.size()); - g_half.resize(g.size()); - m1_half.resize(m1.size()); - m2_half.resize(m2.size()); - w_half.resize(w.size()); - ConvertFloatToMLFloat16(eta.data(), eta_half.data(), int(eta.size())); - ConvertFloatToMLFloat16(g.data(), g_half.data(), int(g.size())); - ConvertFloatToMLFloat16(m1.data(), m1_half.data(), int(m1.size())); - ConvertFloatToMLFloat16(m2.data(), m2_half.data(), int(m2.size())); - ConvertFloatToMLFloat16(w.data(), w_half.data(), int(w.size())); - - m1_new_half.resize(m1_new.size()); - m2_new_half.resize(m2_new.size()); - w_new_half.resize(w_new.size()); - g_new_half.resize(g_new.size()); - ConvertFloatToMLFloat16(m1_new.data(), m1_new_half.data(), int(m1_new.size())); - ConvertFloatToMLFloat16(m2_new.data(), m2_new_half.data(), int(m2_new.size())); - ConvertFloatToMLFloat16(w_new.data(), w_new_half.data(), int(w_new.size())); - ConvertFloatToMLFloat16(g_new.data(), g_new_half.data(), int(g_new.size())); - } - - // Fp32 Inputs - std::vector eta = {0.5f}; - std::vector w = {1.0f, 2.0f, 3.0f}; - std::vector g = {4.0f, 5.0f, 6.0f}; - std::vector m1 = {0.1f, 0.2f, 0.3f}; - std::vector m2 = {0.4f, 0.5f, 0.6f}; - - // Fp16 Inputs - std::vector eta_half; - std::vector w_half; - std::vector g_half; - std::vector m1_half; - std::vector m2_half; - - // FP32 Outptus - std::vector m1_new = {0.49f, 0.68f, 0.87f}; - std::vector m2_new = {0.4156f, 0.5245f, 0.6354f}; - std::vector w_new = {0.6199609f, 1.5305318f, 2.4542853f}; - std::vector g_new = {-0.3800391f, -0.4694682f, -0.5457147f}; - - // FP16 Outptus - std::vector m1_new_half; - std::vector m2_new_half; - std::vector w_new_half; - std::vector g_new_half; -}; - -TEST(OptimizerTest, AdamOptimizerTest) { - OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); - AdamOptimizerInputOutput data; - - test.AddInput("ETA", {}, data.eta); - test.AddInput("Update_Count", {}, {3}); - test.AddInput("W", {3}, data.w); - test.AddInput("G", {3}, data.g); - test.AddInput("Moment_1", {3}, data.m1); - test.AddInput("Moment_2", {3}, data.m2); - - // Verify AdamOptimizer outputs - test.AddOutput("Update_Count_Out", {}, {4}); - test.AddOutput("Moment_1_Out", {3}, data.m1_new); - test.AddOutput("Moment_2_Out", {3}, data.m2_new); - test.AddOutput("W_Out", {3}, data.w_new); - - test.AddAttribute("do_bias_correction", static_cast(0)); - - test.Run(); -} - -TEST(OptimizerTest, AdamOptimizerTest_Gradient) { - OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); - AdamOptimizerInputOutput data; - - test.AddInput("ETA", {}, data.eta); - test.AddInput("Update_Count", {}, {3}); - test.AddInput("W", {3}, data.w); - test.AddInput("G", {3}, data.g); - test.AddInput("Moment_1", {3}, data.m1); - test.AddInput("Moment_2", {3}, data.m2); - - // Verify AdamOptimizer outputs - test.AddOutput("Update_Count_Out", {}, {4}); - test.AddOutput("Moment_1_Out", {3}, data.m1_new); - test.AddOutput("Moment_2_Out", {3}, data.m2_new); - test.AddMissingOptionalOutput(); - test.AddOutput("G_Out", {3}, data.g_new); - - test.AddAttribute("do_bias_correction", static_cast(0)); - - test.Run(); -} - -TEST(OptimizerTest, AdamBiasCorrection) { - OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); - AdamOptimizerInputOutput data; - - test.AddInput("ETA", {}, {1.f}); - test.AddInput("Update_Count", {}, {1}); - test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); - test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); - test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); - test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); - - test.AddOutput("Update_Count_Out", {}, {2}); - test.AddOutput("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f}); - test.AddOutput("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f}); - test.AddOutput("W_Out", {3}, {-1.4634f, -0.6416f, -1.2121f}); - - test.AddAttribute("do_bias_correction", static_cast(1)); - - test.Run(); -} - - TEST(GradientCheckerTest, GatherGrad) { float max_error; GradientChecker gradient_checker; @@ -1378,7 +1242,6 @@ void TestDropoutGradOp(float ratio, TensorShape& x_shape, bool default_ratio = t } #ifdef USE_CUDA - TEST(GradientCheckerTest, DISABLED_Dropout) { { //Ratio 0 @@ -1499,1124 +1362,6 @@ TEST(GradientCheckerTest, GatherNDGrad_int32_indice_unique_float_data_axis_2) { EXPECT_IS_TINY(max_error); } -TEST(OptimizerTest, AdamOptimizerMixPrecisionTest) { - OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); - AdamOptimizerInputOutput data; - - test.AddInput("ETA", {}, data.eta_half); - test.AddInput("Update_Count", {}, {3}); - test.AddInput("W", {3}, data.w); - test.AddInput("G", {3}, data.g_half); - test.AddInput("Moment_1", {3}, data.m1_half); - test.AddInput("Moment_2", {3}, data.m2_half); - - // Verify AdamOptimizer outputs - test.AddOutput("Update_Count_Out", {}, {4}); - test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); - test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); - test.AddOutput("W_Out", {3}, data.w_new); - - test.AddAttribute("do_bias_correction", static_cast(0)); - - test.Run(); -} - -TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_Test) { - OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); - AdamOptimizerInputOutput data; - - test.AddInput("ETA", {}, data.eta_half); - test.AddInput("Update_Count", {}, {3}); - test.AddInput("W", {3}, data.w); - test.AddInput("G", {3}, data.g_half); - test.AddInput("Moment_1", {3}, data.m1_half); - test.AddInput("Moment_2", {3}, data.m2_half); - test.AddInput("FP16_W", {3}, data.w_half); - - // Verify AdamOptimizer outputs - test.AddOutput("Update_Count_Out", {}, {4}); - test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); - test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); - test.AddOutput("W_Out", {3}, data.w_new); - test.AddMissingOptionalOutput(); - test.AddOutput("FP16_W_Out", {3}, data.w_new_half); - - test.AddAttribute("do_bias_correction", static_cast(0)); - - test.Run(); -} - -TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_SkipUpdate_Test) { - OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); - AdamOptimizerInputOutput data; - - test.AddInput("ETA", {}, data.eta_half); - test.AddInput("Update_Count", {}, {3}); - test.AddInput("W", {3}, data.w); - test.AddInput("G", {3}, data.g_half); - test.AddInput("Moment_1", {3}, data.m1_half); - test.AddInput("Moment_2", {3}, data.m2_half); - test.AddInput("FP16_W", {3}, data.w_half); - test.AddInput("loss_scale", {1}, {1.0f}); - // grad clipping should not take effect - test.AddInput("grad_norm", {1}, {0.01f}); - test.AddInput("DoUpdate", {1}, {false}); - - // Verify AdamOptimizer outputs - test.AddOutput("Update_Count_Out", {}, {3}); - test.AddOutput("Moment_1_Out", {3}, data.m1_half); - test.AddOutput("Moment_2_Out", {3}, data.m2_half); - test.AddOutput("W_Out", {3}, data.w); - test.AddMissingOptionalOutput(); - test.AddOutput("FP16_W_Out", {3}, data.w_half); - - test.AddAttribute("do_bias_correction", static_cast(0)); - - test.Run(); -} - -TEST(OptimizerTest, AdamOptimizerMixPrecisionTestFloatEta) { - OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); - AdamOptimizerInputOutput data; - - test.AddInput("ETA", {}, data.eta); - test.AddInput("Update_Count", {}, {3}); - test.AddInput("W", {3}, data.w); - test.AddInput("G", {3}, data.g_half); - test.AddInput("Moment_1", {3}, data.m1_half); - test.AddInput("Moment_2", {3}, data.m2_half); - - // Verify AdamOptimizer outputs - test.AddOutput("Update_Count_Out", {}, {4}); - test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); - test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); - test.AddOutput("W_Out", {3}, data.w_new); - - test.AddAttribute("do_bias_correction", static_cast(0)); - - test.Run(); -} - -TEST(OptimizerTest, AdamOptimizerMixPrecisionTest_Gradient) { - OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); - AdamOptimizerInputOutput data; - - test.AddInput("ETA", {}, data.eta); - test.AddInput("Update_Count", {}, {3}); - test.AddInput("W", {3}, data.w); - test.AddInput("G", {3}, data.g_half); - test.AddInput("Moment_1", {3}, data.m1_half); - test.AddInput("Moment_2", {3}, data.m2_half); - - // Verify AdamOptimizer outputs - test.AddOutput("Update_Count_Out", {}, {4}); - test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); - test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); - test.AddMissingOptionalOutput(); - test.AddOutput("G_Out", {3}, data.g_new_half); - - test.AddAttribute("do_bias_correction", static_cast(0)); - - test.Run(); -} - -// This helper function is a CPU-based LAMB optimizer -// implementation. It mainly focuses on readability. -void compute_lamb( - const std::vector shape, - /* weights */ const std::vector& w, - /* gradient */ const std::vector& g, - /* momentum */ const std::vector& m, - /* 2nd-order momentum */ const std::vector& v, - const float eta, - const float loss_scale, - const float g_norm, - const float lambda, - const float alpha, - const float beta, - const float epsilon, - /* updated weights */ std::vector& w_new, - /* updated gradients */ std::vector& g_new, - /* updated momentum */ std::vector& m_new, - /* updated 2nd-order momentum */ std::vector& v_new, - const int64_t step = 0, - const float ratio_min = -std::numeric_limits::infinity(), - const float ratio_max = std::numeric_limits::infinity()) { - // Element counts of all vector-typed arguments. - const int64_t size = std::accumulate( - shape.begin(), - shape.end(), - (int64_t)1, - std::multiplies()); - - // Buffer to store update direction. - std::vector r(size, 0.0f); - - float g_scale = loss_scale; - if (g_norm > loss_scale) { - g_scale *= g_norm / loss_scale; - } - - const float alpha_correction = step > 0 ? 1.f - std::pow(alpha, static_cast(step)) : 1.f; - const float beta_correction = step > 0 ? 1.f - std::pow(beta, static_cast(step)) : 1.f; - - // Compute new 1st-, 2nd-order momentums, and the update direction. - for (int i = 0; i < size; ++i) { - const float g_scaled = g[i] / g_scale; - m_new[i] = alpha * m[i] + (1.0f - alpha) * g_scaled; - v_new[i] = beta * v[i] + (1.0f - beta) * g_scaled * g_scaled; - const float m_new_tmp = m_new[i] / alpha_correction; - const float v_new_tmp = v_new[i] / beta_correction; - r[i] = lambda * w[i] + m_new_tmp / (std::sqrt(v_new_tmp) + epsilon); - } - - // Compute squared sum of all elements. Note that Eigen sqrt could lead to significant - // numerical error so we use std::sqrt. The std::inner_product produces wrong result - // when std::inner_product(r.begin(), r.end(), r.begin(), 0) so we just use a loop below. - float r_norm = 0.0f; - float w_norm = 0.0f; - for (int i = 0; i < size; ++i) { - r_norm += r[i] * r[i]; - w_norm += w[i] * w[i]; - } - - r_norm = std::sqrt(r_norm); - w_norm = std::sqrt(w_norm); - - float ratio = (w_norm != 0.0f && r_norm != 0.0f) ? w_norm / r_norm : 1.0f; - - if (ratio > ratio_max) { - ratio = ratio_max; - } - - if (ratio < ratio_min) { - ratio = ratio_min; - } - - ratio *= eta; - - // Compute the new weight. - for (int64_t i = 0; i < size; ++i) { - g_new[i] = -ratio * r[i]; - w_new[i] = w[i] + g_new[i]; - } -} - -template -void run_lamb_test_with_baseline( - const std::vector& shape, - const std::vector& eta, - const std::vector& w, - const std::vector& g, - const std::vector& m, - const std::vector& v, - const float alpha, - const float beta, - const float lambda, - const float epsilon, - const std::vector& w_new, - const std::vector& g_new, - const std::vector& m_new, - const std::vector& v_new, - const std::vector& w_half = {}, - const std::vector& w_new_half = {}, - const bool do_update = true, - const int64_t step = 0, - const float ratio_min = -std::numeric_limits::infinity(), - const float ratio_max = std::numeric_limits::infinity()) { - OpTester test("LambOptimizer", 9, onnxruntime::kOnnxDomain, true); - - test.AddInput("update_signal", {1}, {do_update}); - test.AddMissingOptionalInput(); - test.AddMissingOptionalInput(); - test.AddInput("ETA", {1}, eta); - if (step > 0) { - test.AddInput("Step", {}, {step}); - } else { - test.AddMissingOptionalInput(); - } - test.AddInput("W", shape, w); - test.AddInput("G", shape, g); - test.AddInput("Moment_1", shape, m); - test.AddInput("Moment_2", shape, v); - if (!w_half.empty()) { - test.AddInput("FP16_W", shape, w_half); - } else { - test.AddMissingOptionalInput(); - } - - test.AddAttribute("alpha", std::vector(1, alpha)); - test.AddAttribute("beta", std::vector(1, beta)); - test.AddAttribute("lambda", std::vector(1, lambda)); - test.AddAttribute("epsilon", std::vector(1, epsilon)); - test.AddAttribute("ratio_min", ratio_min); - test.AddAttribute("ratio_max", ratio_max); - - if (step > 0) { - test.AddOutput("Step_Out", {}, {do_update ? step + 1 : step}); - } else { - test.AddMissingOptionalOutput(); - } - if (!w_new.empty()) { - test.AddOutput("W_Out", shape, w_new); - } else { - test.AddMissingOptionalOutput(); - } - if (!g_new.empty()) { - test.AddOutput("G_Out", shape, g_new); - } else { - test.AddMissingOptionalOutput(); - } - test.AddOutput("Moment_1_Out", shape, m_new); - test.AddOutput("Moment_2_Out", shape, v_new); - if (!w_new_half.empty()) { - test.AddOutput("FP16_W_Out", shape, w_new_half); - } else { - test.AddMissingOptionalOutput(); - } - - test.Run(); -} - -template -void run_multi_tensor_lamb_test_with_baseline( - const std::vector>& shapes, - const T1 eta, - const T1 loss_scale, - const T1 g_norm, - const std::vector>& ws, - const std::vector>& gs, - const std::vector>& ms, - const std::vector>& vs, - const std::vector& alphas, - const std::vector& betas, - const std::vector& lambdas, - const std::vector& epsilons, - const std::vector>& w_news, - const std::vector>& g_news, - const std::vector>& m_news, - const std::vector>& v_news, - const std::vector>& w_halfs = {}, - const std::vector>& w_new_halfs = {}, - const bool do_update = true, - const int64_t step = 0, - const float ratio_min = -std::numeric_limits::infinity(), - const float ratio_max = std::numeric_limits::infinity()) { - OpTester test("LambOptimizer", 9, onnxruntime::kOnnxDomain, true); - - ORT_ENFORCE(shapes.size() == ws.size()); - ORT_ENFORCE(shapes.size() == gs.size()); - ORT_ENFORCE(shapes.size() == ms.size()); - ORT_ENFORCE(shapes.size() == vs.size()); - ORT_ENFORCE(shapes.size() == alphas.size()); - ORT_ENFORCE(shapes.size() == betas.size()); - ORT_ENFORCE(shapes.size() == lambdas.size()); - ORT_ENFORCE(shapes.size() == epsilons.size()); - if (!w_news.empty()) { - ORT_ENFORCE(shapes.size() == w_news.size()); - } - if (!g_news.empty()) { - ORT_ENFORCE(shapes.size() == g_news.size()); - } - ORT_ENFORCE(shapes.size() == m_news.size()); - ORT_ENFORCE(shapes.size() == v_news.size()); - if (!w_halfs.empty()) { - ORT_ENFORCE(shapes.size() == w_halfs.size()); - } - if (!w_new_halfs.empty()) { - ORT_ENFORCE(shapes.size() == w_new_halfs.size()); - } - - const int group_count = static_cast(ws.size()); - - test.AddInput("update_signal", {}, {do_update}); - test.AddInput("loss_scale", {}, {loss_scale}); - test.AddInput("gradient_norm", {}, {g_norm}); - test.AddInput("ETA", {}, {eta}); - if (step > 0) { - test.AddInput("Step", {}, {step}); - test.AddOutput("Step_Out", {}, {do_update ? step + 1 : step}); - } else { - test.AddMissingOptionalInput(); - test.AddMissingOptionalOutput(); - } - for (int i = 0; i < group_count; ++i) { - std::string w_name = "W_" + std::to_string(i); - std::string g_name = "G_" + std::to_string(i); - std::string m1_name = "Moment_1_" + std::to_string(i); - std::string m2_name = "Moment_2_" + std::to_string(i); - std::string w_fp16_name = "FP16_W_" + std::to_string(i); - std::string w_new_name = "W_Out_" + std::to_string(i); - std::string g_new_name = "G_Out_" + std::to_string(i); - std::string m1_new_name = "Moment_1_Out_" + std::to_string(i); - std::string m2_new_name = "Moment_2_Out_" + std::to_string(i); - std::string w_fp16_new_name = "FP16_W_Out_" + std::to_string(i); - - test.AddInput(w_name.c_str(), shapes[i], ws[i]); - test.AddInput(g_name.c_str(), shapes[i], gs[i]); - test.AddInput(m1_name.c_str(), shapes[i], ms[i]); - test.AddInput(m2_name.c_str(), shapes[i], vs[i]); - if (!w_halfs.empty() && !w_halfs[i].empty()) { - test.AddInput(w_fp16_name.c_str(), shapes[i], w_halfs[i]); - } else { - test.AddMissingOptionalInput(); - } - - if (!w_news.empty() && !w_news[i].empty()) { - test.AddOutput(w_new_name.c_str(), shapes[i], w_news[i]); - } else { - test.AddMissingOptionalOutput(); - } - if (!g_news.empty() && !g_news[i].empty()) { - test.AddOutput(g_new_name.c_str(), shapes[i], g_news[i]); - } else { - test.AddMissingOptionalOutput(); - } - test.AddOutput(m1_new_name.c_str(), shapes[i], m_news[i]); - test.AddOutput(m2_new_name.c_str(), shapes[i], v_news[i]); - if (!w_new_halfs.empty() && !w_new_halfs[i].empty()) { - test.AddOutput(w_fp16_new_name.c_str(), shapes[i], w_new_halfs[i]); - } else { - test.AddMissingOptionalOutput(); - } - } - - test.AddAttribute("alpha", alphas); - test.AddAttribute("beta", betas); - test.AddAttribute("lambda", lambdas); - test.AddAttribute("epsilon", epsilons); - test.AddAttribute("ratio_min", ratio_min); - test.AddAttribute("ratio_max", ratio_max); - - test.Run(); -} - -// Lamb test without baseline. This function computes -// baseline via an internal function and then invoke -// run_lamb_test_with_baseline(...) to check the result. -void run_multi_tensor_lamb_test( - const std::vector> shapes, - const float eta, - const float loss_scale, - const float g_norm, - const std::vector> ws, - const std::vector> gs, - const std::vector> ms, - const std::vector> vs, - const std::vector lambdas, - const std::vector alphas, - const std::vector betas, - const std::vector epsilons, - const int64_t step = 0, - const float ratio_min = -std::numeric_limits::infinity(), - const float ratio_max = std::numeric_limits::infinity()) { - // Check if parallel vectors have the same length. - ORT_ENFORCE(shapes.size() == ws.size()); - ORT_ENFORCE(shapes.size() == gs.size()); - ORT_ENFORCE(shapes.size() == ms.size()); - ORT_ENFORCE(shapes.size() == vs.size()); - ORT_ENFORCE(shapes.size() == alphas.size()); - ORT_ENFORCE(shapes.size() == betas.size()); - ORT_ENFORCE(shapes.size() == lambdas.size()); - ORT_ENFORCE(shapes.size() == epsilons.size()); - - const int group_count = static_cast(ws.size()); - - // Output buffers of the optimizer. - std::vector> w_news(group_count); - std::vector> g_news(group_count); - std::vector> m_news(group_count); - std::vector> v_news(group_count); - - for (int i = 0; i < group_count; ++i) { - w_news[i] = std::vector(ws[i].size(), 0.f); - g_news[i] = std::vector(gs[i].size(), 0.f); - m_news[i] = std::vector(ms[i].size(), 0.f); - v_news[i] = std::vector(vs[i].size(), 0.f); - - // Invoke LAMB's reference implementation to compute baseline output. - compute_lamb( - shapes[i], ws[i], gs[i], ms[i], vs[i], - eta, loss_scale, g_norm, - lambdas[i], alphas[i], betas[i], epsilons[i], - w_news[i], g_news[i], m_news[i], v_news[i], step, - ratio_min, ratio_max); - } - - // Create tests to make sure the output is correct. - - // Output new weights. - run_multi_tensor_lamb_test_with_baseline( - shapes, eta, loss_scale, g_norm, - ws, gs, ms, vs, - alphas, betas, lambdas, epsilons, - w_news, {}, m_news, v_news, {}, {}, true, step, - ratio_min, ratio_max); - - // Output new gradients. - run_multi_tensor_lamb_test_with_baseline( - shapes, eta, loss_scale, g_norm, - ws, gs, ms, vs, - alphas, betas, lambdas, epsilons, - {}, g_news, m_news, v_news, {}, {}, true, step, - ratio_min, ratio_max); -} - -void run_lamb_mix_precision_test( - const std::vector& shape, - const std::vector& eta, - const std::vector& w, - const std::vector& g, - const std::vector& m, - const std::vector& v, - const float lambda, - const float alpha, - const float beta, - const float epsilon, - const int64_t step = 0) { - std::vector w_new(w.size(), 0); - std::vector g_new(g.size(), 0); - std::vector m_new(m.size(), 0); - std::vector v_new(v.size(), 0); - - // Invoke LAMB's reference implementation to compute output. - compute_lamb( - shape, w, g, m, v, - eta[0], 1.f, 1.f, lambda, alpha, beta, epsilon, - w_new, g_new, m_new, v_new, step); - - std::vector eta_half(eta.size()); - std::vector g_half(w.size()); - std::vector m_half(w.size()); - std::vector v_half(w.size()); - std::vector w_half(w.size()); - ConvertFloatToMLFloat16(eta.data(), eta_half.data(), int(eta.size())); - ConvertFloatToMLFloat16(g.data(), g_half.data(), int(g.size())); - ConvertFloatToMLFloat16(m.data(), m_half.data(), int(m.size())); - ConvertFloatToMLFloat16(v.data(), v_half.data(), int(v.size())); - ConvertFloatToMLFloat16(w.data(), w_half.data(), int(w.size())); - - std::vector m_new_half(m_new.size()); - std::vector v_new_half(v_new.size()); - std::vector w_new_half(w_new.size()); - std::vector g_new_half(g_new.size()); - ConvertFloatToMLFloat16(m_new.data(), m_new_half.data(), int(m_new.size())); - ConvertFloatToMLFloat16(v_new.data(), v_new_half.data(), int(v_new.size())); - ConvertFloatToMLFloat16(w_new.data(), w_new_half.data(), int(w_new.size())); - ConvertFloatToMLFloat16(g_new.data(), g_new_half.data(), int(g_new.size())); - - // Half momentums, without fp16 weight - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w_new, {}, m_new_half, v_new_half, {}, {}, true, step); - - // Float momentums, without fp16 weight - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, {}, {}, true, step); - - // Half momentums, with fp16 weight - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w_new, {}, m_new_half, v_new_half, {}, {}, true, step); - - // Float momentums, with fp16 weight - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, w_half, w_new_half, true, step); - - // Half momentums, with fp16 weight, skip weight update - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w, {}, m_half, v_half, w_half, w_half, false, step); - - // Float momentums, with fp16 weight, skip weight update - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w, {}, m, v, w_half, w_half, false, step); - - // Float eta, float momentums, with fp16 weight - run_lamb_test_with_baseline( - shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, w_half, w_new_half, true, step); - - // Float eta, float momentums, with fp16 weight, skip weight update - run_lamb_test_with_baseline( - shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, w, {}, m, v, w_half, w_half, false, step); - - // Float momentums, without fp16 weight, output gradients only - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_new_half, m_new, v_new, {}, {}, true, step); - - // Float momentums, with fp16 weight, output gradients only - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_new_half, m_new, v_new, w_half, {}, true, step); - - // Float momentums, with fp16 weight, output gradients only, skip weight update - run_lamb_test_with_baseline( - shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_half, m, v, w_half, {}, false, step); -} - -// A optimizer test with an 2-element vector. -TEST(OptimizerTest, LambOptimizerTestVector) { - // Input tensors and attributes. - const std::vector shape = {2}; - const float eta = 0.5f; - const std::vector w = {1.0f, 2.0f}; - const std::vector g = {3.0f, 4.0f}; - const std::vector m = {-1.0f, -2.0f}; - const std::vector v = {2.0f, 1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}); -} - -TEST(OptimizerTest, LambOptimizerTestVectorWithZeroWeight) { - // Input tensors and attributes. - const std::vector shape = {2}; - const float eta = 0.5f; - const std::vector w = {0.0f, 0.0f}; - const std::vector g = {1.0f, -1.0f}; - const std::vector m = {-1.0f, -2.0f}; - const std::vector v = {2.0f, 1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}); -} - -TEST(OptimizerTest, LambOptimizerRatioMin) { - // Input tensors and attributes. - const std::vector shape = {2}; - const float eta = 0.5f; - const std::vector w = {-1.0f, 1.0f}; - const std::vector g = {1.0f, -1.0f}; - const std::vector m = {-1.0f, -2.0f}; - const std::vector v = {2.0f, 1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - const float ratio_min = -std::numeric_limits::infinity(); - const float ratio_max = 0.1f; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}, - 0, - ratio_min, - ratio_max); -} - -TEST(OptimizerTest, LambOptimizerRatioMax) { - // Input tensors and attributes. - const std::vector shape = {2}; - const float eta = 0.5f; - const std::vector w = {0.0001f, -0.0001f}; - const std::vector g = {1.0f, -1.0f}; - const std::vector m = {-1.0f, -2.0f}; - const std::vector v = {2.0f, 1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - const float ratio_min = 1.0f; - const float ratio_max = std::numeric_limits::infinity(); - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}, - 0, - ratio_min, - ratio_max); -} - -TEST(OptimizerTest, LambOptimizerTestBiasCorrectionFirst) { - // Input tensors and attributes. - const std::vector shape = {2}; - const float eta = 0.5f; - const std::vector w = {1.0f, 2.0f}; - const std::vector g = {3.0f, 4.0f}; - const std::vector m = {-1.0f, -2.0f}; - const std::vector v = {2.0f, 1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}, - 1); -} - -TEST(OptimizerTest, LambOptimizerTestBiasCorrectionThird) { - // Input tensors and attributes. - const std::vector shape = {2}; - const float eta = 0.5f; - const std::vector w = {1.0f, 2.0f}; - const std::vector g = {3.0f, 4.0f}; - const std::vector m = {-1.0f, -2.0f}; - const std::vector v = {2.0f, 1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}, - 3); -} - -// A optimizer test with an 2-by-1-by-1-by-1 tensor. -TEST(OptimizerTest, LambOptimizerTest4DTensor) { - // Input tensors and attributes. - const std::vector shape = {2, 1, 1, 1}; - const float eta = 0.5f; - const std::vector w = {1.0f, 2.0f}; - const std::vector g = {3.0f, 4.0f}; - const std::vector m = {-1.0f, -2.0f}; - const std::vector v = {2.0f, 1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}); -} - -// A optimizer test with an 2-by-3 tensor. -TEST(OptimizerTest, LambOptimizerTest2by3Tensor) { - // Input tensors and attributes. - const std::vector shape = {2, 3}; - const float eta = 0.5f; - const std::vector w = {1.0f, 2.0f, 1.0f, 1.0f, 2.0f, 2.0f}; - const std::vector g = {3.0f, 4.0f, 3.0f, 3.0f, 4.0f, 4.0f}; - const std::vector m = {-1.0f, -2.0f, 2.0f, 1.0f, 1.0f, -2.0f}; - const std::vector v = {1.0f, 1.0f, 5.0f, 5.0f, 6.0f, 6.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}); -} - -// A optimizer test with an 1-element tensor. -TEST(OptimizerTest, LambOptimizerTestScalar) { - // Input tensors and attributes. - const std::vector shape = {(int64_t)1}; - const float eta = 0.5f; - const std::vector w = {1.0f}; - const std::vector g = {3.0f}; - const std::vector m = {-10.0f}; - const std::vector v = {1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - // Intermediate and output buffers of the optimizer. - std::vector m_new = {0.0f}; - std::vector v_new = {0.0f}; - std::vector w_new = {0.0f}; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}); -} - -TEST(OptimizerTest, LambOptimizerTestScalarScaling) { - // Input tensors and attributes. - const std::vector shape = {(int64_t)1}; - const float eta = 0.5f; - const std::vector w = {1.0f}; - const std::vector g = {3.0f}; - const std::vector m = {-10.0f}; - const std::vector v = {1.0f}; - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - // Intermediate and output buffers of the optimizer. - std::vector m_new = {0.0f}; - std::vector v_new = {0.0f}; - std::vector w_new = {0.0f}; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 8.f, - 4.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}); -} - -TEST(OptimizerTest, LambOptimizerTestExternalBaseline) { - // Input tensors and attributes. - const std::vector shape = {2, 5}; - const std::vector eta = {0.1f}; - const std::vector w = { - 0.01379026f, 0.15308191f, -0.24356517f, -0.21798165f, -0.13770047f, 0.09694599f, - -0.02223516f, 0.2664228f, -0.01177993f, 0.06832688f}; - const std::vector g = { - -6.048543f, 10.569487f, -9.207029f, -0.57407373f, - 5.884985f, -0.21047728f, 3.539946f, -5.957566f, -9.343748f, 1.1502024f}; - const std::vector m = { - -5.9078765f, 9.673933f, -8.731428f, -0.6227454f, 5.284312f, -0.27138948f, - 3.443532f, -5.681713f, -8.72421f, 1.1441823f}; - const std::vector v = { - 4.2659229e+01f, 1.1438165e+02f, 9.3179581e+01f, 4.7399229e-01f, 3.4129276e+01f, - 9.0019435e-02f, 1.4493006e+01f, 3.9455612e+01f, 9.3025581e+01f, 1.6000764e+0f}; - - const float lambda = 0.1f; - const float alpha = 0.1f; - const float beta = 0.01f; - const float epsilon = 0.1f; - - std::vector w_new = { - 0.02979828f, 0.13677707f, -0.22708717f, -0.20361158f, -0.15338624f, 0.1081504f, - -0.03804127f, 0.28198114f, 0.00430069f, 0.05319814f}; - std::vector g_new = { - 0.01600802f, -0.01630484f, 0.01647800f, 0.01437007f, -0.01568577f, 0.01120441f, - -0.01580611f, 0.01555834f, 0.01608062f, -0.01512874f}; - std::vector m_new = { - -6.0344763f, 10.479931f, -9.15947f, -0.57894087f, 5.824918f, -0.2165685f, - 3.5303047f, -5.9299808f, -9.281795f, 1.1496004f}; - std::vector v_new = { - 3.6645618e+01f, 1.1174072e+02f, 8.4853485e+01f, 3.3100498e-01f, 3.4628010e+01f, - 4.4757873e-02f, 1.2550836e+01f, 3.5532223e+01f, 8.7362823e+01f, 1.3257366e+00f}; - - // Output new weights - run_lamb_test_with_baseline( - shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new); - - // Output new gradients - run_lamb_test_with_baseline( - shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, {}, g_new, m_new, v_new); -} - -TEST(OptimizerTest, LambOptimizerTestExternalBaselineDouble) { - // Input tensors and attributes. - const std::vector shape = {2, 5}; - const std::vector eta = {0.1f}; - const std::vector w = { - 0.01379026, 0.15308191, -0.24356517, -0.21798165, -0.13770047, 0.09694599, - -0.02223516, 0.2664228, -0.01177993, 0.06832688}; - const std::vector g = { - -6.048543, 10.569487, -9.207029, -0.57407373, - 5.884985, -0.21047728, 3.539946, -5.957566, -9.343748, 1.1502024}; - const std::vector m = { - -5.9078765, 9.673933, -8.731428, -0.6227454, 5.284312, -0.27138948, - 3.443532, -5.681713, -8.72421, 1.1441823}; - const std::vector v = { - 4.2659229e+01, 1.1438165e+02, 9.3179581e+01, 4.7399229e-01, 3.4129276e+01, - 9.0019435e-02, 1.4493006e+01, 3.9455612e+01, 9.3025581e+01, 1.6000764e+0}; - - const float lambda = 0.1f; - const float alpha = 0.1f; - const float beta = 0.01f; - const float epsilon = 0.1f; - - std::vector w_new = { - 0.02979828, 0.13677707, -0.22708717, -0.20361158, -0.15338624, 0.1081504, - -0.03804127, 0.28198114, 0.00430069, 0.05319814}; - std::vector g_new = { - 0.01600802, -0.01630484, 0.016478, 0.01437007, -0.01568577, 0.01120441, - -0.01580611, 0.01555834, 0.01608062, -0.01512874}; - std::vector m_new = { - -6.0344763, 10.479931, -9.15947, -0.57894087, 5.824918, -0.2165685, - 3.5303047, -5.9299808, -9.281795, 1.1496004}; - std::vector v_new = { - 3.6645618e+01, 1.1174072e+02, 8.4853485e+01, 3.3100498e-01, 3.4628010e+01, - 4.4757873e-02, 1.2550836e+01, 3.5532223e+01, 8.7362823e+01, 1.3257366e+00}; - - // Output new weights - run_lamb_test_with_baseline( - shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new); - - // Output new gradients - run_lamb_test_with_baseline( - shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, {}, g_new, m_new, v_new); -} - -TEST(OptimizerTest, LambOptimizerTest5DTensorMixPrecision32_16) { - const std::vector shape = {2, 2, 2, 1, 1}; - const std::vector eta = {0.5f}; - const std::vector w = {1.0f, 2.0f, 2.5f, 1.5f, 1.0f, 2.0f, 2.0f, 1.5f}; - const std::vector g = {-1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 0.8f}; - const std::vector m = {1.0f, 2.0f, -0.25f, 1.1f, 1.0f, 2.0f, -0.21f, 1.1f}; - const std::vector v = {1.5f, 1.0f, 1.1f, 0.76f, 1.5f, 1.0f, 1.5f, 0.76f}; - - const float lambda = 1.5f; - const float alpha = 1.5f; - const float beta = 1.5f; - const float epsilon = 1.0f; - - run_lamb_mix_precision_test( - shape, eta, w, g, m, v, - lambda, alpha, beta, epsilon); -} - -TEST(OptimizerTest, LambOptimizerTestSimpleBaselineMixPrecision32_16) { - const std::vector shape = {2, 1}; - const std::vector eta = {1.0f}; - const std::vector w = {1.0f, 1.0f}; - const std::vector g = {-1.0f, 1.0f}; - const std::vector m = {1.0f, 1.0f}; - const std::vector v = {0.0f, 0.0f}; - - const float lambda = 0.0f; - const float alpha = 1.0f; - const float beta = 1.0f; - const float epsilon = 1.0f; - - run_lamb_mix_precision_test( - shape, eta, w, g, m, v, - lambda, alpha, beta, epsilon); -} - -TEST(OptimizerTest, LambOptimizerTestBaselineMixPrecision32_16) { - const std::vector shape = {2, 1}; - const std::vector eta = {0.1f}; - const std::vector w = {-1.5f, 2.4f}; - const std::vector g = {-0.75f, 1.2f}; - const std::vector m = {0.87f, -0.94f}; - const std::vector v = {0.12f, 0.28f}; - - const float lambda = 0.25f; - const float alpha = 0.9f; - const float beta = 0.95f; - const float epsilon = 0.33f; - - run_lamb_mix_precision_test( - shape, eta, w, g, m, v, - lambda, alpha, beta, epsilon); -} - -TEST(OptimizerTest, LambOptimizerTestScalarMixPrecision32_16) { - const std::vector shape = {1}; - const std::vector eta = {0.1f}; - const std::vector w = {-1.5f}; - const std::vector g = {-0.75f}; - const std::vector m = {0.87f}; - const std::vector v = {0.12f}; - - const float lambda = 0.25f; - const float alpha = 0.9f; - const float beta = 0.95f; - const float epsilon = 0.33f; - - run_lamb_mix_precision_test( - shape, eta, w, g, m, v, - lambda, alpha, beta, epsilon); - - run_lamb_mix_precision_test( - shape, eta, w, g, m, v, - lambda, alpha, beta, epsilon, 2); -} - -TEST(OptimizerTest, LambOptimizerTestLarge) { - // Input tensors and attributes. - for (const auto& size : {55667, 1944006, 3907584}) { - const std::vector shape = {static_cast(size)}; - const float eta = 0.5f; - std::vector w(size); - std::vector g(size); - std::vector m(size); - std::vector v(size); - - std::random_device random_device; - std::mt19937 random_engine(0); - std::uniform_real_distribution dist(0.1f, 1.0f); - for (int i = 0; i < size; ++i) { - w[i] = dist(random_engine); - g[i] = dist(random_engine); - m[i] = dist(random_engine); - v[i] = dist(random_engine); - } - - const float lambda = 0.5f; - const float alpha = 0.2f; - const float beta = 0.8f; - const float epsilon = 1e-6f; - - run_multi_tensor_lamb_test( - {shape}, - eta, - 1.f, - 1.f, - {w}, - {g}, - {m}, - {v}, - {lambda}, - {alpha}, - {beta}, - {epsilon}); - } -} - -TEST(OptimizerTest, LambOptimizerMultiTensorRatio) { - const int group_count = 127; - std::random_device random_device; - std::mt19937 random_engine(0); - std::uniform_real_distribution dist(0.1f, 1.0f); - std::uniform_int_distribution dist_int(1, 1228); - - std::vector sizes(group_count); - std::vector> shapes(group_count); - - std::vector> ws(group_count); - std::vector> gs(group_count); - std::vector> ms(group_count); - std::vector> vs(group_count); - std::vector alphas(group_count); - std::vector betas(group_count); - std::vector lambdas(group_count); - std::vector epsilons(group_count); - - const float eta = dist(random_engine); - - for (int64_t i = 0; i < group_count; ++i) { - const auto size = dist_int(random_engine); - sizes[i] = size; - shapes[i] = std::vector(1, size); - - ws[i] = std::vector(sizes[i]); - gs[i] = std::vector(sizes[i]); - ms[i] = std::vector(sizes[i]); - vs[i] = std::vector(sizes[i]); - - for (int64_t j = 0; j < sizes[i]; ++j) { - ws[i][j] = dist(random_engine); - gs[i][j] = dist(random_engine); - ms[i][j] = dist(random_engine); - vs[i][j] = dist(random_engine); - } - - alphas[i] = dist(random_engine); - betas[i] = dist(random_engine); - lambdas[i] = dist(random_engine); - epsilons[i] = dist(random_engine); - } - - run_multi_tensor_lamb_test( - shapes, eta, 1.f, 1.f, - ws, gs, ms, vs, - lambdas, alphas, betas, epsilons, - 0, 0.3f, 0.7f); - - run_multi_tensor_lamb_test( - shapes, eta, 1.f, 1.f, - ws, gs, ms, vs, - lambdas, alphas, betas, epsilons); -} - TEST(GradientCheckerTest, LayerNormGrad) { GradientChecker gradient_checker; { @@ -2699,7 +1444,6 @@ TEST(GradientUtilsTest, ZeroGradientFloat16) { test.Run(); } - #endif TEST(GradientCheckerTest, SliceGrad) { diff --git a/orttraining/orttraining/test/gradient/optimizer_ops_test.cc b/orttraining/orttraining/test/gradient/optimizer_ops_test.cc new file mode 100644 index 0000000000..47f75a3195 --- /dev/null +++ b/orttraining/orttraining/test/gradient/optimizer_ops_test.cc @@ -0,0 +1,1367 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime { +namespace test { + +TEST(OptimizerTest, SGDOptimizerTest) { + OpTester test("SGDOptimizer", 9, onnxruntime::kOnnxDomain); + test.AddInput("ETA", {}, {0.5f}); + test.AddInput("W", {3}, {1, 2, 3}); + test.AddInput("G", {3}, {4, 5, 6}); + test.AddOutput("W_New", {3}, {-1.f, -0.5f, 0.f}); + test.Run(); +} + +TEST(OptimizerTest, SGDOptimizerTest_Gradient) { + OpTester test("SGDOptimizer", 9, onnxruntime::kOnnxDomain); + test.AddInput("ETA", {}, {0.5f}); + test.AddInput("W", {3}, {1, 2, 3}); + test.AddInput("G", {3}, {4, 5, 6}); + test.AddMissingOptionalOutput(); + test.AddOutput("G_New", {3}, {-2.f, -2.5f, -3.f}); + test.Run(); +} + +struct AdamOptimizerInputOutput { + AdamOptimizerInputOutput() { + eta_half.resize(eta.size()); + g_half.resize(g.size()); + m1_half.resize(m1.size()); + m2_half.resize(m2.size()); + w_half.resize(w.size()); + ConvertFloatToMLFloat16(eta.data(), eta_half.data(), int(eta.size())); + ConvertFloatToMLFloat16(g.data(), g_half.data(), int(g.size())); + ConvertFloatToMLFloat16(m1.data(), m1_half.data(), int(m1.size())); + ConvertFloatToMLFloat16(m2.data(), m2_half.data(), int(m2.size())); + ConvertFloatToMLFloat16(w.data(), w_half.data(), int(w.size())); + + m1_new_half.resize(m1_new.size()); + m2_new_half.resize(m2_new.size()); + w_new_half.resize(w_new.size()); + g_new_half.resize(g_new.size()); + ConvertFloatToMLFloat16(m1_new.data(), m1_new_half.data(), int(m1_new.size())); + ConvertFloatToMLFloat16(m2_new.data(), m2_new_half.data(), int(m2_new.size())); + ConvertFloatToMLFloat16(w_new.data(), w_new_half.data(), int(w_new.size())); + ConvertFloatToMLFloat16(g_new.data(), g_new_half.data(), int(g_new.size())); + } + + // Fp32 Inputs + std::vector eta = {0.5f}; + std::vector w = {1.0f, 2.0f, 3.0f}; + std::vector g = {4.0f, 5.0f, 6.0f}; + std::vector m1 = {0.1f, 0.2f, 0.3f}; + std::vector m2 = {0.4f, 0.5f, 0.6f}; + + // Fp16 Inputs + std::vector eta_half; + std::vector w_half; + std::vector g_half; + std::vector m1_half; + std::vector m2_half; + + // FP32 Outptus + std::vector m1_new = {0.49f, 0.68f, 0.87f}; + std::vector m2_new = {0.4156f, 0.5245f, 0.6354f}; + std::vector w_new = {0.6199609f, 1.5305318f, 2.4542853f}; + std::vector g_new = {-0.3800391f, -0.4694682f, -0.5457147f}; + + // FP16 Outptus + std::vector m1_new_half; + std::vector m2_new_half; + std::vector w_new_half; + std::vector g_new_half; +}; + +TEST(OptimizerTest, AdamOptimizerTest) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, data.eta); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g); + test.AddInput("Moment_1", {3}, data.m1); + test.AddInput("Moment_2", {3}, data.m2); + + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {4}); + test.AddOutput("Moment_1_Out", {3}, data.m1_new); + test.AddOutput("Moment_2_Out", {3}, data.m2_new); + test.AddOutput("W_Out", {3}, data.w_new); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamOptimizerTest_Gradient) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, data.eta); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g); + test.AddInput("Moment_1", {3}, data.m1); + test.AddInput("Moment_2", {3}, data.m2); + + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {4}); + test.AddOutput("Moment_1_Out", {3}, data.m1_new); + test.AddOutput("Moment_2_Out", {3}, data.m2_new); + test.AddMissingOptionalOutput(); + test.AddOutput("G_Out", {3}, data.g_new); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamBiasCorrection) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, {1.f}); + test.AddInput("Update_Count", {}, {1}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); + test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); + test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); + + test.AddOutput("Update_Count_Out", {}, {2}); + test.AddOutput("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f}); + test.AddOutput("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f}); + test.AddOutput("W_Out", {3}, {-1.4634f, -0.6416f, -1.2121f}); + + test.AddAttribute("do_bias_correction", static_cast(1)); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamWeightDecayMode0NoBiasCorrection) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, {1.f}); + test.AddInput("Update_Count", {}, {1}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); + test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); + test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); + + test.AddOutput("Update_Count_Out", {}, {2}); + test.AddOutput("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f}); + test.AddOutput("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f}); + test.AddOutput("W_Out", {3}, {-3.6210f, -2.8075f, -3.3723f}); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("lambda", 0.01f); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamWeightDecayMode0WithBiasCorrection) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, {1.f}); + test.AddInput("Update_Count", {}, {1}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); + test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); + test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); + + test.AddOutput("Update_Count_Out", {}, {2}); + test.AddOutput("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f}); + test.AddOutput("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f}); + test.AddOutput("W_Out", {3}, {-1.4587f, -0.6452f, -1.2099f}); + + test.AddAttribute("do_bias_correction", static_cast(1)); + test.AddAttribute("lambda", 0.01f); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamWeightDecayMode1NoBiasCorrection) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, {1.f}); + test.AddInput("Update_Count", {}, {1}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); + test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); + test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); + + test.AddOutput("Update_Count_Out", {}, {2}); + test.AddOutput("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f}); + test.AddOutput("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f}); + test.AddOutput("W_Out", {3}, {-3.5894f, -2.7758f, -3.3406f}); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("lambda", 0.01f); + test.AddAttribute("weight_decay_mode", static_cast(1)); + + test.Run(); +} + +TEST(OptimizerTest, AdamWeightDecayMode1WithBiasCorrection) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, {1.f}); + test.AddInput("Update_Count", {}, {1}); + test.AddInput("W", {3}, {-0.4634f, 0.3584f, -0.2121f}); + test.AddInput("G", {3}, {0.4171f, 0.9485f, 1.2289f}); + test.AddInput("Moment_1", {3}, {0.f, 0.f, 0.f}); + test.AddInput("Moment_2", {3}, {0.f, 0.f, 0.f}); + + test.AddOutput("Update_Count_Out", {}, {2}); + test.AddOutput("Moment_1_Out", {3}, {0.0417f, 0.0949f, 0.1229f}); + test.AddOutput("Moment_2_Out", {3}, {1.7400e-04f, 8.9966e-04f, 1.5102e-03f}); + test.AddOutput("W_Out", {3}, {-1.4488f, -0.6352f, -1.1999f}); + + test.AddAttribute("do_bias_correction", static_cast(1)); + test.AddAttribute("lambda", 0.01f); + test.AddAttribute("weight_decay_mode", static_cast(1)); + + test.Run(); +} + +#ifdef USE_CUDA +TEST(OptimizerTest, AdamOptimizerMixPrecisionTest) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, data.eta_half); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g_half); + test.AddInput("Moment_1", {3}, data.m1_half); + test.AddInput("Moment_2", {3}, data.m2_half); + + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {4}); + test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); + test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); + test.AddOutput("W_Out", {3}, data.w_new); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_Test) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, data.eta_half); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g_half); + test.AddInput("Moment_1", {3}, data.m1_half); + test.AddInput("Moment_2", {3}, data.m2_half); + test.AddInput("FP16_W", {3}, data.w_half); + + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {4}); + test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); + test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); + test.AddOutput("W_Out", {3}, data.w_new); + test.AddMissingOptionalOutput(); + test.AddOutput("FP16_W_Out", {3}, data.w_new_half); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamOptimizerMixPrecision_FP16Weight_SkipUpdate_Test) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, data.eta_half); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g_half); + test.AddInput("Moment_1", {3}, data.m1_half); + test.AddInput("Moment_2", {3}, data.m2_half); + test.AddInput("FP16_W", {3}, data.w_half); + test.AddInput("loss_scale", {1}, {1.0f}); + // grad clipping should not take effect + test.AddInput("grad_norm", {1}, {0.01f}); + test.AddInput("DoUpdate", {1}, {false}); + + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {3}); + test.AddOutput("Moment_1_Out", {3}, data.m1_half); + test.AddOutput("Moment_2_Out", {3}, data.m2_half); + test.AddOutput("W_Out", {3}, data.w); + test.AddMissingOptionalOutput(); + test.AddOutput("FP16_W_Out", {3}, data.w_half); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamOptimizerMixPrecisionTestFloatEta) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, data.eta); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g_half); + test.AddInput("Moment_1", {3}, data.m1_half); + test.AddInput("Moment_2", {3}, data.m2_half); + + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {4}); + test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); + test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); + test.AddOutput("W_Out", {3}, data.w_new); + + test.AddAttribute("do_bias_correction", static_cast(0)); + test.AddAttribute("weight_decay_mode", static_cast(0)); + + test.Run(); +} + +TEST(OptimizerTest, AdamOptimizerMixPrecisionTest_Gradient) { + OpTester test("AdamOptimizer", 9, onnxruntime::kOnnxDomain); + AdamOptimizerInputOutput data; + + test.AddInput("ETA", {}, data.eta); + test.AddInput("Update_Count", {}, {3}); + test.AddInput("W", {3}, data.w); + test.AddInput("G", {3}, data.g_half); + test.AddInput("Moment_1", {3}, data.m1_half); + test.AddInput("Moment_2", {3}, data.m2_half); + + // Verify AdamOptimizer outputs + test.AddOutput("Update_Count_Out", {}, {4}); + test.AddOutput("Moment_1_Out", {3}, data.m1_new_half); + test.AddOutput("Moment_2_Out", {3}, data.m2_new_half); + test.AddMissingOptionalOutput(); + test.AddOutput("G_Out", {3}, data.g_new_half); + + test.AddAttribute("do_bias_correction", static_cast(0)); + + test.Run(); +} + +// This helper function is a CPU-based LAMB optimizer +// implementation. It mainly focuses on readability. +void compute_lamb( + const std::vector shape, + /* weights */ const std::vector& w, + /* gradient */ const std::vector& g, + /* momentum */ const std::vector& m, + /* 2nd-order momentum */ const std::vector& v, + const float eta, + const float loss_scale, + const float g_norm, + const float lambda, + const float alpha, + const float beta, + const float epsilon, + /* updated weights */ std::vector& w_new, + /* updated gradients */ std::vector& g_new, + /* updated momentum */ std::vector& m_new, + /* updated 2nd-order momentum */ std::vector& v_new, + const int64_t step = 0, + const float ratio_min = -std::numeric_limits::infinity(), + const float ratio_max = std::numeric_limits::infinity()) { + // Element counts of all vector-typed arguments. + const int64_t size = std::accumulate( + shape.begin(), + shape.end(), + (int64_t)1, + std::multiplies()); + + // Buffer to store update direction. + std::vector r(size, 0.0f); + + float g_scale = loss_scale; + if (g_norm > loss_scale) { + g_scale *= g_norm / loss_scale; + } + + const float alpha_correction = step > 0 ? 1.f - std::pow(alpha, static_cast(step)) : 1.f; + const float beta_correction = step > 0 ? 1.f - std::pow(beta, static_cast(step)) : 1.f; + + // Compute new 1st-, 2nd-order momentums, and the update direction. + for (int i = 0; i < size; ++i) { + const float g_scaled = g[i] / g_scale; + m_new[i] = alpha * m[i] + (1.0f - alpha) * g_scaled; + v_new[i] = beta * v[i] + (1.0f - beta) * g_scaled * g_scaled; + const float m_new_tmp = m_new[i] / alpha_correction; + const float v_new_tmp = v_new[i] / beta_correction; + r[i] = lambda * w[i] + m_new_tmp / (std::sqrt(v_new_tmp) + epsilon); + } + + // Compute squared sum of all elements. Note that Eigen sqrt could lead to significant + // numerical error so we use std::sqrt. The std::inner_product produces wrong result + // when std::inner_product(r.begin(), r.end(), r.begin(), 0) so we just use a loop below. + float r_norm = 0.0f; + float w_norm = 0.0f; + for (int i = 0; i < size; ++i) { + r_norm += r[i] * r[i]; + w_norm += w[i] * w[i]; + } + + r_norm = std::sqrt(r_norm); + w_norm = std::sqrt(w_norm); + + float ratio = (w_norm != 0.0f && r_norm != 0.0f) ? w_norm / r_norm : 1.0f; + + if (ratio > ratio_max) { + ratio = ratio_max; + } + + if (ratio < ratio_min) { + ratio = ratio_min; + } + + ratio *= eta; + + // Compute the new weight. + for (int64_t i = 0; i < size; ++i) { + g_new[i] = -ratio * r[i]; + w_new[i] = w[i] + g_new[i]; + } +} + +template +void run_lamb_test_with_baseline( + const std::vector& shape, + const std::vector& eta, + const std::vector& w, + const std::vector& g, + const std::vector& m, + const std::vector& v, + const float alpha, + const float beta, + const float lambda, + const float epsilon, + const std::vector& w_new, + const std::vector& g_new, + const std::vector& m_new, + const std::vector& v_new, + const std::vector& w_half = {}, + const std::vector& w_new_half = {}, + const bool do_update = true, + const int64_t step = 0, + const float ratio_min = -std::numeric_limits::infinity(), + const float ratio_max = std::numeric_limits::infinity()) { + OpTester test("LambOptimizer", 9, onnxruntime::kOnnxDomain, true); + + test.AddInput("update_signal", {1}, {do_update}); + test.AddMissingOptionalInput(); + test.AddMissingOptionalInput(); + test.AddInput("ETA", {1}, eta); + if (step > 0) { + test.AddInput("Step", {}, {step}); + } else { + test.AddMissingOptionalInput(); + } + test.AddInput("W", shape, w); + test.AddInput("G", shape, g); + test.AddInput("Moment_1", shape, m); + test.AddInput("Moment_2", shape, v); + if (!w_half.empty()) { + test.AddInput("FP16_W", shape, w_half); + } else { + test.AddMissingOptionalInput(); + } + + test.AddAttribute("alpha", std::vector(1, alpha)); + test.AddAttribute("beta", std::vector(1, beta)); + test.AddAttribute("lambda", std::vector(1, lambda)); + test.AddAttribute("epsilon", std::vector(1, epsilon)); + test.AddAttribute("ratio_min", ratio_min); + test.AddAttribute("ratio_max", ratio_max); + + if (step > 0) { + test.AddOutput("Step_Out", {}, {do_update ? step + 1 : step}); + } else { + test.AddMissingOptionalOutput(); + } + if (!w_new.empty()) { + test.AddOutput("W_Out", shape, w_new); + } else { + test.AddMissingOptionalOutput(); + } + if (!g_new.empty()) { + test.AddOutput("G_Out", shape, g_new); + } else { + test.AddMissingOptionalOutput(); + } + test.AddOutput("Moment_1_Out", shape, m_new); + test.AddOutput("Moment_2_Out", shape, v_new); + if (!w_new_half.empty()) { + test.AddOutput("FP16_W_Out", shape, w_new_half); + } else { + test.AddMissingOptionalOutput(); + } + + test.Run(); +} + +template +void run_multi_tensor_lamb_test_with_baseline( + const std::vector>& shapes, + const T1 eta, + const T1 loss_scale, + const T1 g_norm, + const std::vector>& ws, + const std::vector>& gs, + const std::vector>& ms, + const std::vector>& vs, + const std::vector& alphas, + const std::vector& betas, + const std::vector& lambdas, + const std::vector& epsilons, + const std::vector>& w_news, + const std::vector>& g_news, + const std::vector>& m_news, + const std::vector>& v_news, + const std::vector>& w_halfs = {}, + const std::vector>& w_new_halfs = {}, + const bool do_update = true, + const int64_t step = 0, + const float ratio_min = -std::numeric_limits::infinity(), + const float ratio_max = std::numeric_limits::infinity()) { + OpTester test("LambOptimizer", 9, onnxruntime::kOnnxDomain, true); + + ORT_ENFORCE(shapes.size() == ws.size()); + ORT_ENFORCE(shapes.size() == gs.size()); + ORT_ENFORCE(shapes.size() == ms.size()); + ORT_ENFORCE(shapes.size() == vs.size()); + ORT_ENFORCE(shapes.size() == alphas.size()); + ORT_ENFORCE(shapes.size() == betas.size()); + ORT_ENFORCE(shapes.size() == lambdas.size()); + ORT_ENFORCE(shapes.size() == epsilons.size()); + if (!w_news.empty()) { + ORT_ENFORCE(shapes.size() == w_news.size()); + } + if (!g_news.empty()) { + ORT_ENFORCE(shapes.size() == g_news.size()); + } + ORT_ENFORCE(shapes.size() == m_news.size()); + ORT_ENFORCE(shapes.size() == v_news.size()); + if (!w_halfs.empty()) { + ORT_ENFORCE(shapes.size() == w_halfs.size()); + } + if (!w_new_halfs.empty()) { + ORT_ENFORCE(shapes.size() == w_new_halfs.size()); + } + + const int group_count = static_cast(ws.size()); + + test.AddInput("update_signal", {}, {do_update}); + test.AddInput("loss_scale", {}, {loss_scale}); + test.AddInput("gradient_norm", {}, {g_norm}); + test.AddInput("ETA", {}, {eta}); + if (step > 0) { + test.AddInput("Step", {}, {step}); + test.AddOutput("Step_Out", {}, {do_update ? step + 1 : step}); + } else { + test.AddMissingOptionalInput(); + test.AddMissingOptionalOutput(); + } + for (int i = 0; i < group_count; ++i) { + std::string w_name = "W_" + std::to_string(i); + std::string g_name = "G_" + std::to_string(i); + std::string m1_name = "Moment_1_" + std::to_string(i); + std::string m2_name = "Moment_2_" + std::to_string(i); + std::string w_fp16_name = "FP16_W_" + std::to_string(i); + std::string w_new_name = "W_Out_" + std::to_string(i); + std::string g_new_name = "G_Out_" + std::to_string(i); + std::string m1_new_name = "Moment_1_Out_" + std::to_string(i); + std::string m2_new_name = "Moment_2_Out_" + std::to_string(i); + std::string w_fp16_new_name = "FP16_W_Out_" + std::to_string(i); + + test.AddInput(w_name.c_str(), shapes[i], ws[i]); + test.AddInput(g_name.c_str(), shapes[i], gs[i]); + test.AddInput(m1_name.c_str(), shapes[i], ms[i]); + test.AddInput(m2_name.c_str(), shapes[i], vs[i]); + if (!w_halfs.empty() && !w_halfs[i].empty()) { + test.AddInput(w_fp16_name.c_str(), shapes[i], w_halfs[i]); + } else { + test.AddMissingOptionalInput(); + } + + if (!w_news.empty() && !w_news[i].empty()) { + test.AddOutput(w_new_name.c_str(), shapes[i], w_news[i]); + } else { + test.AddMissingOptionalOutput(); + } + if (!g_news.empty() && !g_news[i].empty()) { + test.AddOutput(g_new_name.c_str(), shapes[i], g_news[i]); + } else { + test.AddMissingOptionalOutput(); + } + test.AddOutput(m1_new_name.c_str(), shapes[i], m_news[i]); + test.AddOutput(m2_new_name.c_str(), shapes[i], v_news[i]); + if (!w_new_halfs.empty() && !w_new_halfs[i].empty()) { + test.AddOutput(w_fp16_new_name.c_str(), shapes[i], w_new_halfs[i]); + } else { + test.AddMissingOptionalOutput(); + } + } + + test.AddAttribute("alpha", alphas); + test.AddAttribute("beta", betas); + test.AddAttribute("lambda", lambdas); + test.AddAttribute("epsilon", epsilons); + test.AddAttribute("ratio_min", ratio_min); + test.AddAttribute("ratio_max", ratio_max); + + test.Run(); +} + +// Lamb test without baseline. This function computes +// baseline via an internal function and then invoke +// run_lamb_test_with_baseline(...) to check the result. +void run_multi_tensor_lamb_test( + const std::vector> shapes, + const float eta, + const float loss_scale, + const float g_norm, + const std::vector> ws, + const std::vector> gs, + const std::vector> ms, + const std::vector> vs, + const std::vector lambdas, + const std::vector alphas, + const std::vector betas, + const std::vector epsilons, + const int64_t step = 0, + const float ratio_min = -std::numeric_limits::infinity(), + const float ratio_max = std::numeric_limits::infinity()) { + // Check if parallel vectors have the same length. + ORT_ENFORCE(shapes.size() == ws.size()); + ORT_ENFORCE(shapes.size() == gs.size()); + ORT_ENFORCE(shapes.size() == ms.size()); + ORT_ENFORCE(shapes.size() == vs.size()); + ORT_ENFORCE(shapes.size() == alphas.size()); + ORT_ENFORCE(shapes.size() == betas.size()); + ORT_ENFORCE(shapes.size() == lambdas.size()); + ORT_ENFORCE(shapes.size() == epsilons.size()); + + const int group_count = static_cast(ws.size()); + + // Output buffers of the optimizer. + std::vector> w_news(group_count); + std::vector> g_news(group_count); + std::vector> m_news(group_count); + std::vector> v_news(group_count); + + for (int i = 0; i < group_count; ++i) { + w_news[i] = std::vector(ws[i].size(), 0.f); + g_news[i] = std::vector(gs[i].size(), 0.f); + m_news[i] = std::vector(ms[i].size(), 0.f); + v_news[i] = std::vector(vs[i].size(), 0.f); + + // Invoke LAMB's reference implementation to compute baseline output. + compute_lamb( + shapes[i], ws[i], gs[i], ms[i], vs[i], + eta, loss_scale, g_norm, + lambdas[i], alphas[i], betas[i], epsilons[i], + w_news[i], g_news[i], m_news[i], v_news[i], step, + ratio_min, ratio_max); + } + + // Create tests to make sure the output is correct. + + // Output new weights. + run_multi_tensor_lamb_test_with_baseline( + shapes, eta, loss_scale, g_norm, + ws, gs, ms, vs, + alphas, betas, lambdas, epsilons, + w_news, {}, m_news, v_news, {}, {}, true, step, + ratio_min, ratio_max); + + // Output new gradients. + run_multi_tensor_lamb_test_with_baseline( + shapes, eta, loss_scale, g_norm, + ws, gs, ms, vs, + alphas, betas, lambdas, epsilons, + {}, g_news, m_news, v_news, {}, {}, true, step, + ratio_min, ratio_max); +} + +void run_lamb_mix_precision_test( + const std::vector& shape, + const std::vector& eta, + const std::vector& w, + const std::vector& g, + const std::vector& m, + const std::vector& v, + const float lambda, + const float alpha, + const float beta, + const float epsilon, + const int64_t step = 0) { + std::vector w_new(w.size(), 0); + std::vector g_new(g.size(), 0); + std::vector m_new(m.size(), 0); + std::vector v_new(v.size(), 0); + + // Invoke LAMB's reference implementation to compute output. + compute_lamb( + shape, w, g, m, v, + eta[0], 1.f, 1.f, lambda, alpha, beta, epsilon, + w_new, g_new, m_new, v_new, step); + + std::vector eta_half(eta.size()); + std::vector g_half(w.size()); + std::vector m_half(w.size()); + std::vector v_half(w.size()); + std::vector w_half(w.size()); + ConvertFloatToMLFloat16(eta.data(), eta_half.data(), int(eta.size())); + ConvertFloatToMLFloat16(g.data(), g_half.data(), int(g.size())); + ConvertFloatToMLFloat16(m.data(), m_half.data(), int(m.size())); + ConvertFloatToMLFloat16(v.data(), v_half.data(), int(v.size())); + ConvertFloatToMLFloat16(w.data(), w_half.data(), int(w.size())); + + std::vector m_new_half(m_new.size()); + std::vector v_new_half(v_new.size()); + std::vector w_new_half(w_new.size()); + std::vector g_new_half(g_new.size()); + ConvertFloatToMLFloat16(m_new.data(), m_new_half.data(), int(m_new.size())); + ConvertFloatToMLFloat16(v_new.data(), v_new_half.data(), int(v_new.size())); + ConvertFloatToMLFloat16(w_new.data(), w_new_half.data(), int(w_new.size())); + ConvertFloatToMLFloat16(g_new.data(), g_new_half.data(), int(g_new.size())); + + // Half momentums, without fp16 weight + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w_new, {}, m_new_half, v_new_half, {}, {}, true, step); + + // Float momentums, without fp16 weight + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, {}, {}, true, step); + + // Half momentums, with fp16 weight + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w_new, {}, m_new_half, v_new_half, {}, {}, true, step); + + // Float momentums, with fp16 weight + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, w_half, w_new_half, true, step); + + // Half momentums, with fp16 weight, skip weight update + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m_half, v_half, alpha, beta, lambda, epsilon, w, {}, m_half, v_half, w_half, w_half, false, step); + + // Float momentums, with fp16 weight, skip weight update + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, w, {}, m, v, w_half, w_half, false, step); + + // Float eta, float momentums, with fp16 weight + run_lamb_test_with_baseline( + shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new, w_half, w_new_half, true, step); + + // Float eta, float momentums, with fp16 weight, skip weight update + run_lamb_test_with_baseline( + shape, eta, w, g_half, m, v, alpha, beta, lambda, epsilon, w, {}, m, v, w_half, w_half, false, step); + + // Float momentums, without fp16 weight, output gradients only + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_new_half, m_new, v_new, {}, {}, true, step); + + // Float momentums, with fp16 weight, output gradients only + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_new_half, m_new, v_new, w_half, {}, true, step); + + // Float momentums, with fp16 weight, output gradients only, skip weight update + run_lamb_test_with_baseline( + shape, eta_half, w, g_half, m, v, alpha, beta, lambda, epsilon, {}, g_half, m, v, w_half, {}, false, step); +} + +// A optimizer test with an 2-element vector. +TEST(OptimizerTest, LambOptimizerTestVector) { + // Input tensors and attributes. + const std::vector shape = {2}; + const float eta = 0.5f; + const std::vector w = {1.0f, 2.0f}; + const std::vector g = {3.0f, 4.0f}; + const std::vector m = {-1.0f, -2.0f}; + const std::vector v = {2.0f, 1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}); +} + +TEST(OptimizerTest, LambOptimizerTestVectorWithZeroWeight) { + // Input tensors and attributes. + const std::vector shape = {2}; + const float eta = 0.5f; + const std::vector w = {0.0f, 0.0f}; + const std::vector g = {1.0f, -1.0f}; + const std::vector m = {-1.0f, -2.0f}; + const std::vector v = {2.0f, 1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}); +} + +TEST(OptimizerTest, LambOptimizerRatioMin) { + // Input tensors and attributes. + const std::vector shape = {2}; + const float eta = 0.5f; + const std::vector w = {-1.0f, 1.0f}; + const std::vector g = {1.0f, -1.0f}; + const std::vector m = {-1.0f, -2.0f}; + const std::vector v = {2.0f, 1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + const float ratio_min = -std::numeric_limits::infinity(); + const float ratio_max = 0.1f; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}, + 0, + ratio_min, + ratio_max); +} + +TEST(OptimizerTest, LambOptimizerRatioMax) { + // Input tensors and attributes. + const std::vector shape = {2}; + const float eta = 0.5f; + const std::vector w = {0.0001f, -0.0001f}; + const std::vector g = {1.0f, -1.0f}; + const std::vector m = {-1.0f, -2.0f}; + const std::vector v = {2.0f, 1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + const float ratio_min = 1.0f; + const float ratio_max = std::numeric_limits::infinity(); + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}, + 0, + ratio_min, + ratio_max); +} + +TEST(OptimizerTest, LambOptimizerTestBiasCorrectionFirst) { + // Input tensors and attributes. + const std::vector shape = {2}; + const float eta = 0.5f; + const std::vector w = {1.0f, 2.0f}; + const std::vector g = {3.0f, 4.0f}; + const std::vector m = {-1.0f, -2.0f}; + const std::vector v = {2.0f, 1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}, + 1); +} + +TEST(OptimizerTest, LambOptimizerTestBiasCorrectionThird) { + // Input tensors and attributes. + const std::vector shape = {2}; + const float eta = 0.5f; + const std::vector w = {1.0f, 2.0f}; + const std::vector g = {3.0f, 4.0f}; + const std::vector m = {-1.0f, -2.0f}; + const std::vector v = {2.0f, 1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}, + 3); +} + +// A optimizer test with an 2-by-1-by-1-by-1 tensor. +TEST(OptimizerTest, LambOptimizerTest4DTensor) { + // Input tensors and attributes. + const std::vector shape = {2, 1, 1, 1}; + const float eta = 0.5f; + const std::vector w = {1.0f, 2.0f}; + const std::vector g = {3.0f, 4.0f}; + const std::vector m = {-1.0f, -2.0f}; + const std::vector v = {2.0f, 1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}); +} + +// A optimizer test with an 2-by-3 tensor. +TEST(OptimizerTest, LambOptimizerTest2by3Tensor) { + // Input tensors and attributes. + const std::vector shape = {2, 3}; + const float eta = 0.5f; + const std::vector w = {1.0f, 2.0f, 1.0f, 1.0f, 2.0f, 2.0f}; + const std::vector g = {3.0f, 4.0f, 3.0f, 3.0f, 4.0f, 4.0f}; + const std::vector m = {-1.0f, -2.0f, 2.0f, 1.0f, 1.0f, -2.0f}; + const std::vector v = {1.0f, 1.0f, 5.0f, 5.0f, 6.0f, 6.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}); +} + +// A optimizer test with an 1-element tensor. +TEST(OptimizerTest, LambOptimizerTestScalar) { + // Input tensors and attributes. + const std::vector shape = {(int64_t)1}; + const float eta = 0.5f; + const std::vector w = {1.0f}; + const std::vector g = {3.0f}; + const std::vector m = {-10.0f}; + const std::vector v = {1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + // Intermediate and output buffers of the optimizer. + std::vector m_new = {0.0f}; + std::vector v_new = {0.0f}; + std::vector w_new = {0.0f}; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}); +} + +TEST(OptimizerTest, LambOptimizerTestScalarScaling) { + // Input tensors and attributes. + const std::vector shape = {(int64_t)1}; + const float eta = 0.5f; + const std::vector w = {1.0f}; + const std::vector g = {3.0f}; + const std::vector m = {-10.0f}; + const std::vector v = {1.0f}; + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + // Intermediate and output buffers of the optimizer. + std::vector m_new = {0.0f}; + std::vector v_new = {0.0f}; + std::vector w_new = {0.0f}; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 8.f, + 4.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}); +} + +TEST(OptimizerTest, LambOptimizerTestExternalBaseline) { + // Input tensors and attributes. + const std::vector shape = {2, 5}; + const std::vector eta = {0.1f}; + const std::vector w = { + 0.01379026f, 0.15308191f, -0.24356517f, -0.21798165f, -0.13770047f, 0.09694599f, + -0.02223516f, 0.2664228f, -0.01177993f, 0.06832688f}; + const std::vector g = { + -6.048543f, 10.569487f, -9.207029f, -0.57407373f, + 5.884985f, -0.21047728f, 3.539946f, -5.957566f, -9.343748f, 1.1502024f}; + const std::vector m = { + -5.9078765f, 9.673933f, -8.731428f, -0.6227454f, 5.284312f, -0.27138948f, + 3.443532f, -5.681713f, -8.72421f, 1.1441823f}; + const std::vector v = { + 4.2659229e+01f, 1.1438165e+02f, 9.3179581e+01f, 4.7399229e-01f, 3.4129276e+01f, + 9.0019435e-02f, 1.4493006e+01f, 3.9455612e+01f, 9.3025581e+01f, 1.6000764e+0f}; + + const float lambda = 0.1f; + const float alpha = 0.1f; + const float beta = 0.01f; + const float epsilon = 0.1f; + + std::vector w_new = { + 0.02979828f, 0.13677707f, -0.22708717f, -0.20361158f, -0.15338624f, 0.1081504f, + -0.03804127f, 0.28198114f, 0.00430069f, 0.05319814f}; + std::vector g_new = { + 0.01600802f, -0.01630484f, 0.01647800f, 0.01437007f, -0.01568577f, 0.01120441f, + -0.01580611f, 0.01555834f, 0.01608062f, -0.01512874f}; + std::vector m_new = { + -6.0344763f, 10.479931f, -9.15947f, -0.57894087f, 5.824918f, -0.2165685f, + 3.5303047f, -5.9299808f, -9.281795f, 1.1496004f}; + std::vector v_new = { + 3.6645618e+01f, 1.1174072e+02f, 8.4853485e+01f, 3.3100498e-01f, 3.4628010e+01f, + 4.4757873e-02f, 1.2550836e+01f, 3.5532223e+01f, 8.7362823e+01f, 1.3257366e+00f}; + + // Output new weights + run_lamb_test_with_baseline( + shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new); + + // Output new gradients + run_lamb_test_with_baseline( + shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, {}, g_new, m_new, v_new); +} + +TEST(OptimizerTest, LambOptimizerTestExternalBaselineDouble) { + // Input tensors and attributes. + const std::vector shape = {2, 5}; + const std::vector eta = {0.1f}; + const std::vector w = { + 0.01379026, 0.15308191, -0.24356517, -0.21798165, -0.13770047, 0.09694599, + -0.02223516, 0.2664228, -0.01177993, 0.06832688}; + const std::vector g = { + -6.048543, 10.569487, -9.207029, -0.57407373, + 5.884985, -0.21047728, 3.539946, -5.957566, -9.343748, 1.1502024}; + const std::vector m = { + -5.9078765, 9.673933, -8.731428, -0.6227454, 5.284312, -0.27138948, + 3.443532, -5.681713, -8.72421, 1.1441823}; + const std::vector v = { + 4.2659229e+01, 1.1438165e+02, 9.3179581e+01, 4.7399229e-01, 3.4129276e+01, + 9.0019435e-02, 1.4493006e+01, 3.9455612e+01, 9.3025581e+01, 1.6000764e+0}; + + const float lambda = 0.1f; + const float alpha = 0.1f; + const float beta = 0.01f; + const float epsilon = 0.1f; + + std::vector w_new = { + 0.02979828, 0.13677707, -0.22708717, -0.20361158, -0.15338624, 0.1081504, + -0.03804127, 0.28198114, 0.00430069, 0.05319814}; + std::vector g_new = { + 0.01600802, -0.01630484, 0.016478, 0.01437007, -0.01568577, 0.01120441, + -0.01580611, 0.01555834, 0.01608062, -0.01512874}; + std::vector m_new = { + -6.0344763, 10.479931, -9.15947, -0.57894087, 5.824918, -0.2165685, + 3.5303047, -5.9299808, -9.281795, 1.1496004}; + std::vector v_new = { + 3.6645618e+01, 1.1174072e+02, 8.4853485e+01, 3.3100498e-01, 3.4628010e+01, + 4.4757873e-02, 1.2550836e+01, 3.5532223e+01, 8.7362823e+01, 1.3257366e+00}; + + // Output new weights + run_lamb_test_with_baseline( + shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, w_new, {}, m_new, v_new); + + // Output new gradients + run_lamb_test_with_baseline( + shape, eta, w, g, m, v, alpha, beta, lambda, epsilon, {}, g_new, m_new, v_new); +} + +TEST(OptimizerTest, LambOptimizerTest5DTensorMixPrecision32_16) { + const std::vector shape = {2, 2, 2, 1, 1}; + const std::vector eta = {0.5f}; + const std::vector w = {1.0f, 2.0f, 2.5f, 1.5f, 1.0f, 2.0f, 2.0f, 1.5f}; + const std::vector g = {-1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 0.8f}; + const std::vector m = {1.0f, 2.0f, -0.25f, 1.1f, 1.0f, 2.0f, -0.21f, 1.1f}; + const std::vector v = {1.5f, 1.0f, 1.1f, 0.76f, 1.5f, 1.0f, 1.5f, 0.76f}; + + const float lambda = 1.5f; + const float alpha = 1.5f; + const float beta = 1.5f; + const float epsilon = 1.0f; + + run_lamb_mix_precision_test( + shape, eta, w, g, m, v, + lambda, alpha, beta, epsilon); +} + +TEST(OptimizerTest, LambOptimizerTestSimpleBaselineMixPrecision32_16) { + const std::vector shape = {2, 1}; + const std::vector eta = {1.0f}; + const std::vector w = {1.0f, 1.0f}; + const std::vector g = {-1.0f, 1.0f}; + const std::vector m = {1.0f, 1.0f}; + const std::vector v = {0.0f, 0.0f}; + + const float lambda = 0.0f; + const float alpha = 1.0f; + const float beta = 1.0f; + const float epsilon = 1.0f; + + run_lamb_mix_precision_test( + shape, eta, w, g, m, v, + lambda, alpha, beta, epsilon); +} + +TEST(OptimizerTest, LambOptimizerTestBaselineMixPrecision32_16) { + const std::vector shape = {2, 1}; + const std::vector eta = {0.1f}; + const std::vector w = {-1.5f, 2.4f}; + const std::vector g = {-0.75f, 1.2f}; + const std::vector m = {0.87f, -0.94f}; + const std::vector v = {0.12f, 0.28f}; + + const float lambda = 0.25f; + const float alpha = 0.9f; + const float beta = 0.95f; + const float epsilon = 0.33f; + + run_lamb_mix_precision_test( + shape, eta, w, g, m, v, + lambda, alpha, beta, epsilon); +} + +TEST(OptimizerTest, LambOptimizerTestScalarMixPrecision32_16) { + const std::vector shape = {1}; + const std::vector eta = {0.1f}; + const std::vector w = {-1.5f}; + const std::vector g = {-0.75f}; + const std::vector m = {0.87f}; + const std::vector v = {0.12f}; + + const float lambda = 0.25f; + const float alpha = 0.9f; + const float beta = 0.95f; + const float epsilon = 0.33f; + + run_lamb_mix_precision_test( + shape, eta, w, g, m, v, + lambda, alpha, beta, epsilon); + + run_lamb_mix_precision_test( + shape, eta, w, g, m, v, + lambda, alpha, beta, epsilon, 2); +} + +TEST(OptimizerTest, LambOptimizerTestLarge) { + // Input tensors and attributes. + for (const auto& size : {55667, 1944006, 3907584}) { + const std::vector shape = {static_cast(size)}; + const float eta = 0.5f; + std::vector w(size); + std::vector g(size); + std::vector m(size); + std::vector v(size); + + std::random_device random_device; + std::mt19937 random_engine(0); + std::uniform_real_distribution dist(0.1f, 1.0f); + for (int i = 0; i < size; ++i) { + w[i] = dist(random_engine); + g[i] = dist(random_engine); + m[i] = dist(random_engine); + v[i] = dist(random_engine); + } + + const float lambda = 0.5f; + const float alpha = 0.2f; + const float beta = 0.8f; + const float epsilon = 1e-6f; + + run_multi_tensor_lamb_test( + {shape}, + eta, + 1.f, + 1.f, + {w}, + {g}, + {m}, + {v}, + {lambda}, + {alpha}, + {beta}, + {epsilon}); + } +} + +TEST(OptimizerTest, LambOptimizerMultiTensorRatio) { + const int group_count = 127; + std::random_device random_device; + std::mt19937 random_engine(0); + std::uniform_real_distribution dist(0.1f, 1.0f); + std::uniform_int_distribution dist_int(1, 1228); + + std::vector sizes(group_count); + std::vector> shapes(group_count); + + std::vector> ws(group_count); + std::vector> gs(group_count); + std::vector> ms(group_count); + std::vector> vs(group_count); + std::vector alphas(group_count); + std::vector betas(group_count); + std::vector lambdas(group_count); + std::vector epsilons(group_count); + + const float eta = dist(random_engine); + + for (int64_t i = 0; i < group_count; ++i) { + const auto size = dist_int(random_engine); + sizes[i] = size; + shapes[i] = std::vector(1, size); + + ws[i] = std::vector(sizes[i]); + gs[i] = std::vector(sizes[i]); + ms[i] = std::vector(sizes[i]); + vs[i] = std::vector(sizes[i]); + + for (int64_t j = 0; j < sizes[i]; ++j) { + ws[i][j] = dist(random_engine); + gs[i][j] = dist(random_engine); + ms[i][j] = dist(random_engine); + vs[i][j] = dist(random_engine); + } + + alphas[i] = dist(random_engine); + betas[i] = dist(random_engine); + lambdas[i] = dist(random_engine); + epsilons[i] = dist(random_engine); + } + + run_multi_tensor_lamb_test( + shapes, eta, 1.f, 1.f, + ws, gs, ms, vs, + lambdas, alphas, betas, epsilons, + 0, 0.3f, 0.7f); + + run_multi_tensor_lamb_test( + shapes, eta, 1.f, 1.f, + ws, gs, ms, vs, + lambdas, alphas, betas, epsilons); +} +#endif + +} // namespace test +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cpu/optimizer/optimizers.cc b/orttraining/orttraining/training_ops/cpu/optimizer/optimizers.cc index c48d98e511..ef050ff30f 100644 --- a/orttraining/orttraining/training_ops/cpu/optimizer/optimizers.cc +++ b/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::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(NM2) / beta_correction).sqrt() + epsilon_; - const auto& update = ( (MakeEigenArrayMap(NM1) / alpha_correction) / denom) + (lambda_ * MakeEigenArrayMap(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(NM2) / beta_correction).sqrt() + epsilon_; + const auto& update = ( (MakeEigenArrayMap(NM1) / alpha_correction) / denom) + (lambda_ * MakeEigenArrayMap(W)); + const auto& delta = -eta * update; - // Weight, gradient, and step update. - if (NG != nullptr) { - MakeEigenArrayMap(*NG) = delta; + // Weight, gradient, and step update. + if (NG != nullptr) { + MakeEigenArrayMap(*NG) = delta; + } + if (NW != nullptr) { + MakeEigenArrayMap(*NW) = MakeEigenArrayMap(W) + delta; + } } - if (NW != nullptr) { - MakeEigenArrayMap(*NW) = MakeEigenArrayMap(W) + delta; + else if (weight_decay_mode_ == 1) { + const auto& denom = MakeEigenArrayMap(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(NM1) / denom + - eta * lambda_ * (MakeEigenArrayMap(W) - step_size * MakeEigenArrayMap(NM1) / denom); + + // Weight, gradient, and step update. + if (NG != nullptr) { + MakeEigenArrayMap(*NG) = delta; + } + if (NW != nullptr) { + MakeEigenArrayMap(*NW) = MakeEigenArrayMap(W) + delta; + } } + else { + // Shouldn't reach here + ORT_THROW("Unsupported Adamw optimizer mode."); + } + *NS.template MutableData() = step + 1; - return Status::OK(); } diff --git a/orttraining/orttraining/training_ops/cpu/optimizer/optimizers.h b/orttraining/orttraining/training_ops/cpu/optimizer/optimizers.h index 759a9a6de1..c2c732e510 100644 --- a/orttraining/orttraining/training_ops/cpu/optimizer/optimizers.h +++ b/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(0); ORT_ENFORCE(info.GetAttr("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(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 diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc index c4de4d91b1..e1a15bc4d9 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cc +++ b/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::ComputeInternal(OpKer ToCudaType::FromFloat(lambda_), ToCudaType::FromFloat(epsilon_), do_bias_correction_, + weight_decay_mode_, reinterpret_cast(NM1.template MutableData()), reinterpret_cast(NM2.template MutableData()), NW != nullptr ? reinterpret_cast(NW->template MutableData()) : nullptr, diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu b/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu index c50d8740bb..0e72c60c3f 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.cu +++ b/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 -__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 +__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(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(weights_out[id]); + } + } + + moment_1_out[id] = m1o; + moment_2_out[id] = m2o; +} + template 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<<>>( + + // 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<<>>( eta, weights, grads, @@ -117,6 +193,33 @@ void AdamOptimizerImpl( grads_out, fp16_weights_out, N); + } + else if (weight_decay_mode == 1) { + _AdamOptimizer_mode1<<>>( + 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, \ diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/adam.h b/orttraining/orttraining/training_ops/cuda/optimizer/adam.h index 1df9280975..fcfa617dbc 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/adam.h +++ b/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("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(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 diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc index c632fc8fa8..edf93b433a 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cc +++ b/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 { diff --git a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu b/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu index b035872500..94832b484f 100644 --- a/orttraining/orttraining/training_ops/cuda/optimizer/lamb.cu +++ b/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