[WIP] Feat/checkpoint resharding

example/gpt2/main.cc

@@ -339,6 +339,18 @@ void Train(const nn::parallel::Rank &rank) {
                                                                          model_chunks, ddp_world_size, ddp_rank);
     } else {
         optimizer = optimizer_creator(params_to_optimize);
+        std::unordered_map<const Tensor *, std::string> parameter_name_by_tensor;
+        for (const auto &[name, parameter] : model->NamedParameters()) {
+            parameter_name_by_tensor.emplace(parameter.get(), name);
+        }
+        std::vector<std::string> parameter_names;
+        parameter_names.reserve(params_to_optimize.size());
+        for (const auto &parameter : params_to_optimize) {
+            auto it = parameter_name_by_tensor.find(parameter.get());
+            CHECK(it != parameter_name_by_tensor.end()) << "Optimizer parameter is not registered in the model";
+            parameter_names.push_back(it->second);
+        }
+        optimizer->SetParameterNames(parameter_names);
     }
 
     const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
@@ -369,18 +381,22 @@ void Train(const nn::parallel::Rank &rank) {
                                                      .optimizer = optimizer,
                                                      .model_config = model_config,
                                                      .state = state,
-                                                     .load_optimizer_state = false,
+                                                     .load_optimizer_state = true,
                                                      .lr_scheduler = scheduler});
     start_step = resume_result.global_step;
     size_t consumed_batches = resume_result.consumed_batches;
 
     // 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;
+        const size_t start = train_iter.BatchIndex();
+        CHECK(pp_world_size == 1 || consumed_batches % num_micro_batches == 0);
+        const size_t consumed_loader_batches
+            = pp_world_size > 1 ? consumed_batches / num_micro_batches : consumed_batches;
+        const size_t target = consumed_loader_batches + static_cast<size_t>(ddp_rank);
+        CHECK_GE(target, start);
+        CHECK_EQ((target - start) % ddp_world_size, 0);
+        const size_t num_skips = (target - start) / ddp_world_size;
         for (size_t i = 0; i < num_skips; ++i) { ++train_iter; }
     }
 
@@ -495,7 +511,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_batches = train_iter.BatchIndex() * num_micro_batches;
             x = std::make_shared<Tensor>(x->To(device));
             y = std::make_shared<Tensor>(y->To(device));
 

example/llama3/main.cc

@@ -319,6 +319,18 @@ void Train(const nn::parallel::Rank &rank) {
                                                                          model_chunks, ddp_world_size, ddp_rank);
     } else {
         optimizer = optimizer_creator(params_to_optimize);
+        std::unordered_map<const Tensor *, std::string> parameter_name_by_tensor;
+        for (const auto &[name, parameter] : model->NamedParameters()) {
+            parameter_name_by_tensor.emplace(parameter.get(), name);
+        }
+        std::vector<std::string> parameter_names;
+        parameter_names.reserve(params_to_optimize.size());
+        for (const auto &parameter : params_to_optimize) {
+            auto it = parameter_name_by_tensor.find(parameter.get());
+            CHECK(it != parameter_name_by_tensor.end()) << "Optimizer parameter is not registered in the model";
+            parameter_names.push_back(it->second);
+        }
+        optimizer->SetParameterNames(parameter_names);
     }
 
     const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
@@ -357,10 +369,14 @@ void Train(const nn::parallel::Rank &rank) {
     // 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;
+        const size_t start = train_iter.BatchIndex();
+        CHECK(pp_world_size == 1 || consumed_batches % num_micro_batches == 0);
+        const size_t consumed_loader_batches
+            = pp_world_size > 1 ? consumed_batches / num_micro_batches : consumed_batches;
+        const size_t target = consumed_loader_batches + static_cast<size_t>(ddp_rank);
+        CHECK_GE(target, start);
+        CHECK_EQ((target - start) % ddp_world_size, 0);
+        const size_t num_skips = (target - start) / ddp_world_size;
         for (size_t i = 0; i < num_skips; ++i) { ++train_iter; }
     }
 
@@ -472,7 +488,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_batches = train_iter.BatchIndex() * num_micro_batches;
             x = std::make_shared<Tensor>(x->To(device));
             y = std::make_shared<Tensor>(y->To(device));
 

infini_train/include/checkpoint/checkpoint.h

@@ -6,6 +6,11 @@
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <vector>
+
+#include "infini_train/include/checkpoint/save_planner.h"
+#include "infini_train/include/checkpoint/shard_spec.h"
+#include "infini_train/include/lr_scheduler.h"
 
 namespace infini_train {
 class Optimizer;
@@ -37,9 +42,70 @@ class Checkpoint {
     static void Load(const std::filesystem::path &checkpoint_dir, nn::Module &model, Optimizer *optimizer,
                      TrainerState &state, bool load_optimizer_state, LRScheduler *lr_scheduler);
 
+    static void SaveSharded(const std::filesystem::path &checkpoint_dir, const checkpoint::ShardedStateDict &sharded_sd,
+                            const std::vector<checkpoint::WriteItem> &write_items,
+                            const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict,
+                            const std::unordered_map<std::string, std::shared_ptr<Tensor>> &optimizer_state,
+                            const TrainerState &state, int global_rank);
+
+    static void SaveStateDictFile(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>>
+    LoadStateDictFile(const std::filesystem::path &path);
+
+    struct CheckpointMetadata {
+        int version = 0;
+        int64_t iteration = 0;
+
+        struct ParallelConfig {
+            int tp_size = 1;
+            int pp_size = 1;
+            int dp_size = 1;
+            int sp_size = 1;
+        } parallel_config;
+
+        struct TensorEntry {
+            std::string key;
+            std::string dtype_str;
+            std::vector<int64_t> global_shape;
+            std::vector<int64_t> local_shape;
+            std::vector<int64_t> global_offset;
+            std::vector<int> axis_fragmentations;
+            std::vector<checkpoint::ShardSegment> segments;
+            std::string file;
+            uint64_t offset = 0;
+            uint64_t byte_size = 0;
+            int replica_id = 0;
+            std::vector<int> stored_on_ranks;
+            int pp_rank = 0;
+        };
+
+        std::vector<TensorEntry> tensors;
+        bool has_metadata = false;
+    };
+
+    static CheckpointMetadata LoadMetadata(const std::filesystem::path &checkpoint_dir);
+    static void SaveMetadataFile(const std::filesystem::path &path, const CheckpointMetadata &metadata);
+
+    // Public LR-scheduler serialization helpers used by checkpoint_manager.
+    static void SaveLRSchedulerStateFile(const std::filesystem::path &path, const LRSchedulerStateDict &state_dict);
+    static LRSchedulerStateDict LoadLRSchedulerStateFile(const std::filesystem::path &path);
+
+    // Public trainer-state serialization helpers used by checkpoint_manager.
+    static void SaveTrainerStateFile(const std::filesystem::path &path, const TrainerState &state);
+    static TrainerState LoadTrainerStateFile(const std::filesystem::path &path);
+
 private:
-    static void SaveStateDict(const std::filesystem::path &path,
-                              const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
+    struct SavedTensorLocation {
+        uint64_t data_offset = 0;
+        uint64_t byte_size = 0;
+    };
+    using SavedTensorLocations = std::unordered_map<std::string, SavedTensorLocation>;
+
+    static SavedTensorLocations
+    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>> LoadStateDict(const std::filesystem::path &path);
 

infini_train/include/checkpoint/checkpoint_manager.h

@@ -6,7 +6,6 @@
 #include <memory>
 
 #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"
 #include "infini_train/include/optimizer.h"

infini_train/include/checkpoint/load_planner.h

@@ -0,0 +1,52 @@
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <string>
+#include <vector>
+
+#include "infini_train/include/checkpoint/checkpoint.h"
+#include "infini_train/include/checkpoint/shard_spec.h"
+#include "infini_train/include/datatype.h"
+
+namespace infini_train::checkpoint {
+
+// One storage-region transfer from a saved shard into a target local tensor.
+struct ReadItem {
+    std::string key;
+    std::string filename;
+    DataType dtype = DataType::kFLOAT32;
+    std::vector<int64_t> global_shape;
+    uint64_t byte_size = 0;
+    uint64_t data_offset = 0;
+    int shard_dim = -1;
+    int64_t source_offset = 0;
+    int64_t target_offset = 0;
+    int64_t length = 0;
+    std::vector<int64_t> source_shape;
+};
+
+// All reads required to materialize one target local tensor.
+struct TargetTensorPlan {
+    std::string key;
+    DataType dtype = DataType::kFLOAT32;
+    std::vector<int64_t> global_shape;
+    std::vector<int64_t> target_shape;
+    int shard_dim = -1;
+    int64_t trailing_zero_fill = 0;
+    std::vector<ReadItem> reads;
+};
+
+// Complete load plan for one rank.
+struct LoadPlan {
+    std::map<std::string, TargetTensorPlan> tensors;
+};
+
+class LoadPlanner {
+public:
+    // Compute saved-to-target overlaps from explicit global shard coordinates.
+    static LoadPlan PlanReshard(const Checkpoint::CheckpointMetadata &metadata,
+                                const ShardedStateDict &target_state_dict);
+};
+
+} // namespace infini_train::checkpoint

infini_train/include/checkpoint/load_strategy.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <filesystem>
+#include <memory>
+#include <string>
+#include <unordered_map>
+
+#include "infini_train/include/checkpoint/load_planner.h"
+
+namespace infini_train {
+class Tensor;
+}
+
+namespace infini_train::checkpoint {
+
+using LoadedStateDict = std::unordered_map<std::string, std::shared_ptr<Tensor>>;
+
+class LoadStrategy {
+public:
+    virtual ~LoadStrategy() = default;
+    virtual LoadedStateDict Execute(const std::filesystem::path &checkpoint_dir, const LoadPlan &plan) = 0;
+};
+
+/// Reads source regions directly from metadata offsets while caching one open stream per file.
+class IndexedRegionLoadStrategy final : public LoadStrategy {
+public:
+    LoadedStateDict Execute(const std::filesystem::path &checkpoint_dir, const LoadPlan &plan) override;
+};
+
+} // namespace infini_train::checkpoint

infini_train/include/checkpoint/reshard.h

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <filesystem>
+
+#include "infini_train/include/checkpoint/checkpoint.h"
+
+namespace infini_train {
+class LRScheduler;
+class Optimizer;
+namespace nn {
+class Module;
+}
+} // namespace infini_train
+
+namespace infini_train::checkpoint {
+
+// Restore this rank's target model shards from a distributed checkpoint.
+void LoadDistributedCheckpoint(const std::filesystem::path &checkpoint_dir, nn::Module &model, Optimizer *optimizer,
+                               TrainerState &state, bool load_optimizer_state, LRScheduler *lr_scheduler,
+                               const Checkpoint::CheckpointMetadata &metadata);
+
+} // namespace infini_train::checkpoint

infini_train/include/checkpoint/save_planner.h

@@ -0,0 +1,80 @@
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "infini_train/include/checkpoint/shard_spec.h"
+#include "infini_train/include/datatype.h"
+
+namespace infini_train {
+class Tensor;
+}
+
+namespace infini_train::checkpoint {
+
+// Physical write description for one local tensor shard.
+struct WriteItem {
+    std::string key;
+    std::string filename;   // "model.ckpt" or "optimizer.ckpt"
+    uint64_t offset = 0;    // Planned byte offset in the checkpoint file.
+    uint64_t byte_size = 0; // Tensor payload size in bytes.
+    DataType dtype = DataType::kFLOAT32;
+    std::vector<int64_t> local_shape;
+    std::vector<int64_t> global_offset;
+    std::vector<int> axis_fragmentations;
+    int replica_id = 0;
+    int rank = 0;
+};
+
+// Build the local tensor write layout from a ShardedStateDict.
+class SavePlanner {
+public:
+    static std::vector<WriteItem> Plan(const ShardedStateDict &sd, int rank);
+};
+
+ShardedStateDict
+BuildOptimizerShardedStateDict(const ShardedStateDict &model_state,
+                               const std::unordered_map<std::string, std::shared_ptr<Tensor>> &optimizer_state);
+
+// Return the number of payload bytes required by a tensor.
+inline uint64_t TensorByteSize(DataType dtype, const std::vector<int64_t> &shape) {
+    uint64_t numel = 1;
+    for (auto d : shape) { numel *= static_cast<uint64_t>(d); }
+    switch (dtype) {
+    case DataType::kBFLOAT16:
+    case DataType::kFLOAT16:
+        return numel * 2;
+    case DataType::kFLOAT32:
+        return numel * 4;
+    case DataType::kFLOAT64:
+    case DataType::kINT64:
+    case DataType::kUINT64:
+        return numel * 8;
+    case DataType::kINT32:
+    case DataType::kUINT32:
+        return numel * 4;
+    case DataType::kINT16:
+    case DataType::kUINT16:
+        return numel * 2;
+    case DataType::kINT8:
+    case DataType::kUINT8:
+    case DataType::kBOOL:
+        return numel;
+    default:
+        return numel * 4;
+    }
+}
+
+// Compute one rank's balanced interval, including non-divisible dimensions.
+inline std::pair<int64_t, int64_t> GetRankSliceRange(int64_t global_size, int world_size, int rank) {
+    int64_t per_rank = global_size / world_size;
+    int64_t remainder = global_size % world_size;
+    int64_t start = rank * per_rank + std::min<int64_t>(rank, remainder);
+    int64_t local_size = per_rank + (rank < remainder ? 1 : 0);
+    return {start, local_size};
+}
+
+} // namespace infini_train::checkpoint