[WIP] Feat/checkpoint resharding

example/gpt2/main.cc

@@ -232,6 +232,18 @@ void Train(const nn::parallel::Rank &rank) {
 
     model->To(device);
 
+    // select the data type
+    // TODO(lzm): change to solely rely on the weight file info for determining the dtype when autocast is supported
+    DataType dtype;
+    if (FLAGS_dtype == kDtypeFP32) {
+        dtype = DataType::kFLOAT32;
+    } else if (FLAGS_dtype == kDtypeBF16) {
+        dtype = DataType::kBFLOAT16;
+    } else {
+        LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported.";
+    }
+    model->To(dtype);
+
     utils::PrecisionChecker::BuildNameMap(model.get());
 
     // Get chunk size before wrapping with LoRA (needed for PipelineParallel)
@@ -257,17 +269,6 @@ void Train(const nn::parallel::Rank &rank) {
         nn::lora::PrintLoRASummary(model, rank.GlobalRank());
     }
 
-    // select the data type
-    // TODO(lzm): change to solely rely on the weight file info for determining the dtype when autocast is supported
-    DataType dtype;
-    if (FLAGS_dtype == kDtypeFP32) {
-        dtype = DataType::kFLOAT32;
-    } else if (FLAGS_dtype == kDtypeBF16) {
-        dtype = DataType::kBFLOAT16;
-    } else {
-        LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported.";
-    }
-
     auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size);
 
     // Create optimizer - use GetLoRAParameters if LoRA is enabled
@@ -395,6 +396,7 @@ void Train(const nn::parallel::Rank &rank) {
             .tp_size = tp_world_size,
             .sp_size = sp_world_size,
             .pp_size = pp_world_size,
+            .vpp_size = static_cast<int>(FLAGS_virtual_pipeline_parallel),
             .checkpoint_root_dir = FLAGS_save,
             .max_checkpoint_keep = FLAGS_max_checkpoint_keep,
             .rank = rank,

example/llama3/main.cc

@@ -219,6 +219,18 @@ void Train(const nn::parallel::Rank &rank) {
 
     model->To(device);
 
+    // select the data type
+    // TODO(lzm): change to solely rely on the weight file info for determining the dtype when autocast is supported
+    DataType dtype;
+    if (FLAGS_dtype == kDtypeFP32) {
+        dtype = DataType::kFLOAT32;
+    } else if (FLAGS_dtype == kDtypeBF16) {
+        dtype = DataType::kBFLOAT16;
+    } else {
+        LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported.";
+    }
+    model->To(dtype);
+
     utils::PrecisionChecker::BuildNameMap(model.get());
 
     // Apply LoRA using GetLoRAModel (in-place injection)
@@ -242,15 +254,6 @@ void Train(const nn::parallel::Rank &rank) {
 
     LOG(INFO) << "Rank " << rank.GlobalRank() << ": Model loaded to device.";
 
-    DataType dtype;
-    if (FLAGS_dtype == kDtypeFP32) {
-        dtype = DataType::kFLOAT32;
-    } else if (FLAGS_dtype == kDtypeBF16) {
-        dtype = DataType::kBFLOAT16;
-    } else {
-        LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported.";
-    }
-
     auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size);
 
     if (pp_world_size > 1) {
@@ -377,6 +380,7 @@ void Train(const nn::parallel::Rank &rank) {
             .tp_size = tp_world_size,
             .sp_size = sp_world_size,
             .pp_size = pp_world_size,
+            .vpp_size = static_cast<int>(FLAGS_virtual_pipeline_parallel),
             .checkpoint_root_dir = FLAGS_save,
             .max_checkpoint_keep = FLAGS_max_checkpoint_keep,
             .rank = rank,

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;
@@ -27,6 +32,7 @@ struct TrainerState {
     int tp_size = 1;
     int sp_size = 1;
     int pp_size = 1;
+    int vpp_size = 1;
 };
 
 class Checkpoint {
@@ -37,9 +43,70 @@ class Checkpoint {
     static void Load(const std::filesystem::path &checkpoint_dir, nn::Module &model, Optimizer *optimizer,
                      TrainerState &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;
+            int vpp_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;
+            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"
@@ -50,6 +49,7 @@ struct SaveCheckpointArgs {
     int tp_size = 1;
     int sp_size = 1;
     int pp_size = 1;
+    int vpp_size = 1;
     std::filesystem::path checkpoint_root_dir;
     size_t max_checkpoint_keep = 0;
     const nn::parallel::Rank &rank;

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, LRScheduler *lr_scheduler,
+                               const Checkpoint::CheckpointMetadata &metadata);
+
+} // namespace infini_train::checkpoint

infini_train/include/checkpoint/save_planner.h

@@ -0,0 +1,79 @@
+#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 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