[WIP] Feat/checkpoint resharding

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,69 @@ 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;
+        } 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"

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/modules/module.h

@@ -6,6 +6,7 @@
 #include <unordered_set>
 #include <vector>
 
+#include "infini_train/include/checkpoint/shard_spec.h"
 #include "infini_train/include/datatype.h"
 #include "infini_train/include/device.h"
 
@@ -49,7 +50,7 @@ class Module : public std::enable_shared_from_this<Module> {
 
     // TODO: Change return type to filterable iterator (like PyTorch's named_parameters with prefix matching)
     virtual std::vector<std::shared_ptr<Tensor>> Parameters() const;
-    std::vector<std::pair<std::string, std::shared_ptr<Tensor>>>
+    virtual std::vector<std::pair<std::string, std::shared_ptr<Tensor>>>
     NamedParameters(const std::string &prefix = "", bool recurse = true, bool remove_duplicate = true) const;
     bool has_parameter(const std::string &name) const;
     std::shared_ptr<Tensor> *mutable_parameter(const std::string &name);
@@ -61,11 +62,14 @@ class Module : public std::enable_shared_from_this<Module> {
     std::shared_ptr<Module> &mutable_module(const std::string &name);
     const Module &module(const std::string &name) const;
 
-    std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const;
+    virtual std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const;
+
+    // Return state-dict metadata with global shard coordinates.
+    virtual checkpoint::ShardedStateDict ShardedStateDict(const std::string &prefix = "") const;
 
     // Current behavior: missing keys / shape / dtype mismatches are FATAL errors; unexpected keys in state_dict are
     // WARNING-only and silently ignored.
-    void LoadStateDict(const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
+    virtual void LoadStateDict(const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict);
 
     // operator() calls hooks and Forward
     std::vector<std::shared_ptr<Tensor>> operator()(const std::vector<std::shared_ptr<Tensor>> &input_tensors);

infini_train/include/nn/modules/transformer/causal_self_attention.h

@@ -21,6 +21,8 @@ class CausalSelfAttention : public infini_train::nn::CloneableModule<CausalSelfA
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
 
+    checkpoint::ShardedStateDict ShardedStateDict(const std::string &prefix = "") const override;
+
 private:
     TransformerConfig config_;
     int64_t n_head_ = 0;

infini_train/include/nn/modules/transformer/transformer.h

@@ -78,6 +78,11 @@ class TransformerModel : public CloneableModule<TransformerModel> {
 
     const TransformerConfig &Config() const { return config_; }
 
+    checkpoint::ShardedStateDict ShardedStateDict(const std::string &prefix = "") const override;
+    std::vector<std::pair<std::string, std::shared_ptr<Tensor>>>
+    NamedParameters(const std::string &prefix = "", bool recurse = true, bool remove_duplicate = true) const override;
+    void LoadStateDict(const std::unordered_map<std::string, std::shared_ptr<Tensor>> &state_dict) override;
+
 private:
     const TransformerConfig config_;
     const infini_train::nn::parallel::StageInfo stage_info_;