feat: add validation tests for learning rate schedulers

kinorw · kinorw · commit 2980f930814b · 2026-03-21T04:31:03.000+08:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -220,4 +220,7 @@ add_executable(test_sequential_lr test/lr_scheduler/test_sequential_lr.cc)
 target_link_libraries(test_sequential_lr infini_train)
 
 add_executable(test_chained_lr test/lr_scheduler/test_chained_lr.cc)
-target_link_libraries(test_chained_lr infini_train)
+target_link_libraries(test_chained_lr infini_train)
+
+add_executable(test_lr_scheduler_validation test/lr_scheduler/test_lr_scheduler_validation.cc)
+target_link_libraries(test_lr_scheduler_validation infini_train)
diff --git a/infini_train/include/lr_scheduler.h b/infini_train/include/lr_scheduler.h
@@ -67,6 +67,8 @@ class LRScheduler {
     virtual StateDict State() const;
     virtual void LoadState(const StateDict &state);
 
+    bool SharesOptimizerWith(const std::shared_ptr<Optimizer> &opt) const ;
+
 protected:
     virtual float GetClosedFormLR() const = 0;
     virtual float GetChainedFormLR() const;
diff --git a/infini_train/src/lr_scheduler.cc b/infini_train/src/lr_scheduler.cc
@@ -102,6 +102,8 @@ float LRScheduler::BaseLR() const { return base_lr_; }
 
 int64_t LRScheduler::LastStep() const { return last_step_; }
 
+bool LRScheduler::SharesOptimizerWith(const std::shared_ptr<Optimizer> &opt) const { return optimizer_ == opt; }
+
 void LRScheduler::ResetStep(int64_t step) { last_step_ = step; }
 
 StateDict LRScheduler::State() const {
@@ -126,7 +128,10 @@ namespace lr_schedulers {
 // --- ConstantLR ---
 
 ConstantLR::ConstantLR(std::shared_ptr<Optimizer> optimizer, float factor, int total_iters, int64_t last_step)
-    : LRScheduler(std::move(optimizer), last_step), factor_(factor), total_iters_(total_iters) {}
+    : LRScheduler(std::move(optimizer), last_step), factor_(factor), total_iters_(total_iters) {
+        CHECK_GE(factor_, 0.0f) << "ConstantLR: factor must be >= 0.";
+        CHECK_LE(factor_, 1.0f) << "ConstantLR: factor must be <= 1.";
+    }
 
 float ConstantLR::GetClosedFormLR() const { return last_step_ < total_iters_ ? base_lr_ * factor_ : base_lr_; }
 
@@ -145,7 +150,9 @@ float ConstantLR::GetChainedFormLR() const {
 // --- StepLR ---
 
 StepLR::StepLR(std::shared_ptr<Optimizer> optimizer, int64_t step_size, float gamma, int64_t last_step)
-    : LRScheduler(std::move(optimizer), last_step), step_size_(step_size), gamma_(gamma) {}
+    : LRScheduler(std::move(optimizer), last_step), step_size_(step_size), gamma_(gamma) {
+    CHECK_GT(step_size_, 0) << "StepLR: step_size must be > 0.";
+}
 
 float StepLR::GetClosedFormLR() const {
     return base_lr_
@@ -163,7 +170,13 @@ float StepLR::GetChainedFormLR() const {
 LinearLR::LinearLR(std::shared_ptr<Optimizer> optimizer, float start_factor, float end_factor, int64_t total_iters,
                    int64_t last_step)
     : LRScheduler(std::move(optimizer), last_step), start_factor_(start_factor), end_factor_(end_factor),
-      total_iters_(total_iters) {}
+      total_iters_(total_iters) {
+        CHECK_GT(start_factor_, 0.0f) << "LinearLR: start_factor must be > 0.";
+        CHECK_LE(start_factor_, 1.0f) << "LinearLR: start_factor must be <= 1.";
+        CHECK_GE(end_factor_, 0.0f) << "LinearLR: end_factor must be >= 0.";
+        CHECK_LE(end_factor_, 1.0f) << "LinearLR: end_factor must be <= 1.";
+                CHECK_GT(total_iters_, 0) << "LinearLR: total_iters must be > 0.";
+      }
 
 float LinearLR::GetClosedFormLR() const {
     if (last_step_ >= total_iters_) {
@@ -196,24 +209,35 @@ float LinearLR::GetChainedFormLR() const {
 }
 
 LambdaLR::LambdaLR(std::shared_ptr<Optimizer> optimizer, std::function<float(int64_t)> lr_lambda, int64_t last_step)
-    : LRScheduler(std::move(optimizer), last_step), lr_lambda_(std::move(lr_lambda)) {}
+    : LRScheduler(std::move(optimizer), last_step), lr_lambda_(std::move(lr_lambda)) {
+        CHECK(lr_lambda_) << "LambdaLR: lr_lambda must not be null.";
+    }
 
 float LambdaLR::GetClosedFormLR() const { return base_lr_ * lr_lambda_(last_step_); }
 
 SequentialLR::SequentialLR(std::shared_ptr<Optimizer> optimizer, std::vector<std::shared_ptr<LRScheduler>> schedulers,
                            std::vector<int64_t> milestones, int64_t last_step)
     : LRScheduler(std::move(optimizer), last_step), schedulers_(std::move(schedulers)),
-      milestones_(std::move(milestones)) {}
-
-void SequentialLR::InitialStep() {
+      milestones_(std::move(milestones)) {
     CHECK(!schedulers_.empty()) << "SequentialLR requires at least one scheduler.";
+
+    for (size_t i = 0; i < schedulers_.size(); ++i) {
+        CHECK(schedulers_[i]) << "SequentialLR: scheduler at index " << i << " must not be null.";
+        CHECK(schedulers_[i]->SharesOptimizerWith(optimizer_))
+            << "SequentialLR: scheduler at index " << i << " must share the same optimizer.";
+    }
+
     CHECK_EQ(milestones_.size(), schedulers_.size() - 1)
         << "SequentialLR: milestones count must be schedulers count - 1.";
 
     for (size_t i = 1; i < milestones_.size(); ++i) {
         CHECK_GT(milestones_[i], milestones_[i - 1]) << "Milestones must be strictly increasing.";
     }
 
+}
+
+void SequentialLR::InitialStep() {
+
     optimizer_->SetLearningRate(schedulers_[0]->BaseLR());
 
     UndoChildInitialSteps();
@@ -279,11 +303,17 @@ void SequentialLR::LoadState(const StateDict &state) {
 
 ChainedScheduler::ChainedScheduler(std::shared_ptr<Optimizer> optimizer,
                                    std::vector<std::shared_ptr<LRScheduler>> schedulers, int64_t last_step)
-    : LRScheduler(std::move(optimizer), last_step), schedulers_(std::move(schedulers)) {}
-
-void ChainedScheduler::InitialStep() {
+    : LRScheduler(std::move(optimizer), last_step), schedulers_(std::move(schedulers)) {
     CHECK(!schedulers_.empty()) << "ChainedScheduler requires at least one scheduler.";
 
+    for (size_t i = 0; i < schedulers_.size(); ++i) {
+        CHECK(schedulers_[i]) << "ChainedScheduler: scheduler at index " << i << " must not be null.";
+        CHECK(schedulers_[i]->SharesOptimizerWith(optimizer_))
+            << "ChainedScheduler: scheduler at index " << i << " must share the same optimizer.";
+    }
+}
+
+void ChainedScheduler::InitialStep() {
 }
 
 void ChainedScheduler::Step() {
diff --git a/test/lr_scheduler/test_lr_scheduler_validation.cc b/test/lr_scheduler/test_lr_scheduler_validation.cc
@@ -0,0 +1,161 @@
+#include <functional>
+#include <iostream>
+#include <memory>
+#include <sys/wait.h>
+#include <unistd.h>
+#include <vector>
+
+#include "infini_train/include/lr_scheduler.h"
+#include "test/lr_scheduler/test_helpers.h"
+
+using namespace infini_train;
+using namespace infini_train::lr_schedulers;
+
+namespace {
+
+bool ExpectDeath(const std::function<void()> &fn) {
+    pid_t pid = fork();
+    if (pid == -1) {
+        return false;
+    }
+
+    if (pid == 0) {
+        fn();
+        _exit(0);
+    }
+
+    int status = 0;
+    if (waitpid(pid, &status, 0) == -1) {
+        return false;
+    }
+
+    return !WIFEXITED(status) || WEXITSTATUS(status) != 0;
+}
+
+void TestStepLRRejectsNonPositiveStepSize() {
+    ASSERT_TRUE(ExpectDeath([] {
+        auto opt = MakeDummyOptimizer(0.1f);
+        auto sched = CreateLRScheduler(opt, {
+                                                .type = "step",
+                                                .step_size = 0,
+                                                .step_gamma = 0.1f,
+                                            });
+        (void)sched;
+    }));
+}
+
+void TestLinearLRRejectsNonPositiveTotalIters() {
+    ASSERT_TRUE(ExpectDeath([] {
+        auto opt = MakeDummyOptimizer(0.1f);
+        auto sched = CreateLRScheduler(opt, {
+                                                .type = "linear",
+                                                .linear_start_factor = 0.5f,
+                                                .linear_end_factor = 1.0f,
+                                                .linear_total_iters = 0,
+                                            });
+        (void)sched;
+    }));
+}
+
+void TestLambdaLRRejectsNullLambda() {
+    ASSERT_TRUE(ExpectDeath([] {
+        auto opt = MakeDummyOptimizer(0.1f);
+        auto sched = CreateLRScheduler(opt, {
+                                                .type = "lambda",
+                                            });
+        (void)sched;
+    }));
+}
+
+void TestSequentialLRRejectsMismatchedOptimizer() {
+    ASSERT_TRUE(ExpectDeath([] {
+        auto opt1 = MakeDummyOptimizer(0.1f);
+        auto opt2 = MakeDummyOptimizer(0.1f);
+
+        auto s1 = CreateLRScheduler(opt1, {
+                                              .type = "linear",
+                                              .linear_start_factor = 0.5f,
+                                              .linear_end_factor = 1.0f,
+                                              .linear_total_iters = 2,
+                                          });
+        auto s2 = CreateLRScheduler(opt2, {
+                                              .type = "step",
+                                              .step_size = 2,
+                                              .step_gamma = 0.5f,
+                                          });
+
+        auto sched = LRScheduler::Create<SequentialLR>(
+            opt1, std::vector<std::shared_ptr<LRScheduler>>{s1, s2}, std::vector<int64_t>{1});
+        (void)sched;
+    }));
+}
+
+void TestSequentialLRRejectsNullChild() {
+    ASSERT_TRUE(ExpectDeath([] {
+        auto opt = MakeDummyOptimizer(0.1f);
+        auto sched = LRScheduler::Create<SequentialLR>(opt, std::vector<std::shared_ptr<LRScheduler>>{nullptr},
+                                                       std::vector<int64_t>{});
+        (void)sched;
+    }));
+}
+
+void TestChainedSchedulerRejectsEmptyChildren() {
+    ASSERT_TRUE(ExpectDeath([] {
+        auto opt = MakeDummyOptimizer(0.1f);
+        auto sched = LRScheduler::Create<ChainedScheduler>(opt, std::vector<std::shared_ptr<LRScheduler>>{});
+        (void)sched;
+    }));
+}
+
+void TestChainedSchedulerRejectsMismatchedOptimizer() {
+    ASSERT_TRUE(ExpectDeath([] {
+        auto opt1 = MakeDummyOptimizer(0.1f);
+        auto opt2 = MakeDummyOptimizer(0.1f);
+
+        auto s1 = CreateLRScheduler(opt1, {
+                                              .type = "step",
+                                              .step_size = 2,
+                                              .step_gamma = 0.5f,
+                                          });
+        auto s2 = CreateLRScheduler(opt2, {
+                                              .type = "constant",
+                                              .constant_factor = 0.5f,
+                                              .constant_total_iters = 2,
+                                          });
+
+        auto sched = LRScheduler::Create<ChainedScheduler>(opt1, std::vector<std::shared_ptr<LRScheduler>>{s1, s2});
+        (void)sched;
+    }));
+}
+
+void TestChainedSchedulerRejectsNullChild() {
+    ASSERT_TRUE(ExpectDeath([] {
+        auto opt = MakeDummyOptimizer(0.1f);
+        auto sched = LRScheduler::Create<ChainedScheduler>(opt, std::vector<std::shared_ptr<LRScheduler>>{nullptr});
+        (void)sched;
+    }));
+}
+
+} // namespace
+
+int main(int argc, char *argv[]) {
+    google::InitGoogleLogging(argv[0]);
+
+    std::cout << "=== LR Scheduler Validation Tests ===" << std::endl;
+    TestStepLRRejectsNonPositiveStepSize();
+    TestLinearLRRejectsNonPositiveTotalIters();
+    TestLambdaLRRejectsNullLambda();
+    TestSequentialLRRejectsMismatchedOptimizer();
+    TestSequentialLRRejectsNullChild();
+    TestChainedSchedulerRejectsEmptyChildren();
+    TestChainedSchedulerRejectsMismatchedOptimizer();
+    TestChainedSchedulerRejectsNullChild();
+
+    if (g_fail_count == 0) {
+        std::cout << "All Tests PASSED" << std::endl;
+    } else {
+        std::cout << g_fail_count << " test(s) FAILED" << std::endl;
+    }
+
+    return g_fail_count > 0 ? 1 : 0;
+}
