refactor(checkpoint): centralize config, simplify prune, fix optimizer test

Author: JYMiracle305

CMakeLists.txt

@@ -194,7 +194,6 @@ add_executable(gpt2
   example/gpt2/main.cc
   example/common/tiny_shakespeare_dataset.cc
   example/common/utils.cc
-  example/common/checkpoint_loader.cc
   example/common/tokenizer.cc
   example/gpt2/checkpoint_loader.cc
 )
@@ -204,7 +203,6 @@ add_executable(llama3
   example/llama3/main.cc
   example/common/tiny_shakespeare_dataset.cc
   example/common/utils.cc
-  example/common/checkpoint_loader.cc
   example/common/tokenizer.cc
   example/llama3/checkpoint_loader.cc
 )

example/gpt2/main.cc

@@ -11,7 +11,7 @@
 #include "glog/logging.h"
 
 #include "infini_train/include/autocast.h"
-#include "infini_train/include/checkpoint.h"
+#include "infini_train/include/checkpoint/checkpoint.h"
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
@@ -31,12 +31,12 @@
 #ifdef PROFILE_MODE
 #include "infini_train/include/profiler.h"
 #endif
+#include "infini_train/include/checkpoint/checkpoint_manager.h"
 #include "infini_train/include/nn/parallel/utils.h"
 #include "infini_train/include/utils/global_module_hook_registry.h"
 #include "infini_train/include/utils/precision_check_config.h"
 #include "infini_train/include/utils/precision_checker.h"
 
-#include "example/common/checkpoint_loader.h"
 #include "example/common/tiny_shakespeare_dataset.h"
 #include "example/common/tokenizer.h"
 #include "example/gpt2/checkpoint_loader.h"
@@ -85,7 +85,7 @@ DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables savin
 DEFINE_string(load, "", "checkpoint directory to resume from");
 DEFINE_string(save, "./checkpoints", "root directory used to store checkpoints");
 DEFINE_uint32(max_checkpoint_keep, 3, "max number of checkpoint steps to keep");
-DEFINE_bool(no_save_optim, false, "whether optimizer state is persisted in checkpoints");
+DEFINE_bool(save_optimizer_state, true, "whether optimizer state is persisted in checkpoints");
 // precision check
 DEFINE_string(
     precision_check, "",
@@ -331,7 +331,8 @@ void Train(const nn::parallel::Rank &rank) {
                                                      .model = model,
                                                      .optimizer = optimizer,
                                                      .model_config = model_config,
-                                                     .state = state});
+                                                     .state = state,
+                                                     .load_optimizer_state = false});
     start_step = resume_result.global_step;
     size_t consumed_batches = resume_result.consumed_batches;
 
@@ -345,31 +346,29 @@ void Train(const nn::parallel::Rank &rank) {
         for (size_t i = 0; i < num_skips; ++i) { ++train_iter; }
     }
 
-    auto save_checkpoint
-        = [&](const std::filesystem::path &save_dir, int64_t global_step, bool prune_step_checkpoints) {
-              SaveCheckpoint({
-                  .save_dir = save_dir,
-                  .global_step = global_step,
-                  .consumed_batches = consumed_batches,
-                  .last_lr = FLAGS_learning_rate,
-                  .n_layer = model_config.n_layer,
-                  .n_head = model_config.n_head,
-                  .n_kv_head = model_config.n_kv_head,
-                  .n_embd = model_config.n_embd,
-                  .vocab_size = model_config.vocab_size,
-                  .ddp_size = ddp_world_size,
-                  .tp_size = tp_world_size,
-                  .sp_size = sp_world_size,
-                  .pp_size = pp_world_size,
-                  .no_save_optim = FLAGS_no_save_optim,
-                  .prune_step_checkpoints = prune_step_checkpoints,
-                  .checkpoint_root_dir = FLAGS_save,
-                  .max_checkpoint_keep = FLAGS_max_checkpoint_keep,
-                  .rank = rank,
-                  .model = *model,
-                  .optimizer = *optimizer,
-              });
-          };
+    auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) {
+        SaveCheckpoint({
+            .save_dir = save_dir,
+            .global_step = global_step,
+            .consumed_batches = consumed_batches,
+            .last_lr = FLAGS_learning_rate,
+            .n_layer = model_config.n_layer,
+            .n_head = model_config.n_head,
+            .n_kv_head = model_config.n_kv_head,
+            .n_embd = model_config.n_embd,
+            .vocab_size = model_config.vocab_size,
+            .ddp_size = ddp_world_size,
+            .tp_size = tp_world_size,
+            .sp_size = sp_world_size,
+            .pp_size = pp_world_size,
+            .save_optimizer_state = FLAGS_save_optimizer_state,
+            .checkpoint_root_dir = FLAGS_save,
+            .max_checkpoint_keep = FLAGS_max_checkpoint_keep,
+            .rank = rank,
+            .model = *model,
+            .optimizer = *optimizer,
+        });
+    };
 
     LOG(INFO) << "start training";
 
@@ -496,7 +495,7 @@ void Train(const nn::parallel::Rank &rank) {
             if (rank.IsParallel()) {
                 step_dir /= std::format("rank_{:06d}", rank.GlobalRank());
             }
-            save_checkpoint(step_dir, step + 1, true);
+            save_checkpoint(step_dir, step + 1);
         }
     }
 
@@ -510,7 +509,7 @@ void Train(const nn::parallel::Rank &rank) {
     if (rank.IsParallel()) {
         final_dir /= std::format("rank_{:06d}", rank.GlobalRank());
     }
-    save_checkpoint(final_dir, FLAGS_num_iteration, false);
+    save_checkpoint(final_dir, FLAGS_num_iteration);
 
 #ifdef PROFILE_MODE
     Profiler::Instance().Report("gpt2.report", Profiler::SortBy::DeviceTimePercentage);

example/llama3/main.cc

@@ -9,7 +9,8 @@
 #include "glog/logging.h"
 
 #include "infini_train/include/autocast.h"
-#include "infini_train/include/checkpoint.h"
+#include "infini_train/include/checkpoint/checkpoint.h"
+#include "infini_train/include/checkpoint/checkpoint_manager.h"
 #include "infini_train/include/core/runtime/device_guard.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/device.h"
@@ -35,7 +36,6 @@
 #include "infini_train/include/profiler.h"
 #endif
 
-#include "example/common/checkpoint_loader.h"
 #include "example/common/tiny_shakespeare_dataset.h"
 #include "example/common/tokenizer.h"
 #include "example/llama3/checkpoint_loader.h"
@@ -83,7 +83,7 @@ DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables savin
 DEFINE_string(load, "", "checkpoint directory to resume from");
 DEFINE_string(save, "./checkpoints", "root directory used to store checkpoints");
 DEFINE_uint32(max_checkpoint_keep, 3, "max number of checkpoint steps to keep");
-DEFINE_bool(no_save_optim, false, "whether optimizer state is persisted in checkpoints");
+DEFINE_bool(save_optimizer_state, true, "whether optimizer state is persisted in checkpoints");
 
 // precision check
 DEFINE_string(
@@ -305,14 +305,13 @@ void Train(const nn::parallel::Rank &rank) {
 
     int start_step = 0;
     TrainerState state;
-    const auto resume_result = ResumeFromCheckpoint({
-        .resume_root = FLAGS_load,
-        .rank = rank,
-        .model = model,
-        .optimizer = optimizer,
-        .model_config = model_config,
-        .state = state,
-    });
+    const auto resume_result = ResumeFromCheckpoint({.resume_root = FLAGS_load,
+                                                     .rank = rank,
+                                                     .model = model,
+                                                     .optimizer = optimizer,
+                                                     .model_config = model_config,
+                                                     .state = state,
+                                                     .load_optimizer_state = true});
 
     start_step = resume_result.global_step;
     size_t consumed_batches = resume_result.consumed_batches;
@@ -327,31 +326,29 @@ void Train(const nn::parallel::Rank &rank) {
         for (size_t i = 0; i < num_skips; ++i) { ++train_iter; }
     }
 
-    auto save_checkpoint
-        = [&](const std::filesystem::path &save_dir, int64_t global_step, bool prune_step_checkpoints) {
-              SaveCheckpoint({
-                  .save_dir = save_dir,
-                  .global_step = global_step,
-                  .consumed_batches = consumed_batches,
-                  .last_lr = FLAGS_learning_rate,
-                  .n_layer = model_config.n_layer,
-                  .n_head = model_config.n_head,
-                  .n_kv_head = model_config.n_kv_head,
-                  .n_embd = model_config.n_embd,
-                  .vocab_size = model_config.vocab_size,
-                  .ddp_size = ddp_world_size,
-                  .tp_size = tp_world_size,
-                  .sp_size = sp_world_size,
-                  .pp_size = pp_world_size,
-                  .no_save_optim = FLAGS_no_save_optim,
-                  .prune_step_checkpoints = prune_step_checkpoints,
-                  .checkpoint_root_dir = FLAGS_save,
-                  .max_checkpoint_keep = FLAGS_max_checkpoint_keep,
-                  .rank = rank,
-                  .model = *model,
-                  .optimizer = *optimizer,
-              });
-          };
+    auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) {
+        SaveCheckpoint({
+            .save_dir = save_dir,
+            .global_step = global_step,
+            .consumed_batches = consumed_batches,
+            .last_lr = FLAGS_learning_rate,
+            .n_layer = model_config.n_layer,
+            .n_head = model_config.n_head,
+            .n_kv_head = model_config.n_kv_head,
+            .n_embd = model_config.n_embd,
+            .vocab_size = model_config.vocab_size,
+            .ddp_size = ddp_world_size,
+            .tp_size = tp_world_size,
+            .sp_size = sp_world_size,
+            .pp_size = pp_world_size,
+            .save_optimizer_state = FLAGS_save_optimizer_state,
+            .checkpoint_root_dir = FLAGS_save,
+            .max_checkpoint_keep = FLAGS_max_checkpoint_keep,
+            .rank = rank,
+            .model = *model,
+            .optimizer = *optimizer,
+        });
+    };
 
     for (int step = start_step; step < FLAGS_num_iteration + 1; ++step) {
         // Reset precision check counters at start of each iteration for file overwrite
@@ -475,7 +472,7 @@ void Train(const nn::parallel::Rank &rank) {
             if (rank.IsParallel()) {
                 step_dir /= std::format("rank_{:06d}", rank.GlobalRank());
             }
-            save_checkpoint(step_dir, step + 1, true);
+            save_checkpoint(step_dir, step + 1);
         }
     }
 
@@ -489,7 +486,7 @@ void Train(const nn::parallel::Rank &rank) {
     if (rank.IsParallel()) {
         final_dir /= std::format("rank_{:06d}", rank.GlobalRank());
     }
-    save_checkpoint(final_dir, FLAGS_num_iteration, false);
+    save_checkpoint(final_dir, FLAGS_num_iteration);
 
 #ifdef PROFILE_MODE
     Profiler::Instance().Report("llama3.report", Profiler::SortBy::DeviceTimePercentage);

infini_train/include/checkpoint.h …ni_train/include/checkpoint/checkpoint.hinfini_train/include/checkpoint.h renamed to infini_train/include/checkpoint/checkpoint.h infini_train/include/checkpoint.h renamed to infini_train/include/checkpoint/checkpoint.h

@@ -34,17 +34,16 @@ struct TrainerState {
 class Checkpoint {
 public:
     static void Save(const std::filesystem::path &checkpoint_dir, const nn::Module &model, const Optimizer *optimizer,
-                     const TrainerState &state, bool no_save_optim);
+                     const TrainerState &state, bool save_optimizer_state);
 
     static void Load(const std::filesystem::path &checkpoint_dir, nn::Module &model, Optimizer *optimizer,
                      TrainerState &state, bool load_optimizer_state);
 
 private:
-    static void SaveStateDictBinary(const std::filesystem::path &path,
-                                    const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
+    static void SaveStateDict(const std::filesystem::path &path,
+                              const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
 
-    static std::unordered_map<std::string, std::shared_ptr<Tensor>>
-    LoadStateDictBinary(const std::filesystem::path &path);
+    static std::unordered_map<std::string, std::shared_ptr<Tensor>> LoadStateDict(const std::filesystem::path &path);
 
     static void SaveTrainerState(const std::filesystem::path &path, const TrainerState &state);
     static TrainerState LoadTrainerState(const std::filesystem::path &path);

example/common/checkpoint_loader.h …/include/checkpoint/checkpoint_manager.hexample/common/checkpoint_loader.h renamed to infini_train/include/checkpoint/checkpoint_manager.h example/common/checkpoint_loader.h renamed to infini_train/include/checkpoint/checkpoint_manager.h

@@ -4,7 +4,7 @@
 #include <cstring>
 #include <filesystem>
 
-#include "infini_train/include/checkpoint.h"
+#include "infini_train/include/checkpoint/checkpoint.h"
 #include "infini_train/include/dataloader.h"
 #include "infini_train/include/nn/modules/module.h"
 #include "infini_train/include/nn/parallel/rank.h"
@@ -24,6 +24,7 @@ struct ResumeFromCheckpointArgs {
     std::shared_ptr<Optimizer> optimizer;
     const nn::TransformerConfig &model_config;
     TrainerState &state;
+    bool load_optimizer_state;
 };
 
 struct ResumeFromCheckpointResult {
@@ -45,8 +46,7 @@ struct SaveCheckpointArgs {
     int tp_size = 1;
     int sp_size = 1;
     int pp_size = 1;
-    bool no_save_optim = false;
-    bool prune_step_checkpoints = false;
+    bool save_optimizer_state = true;
     std::filesystem::path checkpoint_root_dir;
     size_t max_checkpoint_keep = 0;
     const nn::parallel::Rank &rank;

infini_train/include/nn/modules/module.h

@@ -47,9 +47,8 @@ class Module : public std::enable_shared_from_this<Module> {
 
     const std::string &type() const;
 
+    // TODO: Change return type to filterable iterator (like PyTorch's named_parameters with prefix matching)
     virtual std::vector<std::shared_ptr<Tensor>> Parameters() const;
-    std::vector<std::pair<std::string, std::shared_ptr<Tensor>>> NamedParameters(const std::string &prefix = "",
-                                                                                 bool remove_duplicate = true) const;
     bool has_parameter(const std::string &name) const;
     std::shared_ptr<Tensor> *mutable_parameter(const std::string &name);
     const std::shared_ptr<Tensor> &parameter(const std::string &name) const;

infini_train/include/optimizer.h

@@ -14,13 +14,10 @@ namespace infini_train {
 class Optimizer;
 
 using OptimizerCreator = std::function<std::shared_ptr<Optimizer>(const std::vector<std::shared_ptr<Tensor>> &params)>;
-using OptimizerCreatorNamed = std::function<std::shared_ptr<Optimizer>(
-    const std::vector<std::pair<std::string, std::shared_ptr<Tensor>>> &named_params)>;
 
 class Optimizer {
 public:
     explicit Optimizer(const std::vector<std::shared_ptr<Tensor>> &params);
-    Optimizer(const std::vector<std::pair<std::string, std::shared_ptr<Tensor>>> &named_params);
 
     virtual void ZeroGrad(bool set_to_none = true);
 
@@ -32,19 +29,16 @@ class Optimizer {
 
 protected:
     std::vector<std::shared_ptr<Tensor>> params_;
-    std::vector<std::string> param_names_;
 };
 
 namespace optimizers {
 class SGD : public Optimizer {
 public:
     SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate);
-    SGD(const std::vector<std::pair<std::string, std::shared_ptr<Tensor>>> &named_params, float learning_rate);
 
     void Step() override;
 
     static OptimizerCreator Create(float learning_rate);
-    static OptimizerCreatorNamed CreateNamed(float learning_rate);
 
 private:
     const float learning_rate_ = 0.0;
@@ -54,8 +48,6 @@ class Adam : public Optimizer {
 public:
     Adam(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate = 1e-3, float beta1 = 0.9,
          float beta2 = 0.999, float eps = 1e-8);
-    Adam(const std::vector<std::pair<std::string, std::shared_ptr<Tensor>>> &named_params, float learning_rate = 1e-3,
-         float beta1 = 0.9, float beta2 = 0.999, float eps = 1e-8);
 
     void Step() override;
 
@@ -64,8 +56,6 @@ class Adam : public Optimizer {
     void LoadStateDict(const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict) override;
     static OptimizerCreator Create(float learning_rate = 1e-3, float beta1 = 0.9, float beta2 = 0.999,
                                    float eps = 1e-8);
-    static OptimizerCreatorNamed CreateNamed(float learning_rate = 1e-3, float beta1 = 0.9, float beta2 = 0.999,
-                                             float eps = 1e-8);
 
 private:
     int64_t t_;