Issue/414: Async model loader (#415)

Batch and parallelize model parameter loading by safetensor shard across rank workers.

Add vLLM-style safetensors loading with shard prefetch/pipelining, state dict reuse, missing-key validation, progress restoration.

Author: qinyiqun

csrc/engine/infer_engine.cpp

@@ -1,6 +1,7 @@
 #include "infer_engine.hpp"
 #include "../config/config_factory.hpp"
 #include "spdlog/spdlog.h"
+#include <future>
 
 namespace infinilm::engine {
 
@@ -54,6 +55,20 @@ void InferEngine::load_param(const std::string &name, const infinicore::Tensor &
         worker->load_param(name, param);
     }
 }
+
+void InferEngine::load_params(const std::unordered_map<std::string, infinicore::Tensor> &params) {
+    std::vector<std::future<void>> futures;
+    futures.reserve(workers_.size());
+    for (auto &worker : workers_) {
+        futures.emplace_back(std::async(std::launch::async, [&worker, &params] {
+            worker->load_params(params);
+        }));
+    }
+    for (auto &future : futures) {
+        future.get();
+    }
+}
+
 //------------------------------------------------------
 // load_param
 //------------------------------------------------------

csrc/engine/infer_engine.hpp

@@ -9,6 +9,7 @@
 #include "rank_worker.hpp"
 
 #include <optional>
+#include <unordered_map>
 #include <vector>
 
 namespace infinilm::engine {
@@ -32,6 +33,9 @@ class InferEngine {
     // Load a parameter to all workers (each can extract its shard inside RankWorker)
     void load_param(const std::string &name, const infinicore::Tensor &param);
 
+    // Load a batch of parameters to all workers, syncing each worker once after the batch.
+    void load_params(const std::unordered_map<std::string, infinicore::Tensor> &params);
+
     // process the weights after loading on all workers (e.g., for quantization)
     void process_weights_after_loading();
 

csrc/engine/rank_worker.cpp

@@ -88,6 +88,31 @@ void RankWorker::load_param(const std::string &name,
     }
 }
 
+//------------------------------------------------------
+// load_params -- synchronous batch load
+//------------------------------------------------------
+void RankWorker::load_params(const std::unordered_map<std::string, infinicore::Tensor> &params) {
+    {
+        std::lock_guard<std::mutex> lock(mutex_);
+        if (should_exit_) {
+            throw std::runtime_error("RankWorker is closing; cannot load_params");
+        }
+
+        pending_params_ = params;
+        job_cmd_ = Command::LOAD_BATCH;
+        has_job_ = true;
+        job_done_ = false;
+    }
+    cv_.notify_all();
+
+    std::unique_lock<std::mutex> lk(mutex_);
+    cv_.wait(lk, [&] { return job_done_ || should_exit_; });
+
+    if (should_exit_) {
+        throw std::runtime_error("RankWorker stopped while loading parameters");
+    }
+}
+
 //------------------------------------------------------
 // process_weights_after_loading -- asynchronous
 //------------------------------------------------------
@@ -266,6 +291,7 @@ void RankWorker::thread_loop() {
             Command local_cmd = Command::INIT;
             std::string local_param_name;
             infinicore::Tensor local_param;
+            std::unordered_map<std::string, infinicore::Tensor> local_params;
             Input local_args;
             std::unique_ptr<cache::CacheConfig> local_cache_config;
 
@@ -283,6 +309,9 @@ void RankWorker::thread_loop() {
                 if (local_cmd == Command::LOAD) {
                     local_param_name = pending_param_name_;
                     local_param = pending_param_;
+                } else if (local_cmd == Command::LOAD_BATCH) {
+                    local_params = std::move(pending_params_);
+                    pending_params_.clear();
                 } else if (local_cmd == Command::PREPROCESS) {
 
                 } else if (local_cmd == Command::RUN) {
@@ -319,6 +348,27 @@ void RankWorker::thread_loop() {
                 }
                 cv_.notify_all();
 
+            } else if (local_cmd == Command::LOAD_BATCH) {
+                try {
+                    model_->load_parameters_no_sync(local_params);
+                    infinicore::context::syncStream();
+                } catch (const std::exception &e) {
+                    {
+                        std::lock_guard<std::mutex> lk(mutex_);
+                        should_exit_ = true;
+                        job_done_ = true;
+                    }
+                    cv_.notify_all();
+                    spdlog::error("[{}] exception during load_parameters_: {}\n", info(), e.what());
+                    break;
+                }
+
+                {
+                    std::lock_guard<std::mutex> lk(mutex_);
+                    job_done_ = true;
+                }
+                cv_.notify_all();
+
             } else if (local_cmd == Command::PREPROCESS) {
                 // Handle preprocess command
                 try {

csrc/engine/rank_worker.hpp

@@ -15,6 +15,7 @@
 #include <random>
 #include <string>
 #include <thread>
+#include <unordered_map>
 #include <vector>
 
 namespace infinilm::engine {
@@ -25,6 +26,7 @@ class RankWorker {
     enum class Command {
         INIT,
         LOAD,
+        LOAD_BATCH,
         PREPROCESS,
         RUN,
         RESET_CACHE,
@@ -83,6 +85,8 @@ class RankWorker {
     void load_param(const std::string &name,
                     const infinicore::Tensor &param);
 
+    void load_params(const std::unordered_map<std::string, infinicore::Tensor> &params);
+
     void process_weights_after_loading();
 
     // return the parameters (i.e. weights and biases).
@@ -139,6 +143,7 @@ class RankWorker {
     // Task payloads (protected by mutex)
     std::string pending_param_name_;
     infinicore::Tensor pending_param_;
+    std::unordered_map<std::string, infinicore::Tensor> pending_params_;
     Input pending_args_;
     std::unique_ptr<cache::CacheConfig> pending_cache_config_;
 

csrc/pybind11/engine/engine.hpp

@@ -59,6 +59,9 @@ inline void bind_infer_engine(py::module &m) {
         .def("load_param", &InferEngine::load_param,
              py::arg("name"), py::arg("param"),
              "Load a parameter tensor into all workers (each worker picks its shard)")
+        .def("load_params", &InferEngine::load_params,
+             py::arg("params"),
+             "Load a batch of parameter tensors into all workers, syncing once per worker")
         .def("state_dict", [](InferEngine &self) {
             py::list state_dict_tp_all;
             for (const auto &state_dict_tp : self.state_dict()) {

python/infinilm/infer_engine.py

@@ -369,11 +369,10 @@ def reset_cache(self, cache_config):
         super().reset_cache(cache_config)
 
     def state_dict_keyname(self):
-        return super().state_dict()[0].keys()
+        return sorted({name for state_dict in super().state_dict() for name in state_dict.keys()})
 
     def load_state_dict(self, state_dict, strict=None):
-        for name, param in state_dict.items():
-            super().load_param(name, param._underlying)
+        super().load_params({name: param._underlying for name, param in state_dict.items()})
 
     def process_weights_after_loading(self):
         super().process_weights_after_loading()