[Inference] Some effective methods to reduce the loading time of pir models (#70219)

aooxin · web-flow · commit 59ebcb7c9728 · 2024-12-17T22:28:44.000-08:00
* params_sync_among_devices_pass support multi threads and multi streams

* add separate_parameters for pir jit.save

* pir model load support multi threads

* del. useless cudaStream and extra load_pir_model_

* params_sync_among_devices_pass: Adjust thread count dynamically

* add unittest for ir and pir inference_save_or_load_sep_model

* del. useless code

* del. useless code

* align num_thread logic

* fix some bugs for model path

* align to ir

* del. useless code
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -18,6 +18,7 @@
 
 #include <algorithm>
 #include <cstdlib>
+#include <filesystem>
 #include <fstream>
 #include <memory>
 #include <set>
@@ -439,8 +440,21 @@ bool AnalysisPredictor::Init(
   paddle::platform::SetNumThreads(config_.cpu_math_library_num_threads());
 
   std::string model_path = config_.prog_file();
-  load_pir_model_ =
-      model_path.substr(model_path.find_last_of(".") + 1) == "json";
+  if (!model_path.empty()) {
+    load_pir_model_ =
+        model_path.substr(model_path.find_last_of(".") + 1) == "json";
+  } else if (!config_.model_dir().empty()) {
+    std::string model_dir = config_.model_dir();
+    load_pir_model_ = false;
+    for (const auto &entry : std::filesystem::directory_iterator(model_dir)) {
+      if (entry.is_regular_file() &&
+          entry.path().filename() == "__model__.json") {
+        load_pir_model_ = true;
+        config_.SetProgFile(config_.model_dir() + "/__model__.json");
+        break;
+      }
+    }
+  }
   if (load_pir_model_) {
     config_.use_pir_ = true;
     config_.use_new_executor_ = true;
@@ -1229,6 +1243,11 @@ bool AnalysisPredictor::SaveOrLoadPirParameters(bool for_save) {
               << " persistable is false, will ignore it when load variables.";
     }
   }
+  bool load_separate_params_ = true;
+  if (!for_save && config_.model_dir().empty()) {
+    // Combine model
+    load_separate_params_ = false;
+  }
 
   if (for_save) {
     std::string optimized_params =
@@ -1239,8 +1258,72 @@ bool AnalysisPredictor::SaveOrLoadPirParameters(bool for_save) {
         const_tensor_out, param_names, optimized_params, true, false, true);
     LOG(INFO) << "Optimized params saved to " << optimized_params;
   } else {
-    pir::LoadCombineFunction(
-        config_.params_file(), filter_param_names, &tensor_out, false, place_);
+    if (load_separate_params_) {
+      std::string params_dir = config_.model_dir();
+
+      auto process_params = [this, &params_dir, &filter_param_names](
+                                size_t start_idx, size_t end_idx) {
+        std::vector<phi::DenseTensor *> local_tensor_out;
+
+        for (size_t j = start_idx; j < end_idx; ++j) {
+          const auto &param_name = filter_param_names[j];
+          std::string param_file = params_dir + "/" + param_name;
+
+          auto *var = sub_scope_->FindVar(param_name);
+          VLOG(4) << "persistable variable's name: " << param_name;
+          if (var == nullptr) {
+            VLOG(4) << "Variable " << param_name << " not found in scope";
+            continue;
+          }
+          auto *tensor_temp = var->GetMutable<phi::DenseTensor>();
+
+          pir::LoadFunction(param_file, -1, {}, false, tensor_temp, place_);
+
+          local_tensor_out.push_back(tensor_temp);
+        }
+
+        return local_tensor_out;
+      };
+
+      size_t num_threads = 8;
+      size_t chunk_size = std::max(static_cast<size_t>(1),
+                                   filter_param_names.size() / num_threads);
+      num_threads =
+          std::min(num_threads, filter_param_names.size() / chunk_size);
+      size_t remain_size = filter_param_names.size() % num_threads;
+      VLOG(4) << "Start Load with multi-thread: " << num_threads
+              << " chund size: " << chunk_size;
+
+      std::vector<std::future<std::vector<phi::DenseTensor *>>> futures;
+
+      for (size_t i = 0; i < num_threads; ++i) {
+        size_t start_idx = i * chunk_size;
+        size_t end_idx = start_idx + chunk_size;
+
+        futures.push_back(
+            std::async(std::launch::async, process_params, start_idx, end_idx));
+      }
+      if (remain_size > 0) {
+        futures.push_back(std::async(std::launch::async,
+                                     process_params,
+                                     filter_param_names.size() - remain_size,
+                                     filter_param_names.size()));
+      }
+
+      std::vector<phi::DenseTensor *> tensor_out;
+      for (auto &future : futures) {
+        auto local_tensor_out = future.get();
+        tensor_out.insert(
+            tensor_out.end(), local_tensor_out.begin(), local_tensor_out.end());
+      }
+
+    } else {
+      pir::LoadCombineFunction(config_.params_file(),
+                               filter_param_names,
+                               &tensor_out,
+                               false,
+                               place_);
+    }
   }
   return true;
 }
diff --git a/paddle/fluid/pir/transforms/general/params_sync_among_devices_pass.cc b/paddle/fluid/pir/transforms/general/params_sync_among_devices_pass.cc
@@ -64,35 +64,71 @@ class ParamsSyncAmongDevicesPass : public pir::Pass {
             "params_sync_among_devices_pass should run on module op."));
     auto& block = module_op.block();
     int64_t num_rewrites_{0};
+
+    std::vector<phi::DenseTensor*> dense_tensors;
     for (auto& inner_op : block) {
-      if (inner_op.isa<pir::ParameterOp>() && inner_op.num_results() > 0) {
+      if (inner_op.template isa<pir::ParameterOp>() &&
+          inner_op.num_results() > 0) {
         auto var = inner_op.result(0);
         auto bool_attr =
-            var.attribute<::pir::BoolAttribute>(kAttrIsPersistable);
+            var.template attribute<::pir::BoolAttribute>(kAttrIsPersistable);
         if (!bool_attr || !bool_attr.data()) {
           continue;
         }
         std::string param_name = inner_op.attributes()
                                      .at("parameter_name")
-                                     .dyn_cast<pir::StrAttribute>()
+                                     .template dyn_cast<pir::StrAttribute>()
                                      .AsString();
         auto* param_var = scope_->FindVar(param_name);
         PADDLE_ENFORCE_NOT_NULL(
             param_var,
             common::errors::InvalidArgument("Parameter var [%s] not in scope.",
                                             param_name));
+
         if (param_var->IsType<phi::DenseTensor>()) {
-          auto* param_tensor = param_var->GetMutable<phi::DenseTensor>();
-          paddle::framework::TensorCopySync(
-              *param_tensor, place_, param_tensor);
-          num_rewrites_++;
+          dense_tensors.push_back(param_var->GetMutable<phi::DenseTensor>());
         } else {
           PADDLE_THROW(common::errors::Unimplemented(
               "params_sync_among_devices_pass only support DenseTensor type of "
               "parameter var."));
         }
       }
     }
+    num_rewrites_ = dense_tensors.size();
+
+    size_t num_threads = 8;
+    const size_t chunk_size =
+        std::max(static_cast<size_t>(1), dense_tensors.size() / num_threads);
+    num_threads = std::min(num_threads, dense_tensors.size() / chunk_size);
+    size_t remain_size = dense_tensors.size() % num_threads;
+
+    auto sync_handler = [&](const std::vector<phi::DenseTensor*>& tensors) {
+      for (auto* tensor : tensors) {
+        paddle::framework::TensorCopySync(*tensor, place_, tensor);
+      }
+    };
+
+    std::vector<std::future<void>> futures;
+    for (size_t i = 0; i < num_threads; ++i) {
+      auto start_it = dense_tensors.begin() + i * chunk_size;
+      auto end_it = start_it + chunk_size;
+
+      futures.push_back(
+          std::async(std::launch::async,
+                     sync_handler,
+                     std::vector<phi::DenseTensor*>(start_it, end_it)));
+    }
+    if (remain_size > 0) {
+      futures.push_back(std::async(
+          std::launch::async,
+          sync_handler,
+          std::vector<phi::DenseTensor*>(
+              dense_tensors.rbegin(), dense_tensors.rbegin() + remain_size)));
+    }
+
+    for (auto& future : futures) {
+      future.wait();
+    }
     AddStatistics(num_rewrites_);
   }
 
diff --git a/python/paddle/static/pir_io.py b/python/paddle/static/pir_io.py
@@ -765,7 +765,16 @@ def save_inference_model_pir(
     readable = kwargs.get('readable', False)
     trainable = kwargs.get('trainable', True)
     paddle.core.serialize_pir_program(
-        program, model_path, 1, True, readable, trainable
+        program,
+        (
+            os.path.join(os.path.dirname(model_path), "__model__.json")
+            if kwargs.get('separate_parameters', False)
+            else model_path
+        ),
+        1,
+        True,
+        readable,
+        trainable,
     )
 
     # serialize and save params
@@ -774,7 +783,11 @@ def save_inference_model_pir(
     save_vars_pir(
         dirname=save_dirname,
         main_program=program,
-        filename=params_filename,
+        filename=(
+            None
+            if kwargs.get('separate_parameters', False)
+            else params_filename
+        ),
     )
 
 
diff --git a/test/ir/inference/test_inference_save_or_load_sep_model.py b/test/ir/inference/test_inference_save_or_load_sep_model.py
