feat: track checkpoint progress in micro-batches

example/gpt2/main.cc

@@ -305,9 +305,9 @@ void Train(const nn::parallel::Rank &rank) {
         model = std::make_shared<DistributedDataParallel>(model, rank, ddp_config);
     }
 
+    const size_t train_loader_batch_size = pp_world_size > 1 ? FLAGS_batch_size * num_micro_batches : FLAGS_batch_size;
     DistributedDataLoader train_loader(std::make_shared<TinyShakespeareDataset>(FLAGS_input_bin, FLAGS_sequence_length),
-                                       pp_world_size > 1 ? FLAGS_batch_size * num_micro_batches : FLAGS_batch_size,
-                                       ddp_rank, ddp_world_size);
+                                       train_loader_batch_size, ddp_rank, ddp_world_size);
 
     std::optional<DistributedDataLoader> val_loader = std::nullopt;
     if (!FLAGS_input_val_bin.empty()) {
@@ -382,23 +382,21 @@ void Train(const nn::parallel::Rank &rank) {
                                                      .state = state,
                                                      .lr_scheduler = scheduler});
     start_step = resume_result.global_step;
-    size_t consumed_batches = resume_result.consumed_batches;
+    size_t consumed_train_samples = resume_result.consumed_train_samples;
 
     // TODO(jym): Replace with Sampler abstraction when available.
     // Skip dataloader to resume from the correct batch position.
-    if (consumed_batches > 0) {
-        size_t start = train_iter.BatchIndex();
-        // Each rank processes every ddp_world_size-th batch starting from its own rank.
-        // num_skips calculates how many ++ iterations to reach the saved batch position.
-        size_t num_skips = (consumed_batches - start) / ddp_world_size;
+    if (consumed_train_samples > 0) {
+        const size_t num_skips
+            = DataLoaderBatchesToSkip(consumed_train_samples, train_loader_batch_size, ddp_world_size);
         for (size_t i = 0; i < num_skips; ++i) { ++train_iter; }
     }
 
     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,
+            .consumed_train_samples = consumed_train_samples,
             .n_layer = model_config.n_layer,
             .n_head = model_config.n_head,
             .n_kv_head = model_config.n_kv_head,
@@ -470,7 +468,7 @@ void Train(const nn::parallel::Rank &rank) {
                 // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
                 // TODO(dcj): support dataloader.reset() later
                 ++train_iter;
-                consumed_batches = train_iter.BatchIndex();
+                consumed_train_samples += static_cast<size_t>(FLAGS_batch_size) * ddp_world_size;
                 x = std::make_shared<Tensor>(x->To(device));
                 y = std::make_shared<Tensor>(y->To(device));
 
@@ -504,7 +502,7 @@ void Train(const nn::parallel::Rank &rank) {
             // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
             // TODO(dcj): support dataloader.reset() later
             ++train_iter;
-            consumed_batches = train_iter.BatchIndex();
+            consumed_train_samples += train_loader_batch_size * ddp_world_size;
             x = std::make_shared<Tensor>(x->To(device));
             y = std::make_shared<Tensor>(y->To(device));
 

example/llama3/main.cc

@@ -279,9 +279,9 @@ void Train(const nn::parallel::Rank &rank) {
         model = std::make_shared<DistributedDataParallel>(model, rank, ddp_config);
     }
 
+    const size_t train_loader_batch_size = pp_world_size > 1 ? FLAGS_batch_size * num_micro_batches : FLAGS_batch_size;
     DistributedDataLoader train_loader(std::make_shared<TinyShakespeareDataset>(FLAGS_input_bin, FLAGS_sequence_length),
-                                       pp_world_size > 1 ? FLAGS_batch_size * num_micro_batches : FLAGS_batch_size,
-                                       ddp_rank, ddp_world_size);
+                                       train_loader_batch_size, ddp_rank, ddp_world_size);
 
     std::optional<DistributedDataLoader> val_loader = std::nullopt;
     if (!FLAGS_input_val_bin.empty()) {
@@ -364,23 +364,21 @@ void Train(const nn::parallel::Rank &rank) {
                                                      .lr_scheduler = scheduler});
 
     start_step = resume_result.global_step;
-    size_t consumed_batches = resume_result.consumed_batches;
+    size_t consumed_train_samples = resume_result.consumed_train_samples;
 
     // TODO(jym): Replace with Sampler abstraction when available.
     // Skip dataloader to resume from the correct batch position.
-    if (consumed_batches > 0) {
-        size_t start = train_iter.BatchIndex();
-        // Each rank processes every ddp_world_size-th batch starting from its own rank.
-        // num_skips calculates how many ++ iterations to reach the saved batch position.
-        size_t num_skips = (consumed_batches - start) / ddp_world_size;
+    if (consumed_train_samples > 0) {
+        const size_t num_skips
+            = DataLoaderBatchesToSkip(consumed_train_samples, train_loader_batch_size, ddp_world_size);
         for (size_t i = 0; i < num_skips; ++i) { ++train_iter; }
     }
 
     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,
+            .consumed_train_samples = consumed_train_samples,
             .n_layer = model_config.n_layer,
             .n_head = model_config.n_head,
             .n_kv_head = model_config.n_kv_head,
@@ -450,7 +448,7 @@ void Train(const nn::parallel::Rank &rank) {
                 // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
                 // TODO(dcj): support dataloader.reset() later
                 ++train_iter;
-                consumed_batches = train_iter.BatchIndex();
+                consumed_train_samples += static_cast<size_t>(FLAGS_batch_size) * ddp_world_size;
                 x = std::make_shared<Tensor>(x->To(device));
                 y = std::make_shared<Tensor>(y->To(device));
 
@@ -483,7 +481,7 @@ void Train(const nn::parallel::Rank &rank) {
             // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
             // TODO(dcj): support dataloader.reset() later
             ++train_iter;
-            consumed_batches = train_iter.BatchIndex();
+            consumed_train_samples += train_loader_batch_size * ddp_world_size;
             x = std::make_shared<Tensor>(x->To(device));
             y = std::make_shared<Tensor>(y->To(device));
 

infini_train/include/checkpoint/checkpoint.h

@@ -17,7 +17,7 @@ class Module;
 
 struct TrainerState {
     int64_t global_step = 0;
-    int64_t consumed_batches = 0;
+    int64_t consumed_train_samples = 0;
     int64_t n_layer = 0;
     int64_t n_head = 0;
     int64_t n_kv_head = 0;

infini_train/include/checkpoint/checkpoint_manager.h

@@ -34,13 +34,13 @@ struct ResumeFromCheckpointArgs {
 
 struct ResumeFromCheckpointResult {
     int global_step = 0;
-    size_t consumed_batches = 0;
+    size_t consumed_train_samples = 0;
 };
 
 struct SaveCheckpointArgs {
     std::filesystem::path save_dir;
     int64_t global_step = 0;
-    size_t consumed_batches = 0;
+    size_t consumed_train_samples = 0;
     int64_t n_layer = 0;
     int64_t n_head = 0;
     int64_t n_kv_head = 0;
@@ -61,3 +61,5 @@ struct SaveCheckpointArgs {
 ResumeFromCheckpointResult ResumeFromCheckpoint(const ResumeFromCheckpointArgs &args);
 
 void SaveCheckpoint(const SaveCheckpointArgs &args);
+
+size_t DataLoaderBatchesToSkip(size_t consumed_train_samples, size_t local_batch_size, size_t ddp_world_size);

infini_train/src/checkpoint/checkpoint.cc

@@ -237,8 +237,8 @@ void Checkpoint::Load(const std::filesystem::path &checkpoint_dir, nn::Module &m
     }
 
     LOG(ERROR) << "[CKPT] Load done: global_step=" << state.global_step
-               << ", consumed_batches =" << state.consumed_batches << ", topology(ddp,tp,sp,pp)=(" << state.ddp_size
-               << "," << state.tp_size << "," << state.sp_size << "," << state.pp_size << ")";
+               << ", consumed_train_samples=" << state.consumed_train_samples << ", topology(ddp,tp,sp,pp)=("
+               << state.ddp_size << "," << state.tp_size << "," << state.sp_size << "," << state.pp_size << ")";
 }
 
 void Checkpoint::SaveStateDict(const std::filesystem::path &path,
@@ -320,7 +320,7 @@ void Checkpoint::SaveTrainerState(const std::filesystem::path &path, const Train
     ofs << "  \"n_embd\": " << state.n_embd << ",\n";
     ofs << "  \"vocab_size\": " << state.vocab_size << ",\n";
     ofs << "  \"global_step\": " << state.global_step << ",\n";
-    ofs << "  \"consumed_batches\": " << state.consumed_batches << ",\n";
+    ofs << "  \"consumed_train_samples\": " << state.consumed_train_samples << ",\n";
     ofs << "  \"ddp_size\": " << state.ddp_size << ",\n";
     ofs << "  \"tp_size\": " << state.tp_size << ",\n";
     ofs << "  \"sp_size\": " << state.sp_size << ",\n";
@@ -341,7 +341,7 @@ TrainerState Checkpoint::LoadTrainerState(const std::filesystem::path &path) {
     state.n_embd = ExtractNumberField<int64_t>(content, "n_embd", 0);
     state.vocab_size = ExtractNumberField<int64_t>(content, "vocab_size", 0);
     state.global_step = ExtractNumberField<int64_t>(content, "global_step", 0);
-    state.consumed_batches = ExtractNumberField<int64_t>(content, "consumed_batches", 0);
+    state.consumed_train_samples = ExtractNumberField<int64_t>(content, "consumed_train_samples", 0);
     state.ddp_size = ExtractNumberField<int>(content, "ddp_size", 1);
     state.tp_size = ExtractNumberField<int>(content, "tp_size", 1);
     state.sp_size = ExtractNumberField<int>(content, "sp_size", 1);

infini_train/src/checkpoint/checkpoint_manager.cc

@@ -4,6 +4,7 @@
 #include <cstdlib>
 #include <filesystem>
 #include <format>
+#include <limits>
 #include <memory>
 #include <string>
 #include <vector>
@@ -63,10 +64,10 @@ ResumeFromCheckpointResult ResumeFromCheckpoint(const ResumeFromCheckpointArgs &
     CHECK_EQ(args.state.pp_size, pp_world_size)
         << "PP size mismatch: checkpoint has PP=" << args.state.pp_size << ", but current run has PP=" << pp_world_size;
 
-    result.consumed_batches = static_cast<size_t>(std::max<int64_t>(args.state.consumed_batches, 0));
+    result.consumed_train_samples = static_cast<size_t>(std::max<int64_t>(args.state.consumed_train_samples, 0));
     if (args.rank.IsMainRank()) {
-        LOG(INFO) << std::format("Resume training from step {}, consumed_batches  {}", args.state.global_step,
-                                 args.state.consumed_batches);
+        LOG(INFO) << std::format("Resume training from step {}, consumed_train_samples {}", args.state.global_step,
+                                 args.state.consumed_train_samples);
     }
 
     return result;
@@ -77,7 +78,7 @@ void SaveCheckpoint(const SaveCheckpointArgs &args) {
 
     TrainerState state;
     state.global_step = args.global_step;
-    state.consumed_batches = static_cast<int64_t>(args.consumed_batches);
+    state.consumed_train_samples = static_cast<int64_t>(args.consumed_train_samples);
     state.n_layer = args.n_layer;
     state.n_head = args.n_head;
     state.n_kv_head = args.n_kv_head;
@@ -118,3 +119,16 @@ void SaveCheckpoint(const SaveCheckpointArgs &args) {
         }
     }
 }
+
+size_t DataLoaderBatchesToSkip(size_t consumed_train_samples, size_t local_batch_size, size_t ddp_world_size) {
+    CHECK_GT(local_batch_size, 0);
+    CHECK_GT(ddp_world_size, 0);
+    CHECK_LE(local_batch_size, std::numeric_limits<size_t>::max() / ddp_world_size)
+        << "Data loader batch size overflows size_t";
+    const size_t global_loader_batch_size = local_batch_size * ddp_world_size;
+    CHECK_EQ(consumed_train_samples % global_loader_batch_size, 0)
+        << "consumed_train_samples=" << consumed_train_samples
+        << " does not align with current local_batch_size=" << local_batch_size
+        << " and ddp_world_size=" << ddp_world_size;
+    return consumed_train_samples / global_loader_batch_size;
+}

tests/checkpoint/test_checkpoint_serialization.cc

@@ -29,7 +29,7 @@ TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) {
     *model1->mutable_parameter("bias") = p2;
 
     auto opt1 = std::make_shared<optimizers::Adam>(model1->Parameters(), 0.01);
-    TrainerState saved{.global_step = 42, .consumed_batches = 100};
+    TrainerState saved{.global_step = 42, .consumed_train_samples = 100};
     Checkpoint::Save(dir, *model1, opt1.get(), saved, nullptr);
 
     auto model2 = std::make_shared<nn::Linear>(3, 2, true, GetDevice());
@@ -45,7 +45,7 @@ TEST_P(CheckpointSerializationTest, SaveAndLoadModelFP32) {
     Checkpoint::Load(dir, *model2, opt2.get(), loaded, nullptr);
 
     EXPECT_EQ(loaded.global_step, 42);
-    EXPECT_EQ(loaded.consumed_batches, 100);
+    EXPECT_EQ(loaded.consumed_train_samples, 100);
 
     auto w1_cpu = model2->parameter("weight")->To(Device());
     const float *data = static_cast<const float *>(w1_cpu.DataPtr());

tests/checkpoint/test_lr_scheduler_state.cc

@@ -59,7 +59,7 @@ TEST_P(LRSchedulerCheckpointTest, SaveAndLoadLRSchedulerState) {
     auto sched1 = CreateLRScheduler(opt1, MakeSchedulerConfig());
     StepTimes(sched1, 3);
 
-    TrainerState saved{.global_step = 3, .consumed_batches = 12};
+    TrainerState saved{.global_step = 3, .consumed_train_samples = 12};
     Checkpoint::Save(dir, *model1, nullptr, saved, sched1.get());
     EXPECT_TRUE(std::filesystem::exists(dir / "lr_scheduler.ckpt"));
 
@@ -71,7 +71,7 @@ TEST_P(LRSchedulerCheckpointTest, SaveAndLoadLRSchedulerState) {
     Checkpoint::Load(dir, *model2, nullptr, loaded, sched2.get());
 
     EXPECT_EQ(loaded.global_step, 3);
-    EXPECT_EQ(loaded.consumed_batches, 12);
+    EXPECT_EQ(loaded.consumed_train_samples, 12);
     EXPECT_EQ(sched2->last_step(), sched1->last_step());
     EXPECT_NEAR(sched2->learning_rate(), sched1->learning_rate(), kEps);
 

tests/checkpoint/test_trainer_state.cc

@@ -5,6 +5,7 @@
 #include "gtest/gtest.h"
 
 #include "infini_train/include/checkpoint/checkpoint.h"
+#include "infini_train/include/checkpoint/checkpoint_manager.h"
 #include "infini_train/include/nn/modules/linear.h"
 #include "infini_train/include/nn/modules/module.h"
 #include "infini_train/include/optimizer.h"
@@ -20,7 +21,7 @@ class TrainerStateTest : public test::InfiniTrainTest {};
 TEST_P(TrainerStateTest, DefaultValues) {
     TrainerState state;
     EXPECT_EQ(state.global_step, 0);
-    EXPECT_EQ(state.consumed_batches, 0);
+    EXPECT_EQ(state.consumed_train_samples, 0);
     EXPECT_EQ(state.n_layer, 0);
     EXPECT_EQ(state.n_head, 0);
     EXPECT_EQ(state.n_kv_head, 0);
@@ -36,7 +37,7 @@ TEST_P(TrainerStateTest, TrainerStateFileCreated) {
     auto dir = std::filesystem::temp_directory_path() / "test_trainer_json";
     std::filesystem::remove_all(dir);
 
-    TrainerState saved{.global_step = 30, .consumed_batches = 1200};
+    TrainerState saved{.global_step = 30, .consumed_train_samples = 1200};
 
     auto model = std::make_shared<nn::Linear>(1, 2, true, GetDevice());
     auto p = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kFLOAT32, GetDevice());
@@ -51,7 +52,7 @@ TEST_P(TrainerStateTest, TrainerStateFileCreated) {
     std::ifstream ifs(dir / "trainer_state.json");
     std::string content((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
     EXPECT_NE(content.find("\"global_step\""), std::string::npos);
-    EXPECT_NE(content.find("\"consumed_batches\""), std::string::npos);
+    EXPECT_NE(content.find("\"consumed_train_samples\""), std::string::npos);
 
     std::filesystem::remove_all(dir);
 }
@@ -62,7 +63,7 @@ TEST_P(TrainerStateTest, RoundTrip) {
 
     TrainerState saved{
         .global_step = 99,
-        .consumed_batches = 5000,
+        .consumed_train_samples = 5000,
         .n_layer = 24,
         .n_head = 16,
         .n_kv_head = 8,
@@ -92,7 +93,7 @@ TEST_P(TrainerStateTest, RoundTrip) {
     Checkpoint::Load(dir, *model2, nullptr, loaded, nullptr);
 
     EXPECT_EQ(loaded.global_step, 99);
-    EXPECT_EQ(loaded.consumed_batches, 5000);
+    EXPECT_EQ(loaded.consumed_train_samples, 5000);
     EXPECT_EQ(loaded.n_layer, 24);
     EXPECT_EQ(loaded.n_head, 16);
     EXPECT_EQ(loaded.n_kv_head, 8);
@@ -104,4 +105,19 @@ TEST_P(TrainerStateTest, RoundTrip) {
     std::filesystem::remove_all(dir);
 }
 
+TEST_P(TrainerStateTest, DataLoaderSkipUsesCurrentBatchConfiguration) {
+    EXPECT_EQ(DataLoaderBatchesToSkip(/*consumed_train_samples=*/400, /*local_batch_size=*/4,
+                                      /*ddp_world_size=*/2),
+              50);
+    EXPECT_EQ(DataLoaderBatchesToSkip(/*consumed_train_samples=*/400, /*local_batch_size=*/8,
+                                      /*ddp_world_size=*/2),
+              25);
+}
+
+TEST_P(TrainerStateTest, DataLoaderSkipRejectsUnalignedConfiguration) {
+    EXPECT_DEATH(DataLoaderBatchesToSkip(/*consumed_train_samples=*/400, /*local_batch_size=*/6,
+                                         /*ddp_world_size=*/4),
+                 "does not align");
+}
+
 INFINI_TRAIN_REGISTER_TEST(TrainerStateTest);