diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 60c0c908..fd2522e3 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -92,7 +92,6 @@ DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables savin DEFINE_string(load, "", "checkpoint directory to resume from"); DEFINE_string(save, "", "root directory used to store checkpoints"); DEFINE_uint32(max_checkpoint_keep, 3, "max number of checkpoint steps to keep"); -DEFINE_bool(save_optimizer_state, true, "whether optimizer state is persisted in checkpoints"); // precision check DEFINE_string( precision_check, "", @@ -366,10 +365,9 @@ void Train(const nn::parallel::Rank &rank) { const auto resume_result = ResumeFromCheckpoint({.resume_root = FLAGS_load, .rank = rank, .model = model, - .optimizer = optimizer, + .optimizer = nullptr, .model_config = model_config, .state = state, - .load_optimizer_state = false, .lr_scheduler = scheduler}); start_step = resume_result.global_step; size_t consumed_batches = resume_result.consumed_batches; @@ -398,12 +396,11 @@ void Train(const nn::parallel::Rank &rank) { .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, + .optimizer = nullptr, .lr_scheduler = scheduler.get(), }); }; diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 302e0808..05fad4a8 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -90,6 +90,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, "", "root directory used to store checkpoints"); DEFINE_uint32(max_checkpoint_keep, 3, "max number of checkpoint steps to keep"); +DEFINE_bool(load_optimizer_state, true, "whether optimizer state is restored from checkpoints"); DEFINE_bool(save_optimizer_state, true, "whether optimizer state is persisted in checkpoints"); // precision check @@ -345,10 +346,9 @@ void Train(const nn::parallel::Rank &rank) { const auto resume_result = ResumeFromCheckpoint({.resume_root = FLAGS_load, .rank = rank, .model = model, - .optimizer = optimizer, + .optimizer = FLAGS_load_optimizer_state ? optimizer : nullptr, .model_config = model_config, .state = state, - .load_optimizer_state = true, .lr_scheduler = scheduler}); start_step = resume_result.global_step; @@ -378,12 +378,11 @@ void Train(const nn::parallel::Rank &rank) { .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, + .optimizer = FLAGS_save_optimizer_state ? optimizer.get() : nullptr, .lr_scheduler = scheduler.get(), }); }; diff --git a/infini_train/include/checkpoint/checkpoint.h b/infini_train/include/checkpoint/checkpoint.h index 7d6f078f..0d3cee6b 100644 --- a/infini_train/include/checkpoint/checkpoint.h +++ b/infini_train/include/checkpoint/checkpoint.h @@ -32,10 +32,10 @@ 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 save_optimizer_state, const LRScheduler *lr_scheduler); + const TrainerState &state, const LRScheduler *lr_scheduler); static void Load(const std::filesystem::path &checkpoint_dir, nn::Module &model, Optimizer *optimizer, - TrainerState &state, bool load_optimizer_state, LRScheduler *lr_scheduler); + TrainerState &state, LRScheduler *lr_scheduler); private: static void SaveStateDict(const std::filesystem::path &path, diff --git a/infini_train/include/checkpoint/checkpoint_manager.h b/infini_train/include/checkpoint/checkpoint_manager.h index 8c1aa547..7e7b5eb7 100644 --- a/infini_train/include/checkpoint/checkpoint_manager.h +++ b/infini_train/include/checkpoint/checkpoint_manager.h @@ -29,7 +29,6 @@ struct ResumeFromCheckpointArgs { std::shared_ptr optimizer; const nn::TransformerConfig &model_config; TrainerState &state; - bool load_optimizer_state; std::shared_ptr lr_scheduler = nullptr; }; @@ -51,12 +50,11 @@ struct SaveCheckpointArgs { int tp_size = 1; int sp_size = 1; int pp_size = 1; - bool save_optimizer_state = true; std::filesystem::path checkpoint_root_dir; size_t max_checkpoint_keep = 0; const nn::parallel::Rank &rank; const nn::Module &model; - const Optimizer &optimizer; + const Optimizer *optimizer = nullptr; const LRScheduler *lr_scheduler = nullptr; }; diff --git a/infini_train/src/checkpoint/checkpoint.cc b/infini_train/src/checkpoint/checkpoint.cc index 2a026691..988b5120 100644 --- a/infini_train/src/checkpoint/checkpoint.cc +++ b/infini_train/src/checkpoint/checkpoint.cc @@ -182,7 +182,7 @@ template T ExtractNumberField(const std::string &content, const std } // namespace void Checkpoint::Save(const std::filesystem::path &checkpoint_dir, const nn::Module &model, const Optimizer *optimizer, - const TrainerState &state, bool save_optimizer_state, const LRScheduler *lr_scheduler) { + const TrainerState &state, const LRScheduler *lr_scheduler) { std::filesystem::create_directories(checkpoint_dir); LOG(INFO) << "[CKPT] Save begin: dir=" << checkpoint_dir << ", global_step=" << state.global_step; @@ -190,8 +190,7 @@ void Checkpoint::Save(const std::filesystem::path &checkpoint_dir, const nn::Mod SaveStateDict(model_path, model.StateDict()); - if (save_optimizer_state) { - CHECK(optimizer != nullptr) << "Optimizer pointer is null, cannot save optimizer state."; + if (optimizer != nullptr) { auto opt_state = optimizer->StateDict(); if (!opt_state.empty()) { const auto opt_path = checkpoint_dir / "optimizer.ckpt"; @@ -208,14 +207,13 @@ void Checkpoint::Save(const std::filesystem::path &checkpoint_dir, const nn::Mod } void Checkpoint::Load(const std::filesystem::path &checkpoint_dir, nn::Module &model, Optimizer *optimizer, - TrainerState &state, bool load_optimizer_state, LRScheduler *lr_scheduler) { + TrainerState &state, LRScheduler *lr_scheduler) { const auto model_path = checkpoint_dir / "model.ckpt"; LOG(INFO) << "[CKPT] Loading model: " << model_path; model.LoadStateDict(LoadStateDict(model_path)); - if (load_optimizer_state) { - CHECK(optimizer != nullptr) << "Optimizer pointer is null, cannot load optimizer state."; + if (optimizer != nullptr) { const auto opt_path = checkpoint_dir / "optimizer.ckpt"; if (std::filesystem::exists(opt_path)) { LOG(INFO) << "[CKPT] Loading optimizer: " << opt_path; diff --git a/infini_train/src/checkpoint/checkpoint_manager.cc b/infini_train/src/checkpoint/checkpoint_manager.cc index cb40c7ad..51c20f12 100644 --- a/infini_train/src/checkpoint/checkpoint_manager.cc +++ b/infini_train/src/checkpoint/checkpoint_manager.cc @@ -39,8 +39,7 @@ ResumeFromCheckpointResult ResumeFromCheckpoint(const ResumeFromCheckpointArgs & } } - Checkpoint::Load(resume_dir, *args.model, args.optimizer.get(), args.state, args.load_optimizer_state, - args.lr_scheduler.get()); + Checkpoint::Load(resume_dir, *args.model, args.optimizer.get(), args.state, args.lr_scheduler.get()); result.global_step = static_cast(args.state.global_step); @@ -89,7 +88,7 @@ void SaveCheckpoint(const SaveCheckpointArgs &args) { state.sp_size = args.sp_size; state.pp_size = args.pp_size; - Checkpoint::Save(args.save_dir, args.model, &args.optimizer, state, args.save_optimizer_state, args.lr_scheduler); + Checkpoint::Save(args.save_dir, args.model, args.optimizer, state, args.lr_scheduler); const auto ckpt_end = std::chrono::high_resolution_clock::now(); const double ckpt_ms = std::chrono::duration(ckpt_end - ckpt_start).count(); diff --git a/scripts/test_config.json b/scripts/test_config.json index 992cf5d2..9bd7780e 100644 --- a/scripts/test_config.json +++ b/scripts/test_config.json @@ -839,7 +839,6 @@ "pipeline_parallel": 2, "virtual_pipeline_parallel": 2, "save": "@CKPT_ROOT_DIR@/3d_ddp2_tp2_pp2_no_resume", - "save_optimizer_state": true, "max_checkpoint_keep": 5 } }, @@ -857,7 +856,6 @@ "virtual_pipeline_parallel": 2, "load": "@CKPT_ROOT_DIR@/3d_ddp2_tp2_pp2_no_resume/checkpoint_step_000030", "save": "@CKPT_ROOT_DIR@/3d_tp2_pp2_resume", - "save_optimizer_state": true, "max_checkpoint_keep": 5 } } diff --git a/tests/checkpoint/test_checkpoint_serialization.cc b/tests/checkpoint/test_checkpoint_serialization.cc index e21952f4..495dcf27 100644 --- a/tests/checkpoint/test_checkpoint_serialization.cc +++ b/tests/checkpoint/test_checkpoint_serialization.cc @@ -30,7 +30,7 @@ TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) { auto opt1 = std::make_shared(model1->Parameters(), 0.01); TrainerState saved{.global_step = 42, .consumed_batches = 100}; - Checkpoint::Save(dir, *model1, opt1.get(), saved, /*save_optimizer_state=*/true, nullptr); + Checkpoint::Save(dir, *model1, opt1.get(), saved, nullptr); auto model2 = std::make_shared(3, 2, true, GetDevice()); auto q1 = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); @@ -42,7 +42,7 @@ TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) { auto opt2 = std::make_shared(model2->Parameters(), 0.01); TrainerState loaded; - Checkpoint::Load(dir, *model2, opt2.get(), loaded, /*load_optimizer_state=*/true, nullptr); + Checkpoint::Load(dir, *model2, opt2.get(), loaded, nullptr); EXPECT_EQ(loaded.global_step, 42); EXPECT_EQ(loaded.consumed_batches, 100); diff --git a/tests/checkpoint/test_lr_scheduler_state.cc b/tests/checkpoint/test_lr_scheduler_state.cc index f12b60b2..fc49fb3d 100644 --- a/tests/checkpoint/test_lr_scheduler_state.cc +++ b/tests/checkpoint/test_lr_scheduler_state.cc @@ -60,7 +60,7 @@ TEST_P(LRSchedulerCheckpointTest, SaveAndLoadLRSchedulerState) { StepTimes(sched1, 3); TrainerState saved{.global_step = 3, .consumed_batches = 12}; - Checkpoint::Save(dir, *model1, opt1.get(), saved, /*save_optimizer_state=*/false, sched1.get()); + Checkpoint::Save(dir, *model1, nullptr, saved, sched1.get()); EXPECT_TRUE(std::filesystem::exists(dir / "lr_scheduler.ckpt")); auto model2 = MakeModel(GetDevice()); @@ -68,7 +68,7 @@ TEST_P(LRSchedulerCheckpointTest, SaveAndLoadLRSchedulerState) { auto sched2 = CreateLRScheduler(opt2, MakeSchedulerConfig()); TrainerState loaded; - Checkpoint::Load(dir, *model2, opt2.get(), loaded, /*load_optimizer_state=*/false, sched2.get()); + Checkpoint::Load(dir, *model2, nullptr, loaded, sched2.get()); EXPECT_EQ(loaded.global_step, 3); EXPECT_EQ(loaded.consumed_batches, 12); @@ -91,7 +91,7 @@ TEST_P(LRSchedulerCheckpointTest, SkipsLRSchedulerStateWhenSchedulerIsNull) { auto opt1 = std::make_shared(model1->Parameters(), kBaseLR); TrainerState saved{.global_step = 3}; - Checkpoint::Save(dir, *model1, opt1.get(), saved, /*save_optimizer_state=*/false, nullptr); + Checkpoint::Save(dir, *model1, nullptr, saved, nullptr); EXPECT_FALSE(std::filesystem::exists(dir / "lr_scheduler.ckpt")); std::filesystem::remove_all(dir); diff --git a/tests/checkpoint/test_trainer_state.cc b/tests/checkpoint/test_trainer_state.cc index 1a390788..b5352556 100644 --- a/tests/checkpoint/test_trainer_state.cc +++ b/tests/checkpoint/test_trainer_state.cc @@ -44,7 +44,7 @@ TEST_P(TrainerStateTest, TrainerStateFileCreated) { *model->mutable_parameter("weight") = p; auto opt = std::make_shared(model->Parameters(), 0.01); - Checkpoint::Save(dir, *model, opt.get(), saved, /*save_optimizer_state=*/true, nullptr); + Checkpoint::Save(dir, *model, opt.get(), saved, nullptr); EXPECT_TRUE(std::filesystem::exists(dir / "trainer_state.json")); @@ -80,7 +80,7 @@ TEST_P(TrainerStateTest, RoundTrip) { *model1->mutable_parameter("weight") = p1; auto opt1 = std::make_shared(model1->Parameters(), 0.01); - Checkpoint::Save(dir, *model1, opt1.get(), saved, /*save_optimizer_state=*/false, nullptr); + Checkpoint::Save(dir, *model1, nullptr, saved, nullptr); auto model2 = std::make_shared(1, 3, true, GetDevice()); auto p2 = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice()); @@ -89,7 +89,7 @@ TEST_P(TrainerStateTest, RoundTrip) { auto opt2 = std::make_shared(model2->Parameters(), 0.01); TrainerState loaded; - Checkpoint::Load(dir, *model2, opt2.get(), loaded, /*load_optimizer_state=*/false, nullptr); + Checkpoint::Load(dir, *model2, nullptr, loaded, nullptr); EXPECT_EQ(loaded.global_step, 99); EXPECT_EQ(loaded.consumed_batches, 5000);