fix single machine gbdt

microsoft · Oct 25, 2024 · b56b39e · b56b39e
1 parent 4bb4411
commit b56b39e
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 6 deletions.
diff --git a/include/LightGBM/cuda/cuda_column_data.hpp b/include/LightGBM/cuda/cuda_column_data.hpp
@@ -13,6 +13,7 @@
 #include <LightGBM/bin.h>
 #include <LightGBM/utils/openmp_wrapper.h>
 
+#include <memory>
 #include <vector>
 
 namespace LightGBM {

diff --git a/include/LightGBM/cuda/cuda_nccl_topology.hpp b/include/LightGBM/cuda/cuda_nccl_topology.hpp
@@ -9,6 +9,7 @@
 #include <LightGBM/network.h>
 #include <LightGBM/utils/common.h>
 
+#include <memory>
 #include <set>
 #include <string>
 #include <vector>

diff --git a/src/boosting/cuda/nccl_gbdt.cpp b/src/boosting/cuda/nccl_gbdt.cpp
@@ -28,6 +28,8 @@ void NCCLGBDT<GBDT_T>::Init(
   const std::vector<const Metric*>& training_metrics) {
   GBDT_T::Init(gbdt_config, train_data, objective_function, training_metrics);
 
+  this->tree_learner_.reset();
+
   nccl_topology_.reset(new NCCLTopology(this->config_->gpu_device_id, this->config_->num_gpus, this->config_->gpu_device_id_list, train_data->num_data()));
 
   nccl_topology_->InitNCCL();

diff --git a/src/boosting/cuda/nccl_gbdt.hpp b/src/boosting/cuda/nccl_gbdt.hpp
@@ -8,12 +8,13 @@
 
 #ifdef USE_CUDA
 
+#include <LightGBM/objective_function.h>
+#include <LightGBM/network.h>
+
 #include "cuda_score_updater.hpp"
 #include "nccl_gbdt_component.hpp"
 
 #include <LightGBM/cuda/cuda_nccl_topology.hpp>
-#include <LightGBM/objective_function.h>
-#include <LightGBM/network.h>
 
 #include <pthread.h>
 #include <memory>

diff --git a/src/boosting/gbdt.cpp b/src/boosting/gbdt.cpp
@@ -104,12 +104,12 @@ void GBDT::Init(const Config* config, const Dataset* train_data, const Objective
   boosting_on_gpu_ = objective_function_ != nullptr && objective_function_->IsCUDAObjective() &&
                      !data_sample_strategy_->IsHessianChange();  // for sample strategy with Hessian change, fall back to boosting on CPU
 
-  tree_learner_ = nullptr; // std::unique_ptr<TreeLearner>(TreeLearner::CreateTreeLearner(config_->tree_learner, config_->device_type,
-                                                                             // config_.get(), boosting_on_gpu_));
+  tree_learner_ = std::unique_ptr<TreeLearner>(TreeLearner::CreateTreeLearner(config_->tree_learner, config_->device_type,
+                                                                              config_.get(), boosting_on_gpu_));
 
   // init tree learner
-  // tree_learner_->Init(train_data_, is_constant_hessian_);
-  // tree_learner_->SetForcedSplit(&forced_splits_json_);
+  tree_learner_->Init(train_data_, is_constant_hessian_);
+  tree_learner_->SetForcedSplit(&forced_splits_json_);
 
   // push training metrics
   training_metrics_.clear();