Feat/checkpoint resharding

example/gpt2/main.cc

@@ -405,6 +405,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

@@ -387,6 +387,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

infini_train/include/checkpoint/shard_spec.h

@@ -0,0 +1,56 @@
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <string>
+#include <vector>
+
+#include "glog/logging.h"
+
+#include "infini_train/include/datatype.h"
+
+namespace infini_train::checkpoint {
+
+struct ShardSegment {
+    int64_t global_offset = 0;
+    int64_t local_offset = 0;
+    int64_t length = 0;
+
+    bool operator==(const ShardSegment &other) const = default;
+};
+
+// Logical tensor shard metadata, aligned with Megatron-LM's ShardedTensor model.
+struct ShardedTensor {
+    std::string key;
+    std::string local_key;
+    DataType dtype = DataType::kFLOAT32;
+    std::vector<int64_t> global_shape;
+    std::vector<int64_t> local_shape;
+    std::vector<int64_t> global_offset;
+    std::vector<int> axis_fragmentations;
+    // Optional disjoint regions along the single fragmented axis. This is used
+    // by layouts such as rank-local [Q, K, V], which are not one contiguous
+    // slice of the logical global [Q, K, V] tensor.
+    std::vector<ShardSegment> segments;
+
+    bool operator==(const ShardedTensor &other) const {
+        return key == other.key && local_key == other.local_key && dtype == other.dtype
+            && global_shape == other.global_shape && local_shape == other.local_shape
+            && global_offset == other.global_offset && axis_fragmentations == other.axis_fragmentations
+            && segments == other.segments;
+    }
+};
+
+struct ShardedStateDict {
+    std::map<std::string, ShardedTensor> tensors;
+
+    void Merge(ShardedStateDict &&other) {
+        for (auto &[key, info] : other.tensors) {
+            const auto display_key = key;
+            const auto [_, inserted] = tensors.emplace(std::move(key), std::move(info));
+            CHECK(inserted) << "Duplicate sharded state-dict key: " << display_key;
+        }
+    }
+};
+
+} // namespace infini_train::checkpoint

infini_train/include/nn/lora/lora_parallel_linear.h

@@ -34,6 +34,8 @@ class LoRAColumnParallelLinear : public nn::parallel::ColumnParallelLinear {
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 
+    checkpoint::ShardedStateDict ShardedStateDict(const std::string &prefix = "") const override;
+
     void MergeWeights();
     void UnmergeWeights();
     bool IsMerged() const;
@@ -74,6 +76,8 @@ class LoRARowParallelLinear : public nn::parallel::RowParallelLinear {
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 
+    checkpoint::ShardedStateDict ShardedStateDict(const std::string &prefix = "") const override;
+
     void MergeWeights();
     void UnmergeWeights();
     bool IsMerged() const;