data_parallel_tree_learner.cpp 11.1 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
#include <cstring>
#include <tuple>
#include <vector>

5
6
#include "parallel_tree_learner.h"

Guolin Ke's avatar
Guolin Ke committed
7
8
namespace LightGBM {

9
template <typename TREELEARNER_T>
Guolin Ke's avatar
Guolin Ke committed
10
11
DataParallelTreeLearner<TREELEARNER_T>::DataParallelTreeLearner(const Config* config)
  :TREELEARNER_T(config) {
Guolin Ke's avatar
Guolin Ke committed
12
13
}

14
15
template <typename TREELEARNER_T>
DataParallelTreeLearner<TREELEARNER_T>::~DataParallelTreeLearner() {
Guolin Ke's avatar
Guolin Ke committed
16
17
}

18
19
template <typename TREELEARNER_T>
void DataParallelTreeLearner<TREELEARNER_T>::Init(const Dataset* train_data, bool is_constant_hessian) {
Guolin Ke's avatar
Guolin Ke committed
20
  // initialize SerialTreeLearner
21
  TREELEARNER_T::Init(train_data, is_constant_hessian);
Guolin Ke's avatar
Guolin Ke committed
22
23
24
25
  // Get local rank and global machine size
  rank_ = Network::rank();
  num_machines_ = Network::num_machines();
  // allocate buffer for communication
26
  size_t buffer_size = this->train_data_->NumTotalBin() * sizeof(HistogramBinEntry);
Guolin Ke's avatar
Guolin Ke committed
27

Guolin Ke's avatar
Guolin Ke committed
28
29
  input_buffer_.resize(buffer_size);
  output_buffer_.resize(buffer_size);
Guolin Ke's avatar
Guolin Ke committed
30

31
  is_feature_aggregated_.resize(this->num_features_);
Guolin Ke's avatar
Guolin Ke committed
32

Guolin Ke's avatar
Guolin Ke committed
33
34
  block_start_.resize(num_machines_);
  block_len_.resize(num_machines_);
Guolin Ke's avatar
Guolin Ke committed
35

36
37
  buffer_write_start_pos_.resize(this->num_features_);
  buffer_read_start_pos_.resize(this->num_features_);
Guolin Ke's avatar
Guolin Ke committed
38
  global_data_count_in_leaf_.resize(this->config_->num_leaves);
Guolin Ke's avatar
Guolin Ke committed
39
40
}

41
template <typename TREELEARNER_T>
Guolin Ke's avatar
Guolin Ke committed
42
43
44
void DataParallelTreeLearner<TREELEARNER_T>::ResetConfig(const Config* config) {
  TREELEARNER_T::ResetConfig(config);
  global_data_count_in_leaf_.resize(this->config_->num_leaves);
Guolin Ke's avatar
Guolin Ke committed
45
}
Guolin Ke's avatar
Guolin Ke committed
46

47
48
49
template <typename TREELEARNER_T>
void DataParallelTreeLearner<TREELEARNER_T>::BeforeTrain() {
  TREELEARNER_T::BeforeTrain();
Guolin Ke's avatar
Guolin Ke committed
50
51
52
  // generate feature partition for current tree
  std::vector<std::vector<int>> feature_distribution(num_machines_, std::vector<int>());
  std::vector<int> num_bins_distributed(num_machines_, 0);
53
54
  for (int i = 0; i < this->train_data_->num_total_features(); ++i) {
    int inner_feature_index = this->train_data_->InnerFeatureIndex(i);
Guolin Ke's avatar
Guolin Ke committed
55
    if (inner_feature_index == -1) { continue; }
56
    if (this->is_feature_used_[inner_feature_index]) {
Guolin Ke's avatar
Guolin Ke committed
57
      int cur_min_machine = static_cast<int>(ArrayArgs<int>::ArgMin(num_bins_distributed));
Guolin Ke's avatar
Guolin Ke committed
58
      feature_distribution[cur_min_machine].push_back(inner_feature_index);
59
60
      auto num_bin = this->train_data_->FeatureNumBin(inner_feature_index);
      if (this->train_data_->FeatureBinMapper(inner_feature_index)->GetDefaultBin() == 0) {
Guolin Ke's avatar
Guolin Ke committed
61
62
63
        num_bin -= 1;
      }
      num_bins_distributed[cur_min_machine] += num_bin;
Guolin Ke's avatar
Guolin Ke committed
64
    }
Guolin Ke's avatar
Guolin Ke committed
65
    is_feature_aggregated_[inner_feature_index] = false;
Guolin Ke's avatar
Guolin Ke committed
66
67
68
69
70
71
72
73
74
75
76
  }
  // get local used feature
  for (auto fid : feature_distribution[rank_]) {
    is_feature_aggregated_[fid] = true;
  }

  // get block start and block len for reduce scatter
  reduce_scatter_size_ = 0;
  for (int i = 0; i < num_machines_; ++i) {
    block_len_[i] = 0;
    for (auto fid : feature_distribution[i]) {
77
78
      auto num_bin = this->train_data_->FeatureNumBin(fid);
      if (this->train_data_->FeatureBinMapper(fid)->GetDefaultBin() == 0) {
Guolin Ke's avatar
Guolin Ke committed
79
80
81
        num_bin -= 1;
      }
      block_len_[i] += num_bin * sizeof(HistogramBinEntry);
Guolin Ke's avatar
Guolin Ke committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    }
    reduce_scatter_size_ += block_len_[i];
  }

  block_start_[0] = 0;
  for (int i = 1; i < num_machines_; ++i) {
    block_start_[i] = block_start_[i - 1] + block_len_[i - 1];
  }

  // get buffer_write_start_pos_
  int bin_size = 0;
  for (int i = 0; i < num_machines_; ++i) {
    for (auto fid : feature_distribution[i]) {
      buffer_write_start_pos_[fid] = bin_size;
96
97
      auto num_bin = this->train_data_->FeatureNumBin(fid);
      if (this->train_data_->FeatureBinMapper(fid)->GetDefaultBin() == 0) {
Guolin Ke's avatar
Guolin Ke committed
98
99
100
        num_bin -= 1;
      }
      bin_size += num_bin * sizeof(HistogramBinEntry);
Guolin Ke's avatar
Guolin Ke committed
101
102
103
104
105
106
107
    }
  }

  // get buffer_read_start_pos_
  bin_size = 0;
  for (auto fid : feature_distribution[rank_]) {
    buffer_read_start_pos_[fid] = bin_size;
108
109
    auto num_bin = this->train_data_->FeatureNumBin(fid);
    if (this->train_data_->FeatureBinMapper(fid)->GetDefaultBin() == 0) {
Guolin Ke's avatar
Guolin Ke committed
110
111
112
      num_bin -= 1;
    }
    bin_size += num_bin * sizeof(HistogramBinEntry);
Guolin Ke's avatar
Guolin Ke committed
113
114
115
  }

  // sync global data sumup info
116
117
  std::tuple<data_size_t, double, double> data(this->smaller_leaf_splits_->num_data_in_leaf(),
                                               this->smaller_leaf_splits_->sum_gradients(), this->smaller_leaf_splits_->sum_hessians());
Guolin Ke's avatar
Guolin Ke committed
118
  int size = sizeof(data);
Guolin Ke's avatar
Guolin Ke committed
119
  std::memcpy(input_buffer_.data(), &data, size);
Guolin Ke's avatar
Guolin Ke committed
120
  // global sumup reduce
Guolin Ke's avatar
Guolin Ke committed
121
122
  Network::Allreduce(input_buffer_.data(), size, sizeof(std::tuple<data_size_t, double, double>), output_buffer_.data(), [](const char *src, char *dst, int type_size, comm_size_t len) {
    comm_size_t used_size = 0;
Guolin Ke's avatar
Guolin Ke committed
123
124
    const std::tuple<data_size_t, double, double> *p1;
    std::tuple<data_size_t, double, double> *p2;
Guolin Ke's avatar
Guolin Ke committed
125
    while (used_size < len) {
Guolin Ke's avatar
Guolin Ke committed
126
127
      p1 = reinterpret_cast<const std::tuple<data_size_t, double, double> *>(src);
      p2 = reinterpret_cast<std::tuple<data_size_t, double, double> *>(dst);
Guolin Ke's avatar
Guolin Ke committed
128
129
130
131
132
133
134
135
136
      std::get<0>(*p2) = std::get<0>(*p2) + std::get<0>(*p1);
      std::get<1>(*p2) = std::get<1>(*p2) + std::get<1>(*p1);
      std::get<2>(*p2) = std::get<2>(*p2) + std::get<2>(*p1);
      src += type_size;
      dst += type_size;
      used_size += type_size;
    }
  });
  // copy back
Tsukasa OMOTO's avatar
Tsukasa OMOTO committed
137
  std::memcpy((void*)&data, output_buffer_.data(), size);
Guolin Ke's avatar
Guolin Ke committed
138
  // set global sumup info
139
  this->smaller_leaf_splits_->Init(std::get<1>(data), std::get<2>(data));
Guolin Ke's avatar
Guolin Ke committed
140
141
142
143
  // init global data count in leaf
  global_data_count_in_leaf_[0] = std::get<0>(data);
}

144
template <typename TREELEARNER_T>
Guolin Ke's avatar
Guolin Ke committed
145
146
void DataParallelTreeLearner<TREELEARNER_T>::FindBestSplits() {
  TREELEARNER_T::ConstructHistograms(this->is_feature_used_, true);
Guolin Ke's avatar
Guolin Ke committed
147
  // construct local histograms
148
  #pragma omp parallel for schedule(static)
149
150
  for (int feature_index = 0; feature_index < this->num_features_; ++feature_index) {
    if ((!this->is_feature_used_.empty() && this->is_feature_used_[feature_index] == false)) continue;
Guolin Ke's avatar
Guolin Ke committed
151
    // copy to buffer
Guolin Ke's avatar
Guolin Ke committed
152
    std::memcpy(input_buffer_.data() + buffer_write_start_pos_[feature_index],
153
154
                this->smaller_leaf_histogram_array_[feature_index].RawData(),
                this->smaller_leaf_histogram_array_[feature_index].SizeOfHistgram());
Guolin Ke's avatar
Guolin Ke committed
155
156
  }
  // Reduce scatter for histogram
Guolin Ke's avatar
Guolin Ke committed
157
158
  Network::ReduceScatter(input_buffer_.data(), reduce_scatter_size_, sizeof(HistogramBinEntry), block_start_.data(),
                         block_len_.data(), output_buffer_.data(), static_cast<comm_size_t>(output_buffer_.size()), &HistogramBinEntry::SumReducer);
Guolin Ke's avatar
Guolin Ke committed
159
160
161
162
163
164
165
  this->FindBestSplitsFromHistograms(this->is_feature_used_, true);
}

template <typename TREELEARNER_T>
void DataParallelTreeLearner<TREELEARNER_T>::FindBestSplitsFromHistograms(const std::vector<int8_t>&, bool) {
  std::vector<SplitInfo> smaller_bests_per_thread(this->num_threads_, SplitInfo());
  std::vector<SplitInfo> larger_bests_per_thread(this->num_threads_, SplitInfo());
Guolin Ke's avatar
Guolin Ke committed
166

167
168
  OMP_INIT_EX();
  #pragma omp parallel for schedule(static)
169
  for (int feature_index = 0; feature_index < this->num_features_; ++feature_index) {
170
    OMP_LOOP_EX_BEGIN();
Guolin Ke's avatar
Guolin Ke committed
171
    if (!is_feature_aggregated_[feature_index]) continue;
Guolin Ke's avatar
Guolin Ke committed
172
    const int tid = omp_get_thread_num();
Guolin Ke's avatar
Guolin Ke committed
173
    const int real_feature_index = this->train_data_->RealFeatureIndex(feature_index);
Guolin Ke's avatar
Guolin Ke committed
174
    // restore global histograms from buffer
175
    this->smaller_leaf_histogram_array_[feature_index].FromMemory(
Guolin Ke's avatar
Guolin Ke committed
176
      output_buffer_.data() + buffer_read_start_pos_[feature_index]);
Guolin Ke's avatar
Guolin Ke committed
177

178
    this->train_data_->FixHistogram(feature_index,
Guolin Ke's avatar
Guolin Ke committed
179
180
181
                                    this->smaller_leaf_splits_->sum_gradients(), this->smaller_leaf_splits_->sum_hessians(),
                                    GetGlobalDataCountInLeaf(this->smaller_leaf_splits_->LeafIndex()),
                                    this->smaller_leaf_histogram_array_[feature_index].RawData());
Guolin Ke's avatar
Guolin Ke committed
182
    SplitInfo smaller_split;
Guolin Ke's avatar
Guolin Ke committed
183
    // find best threshold for smaller child
184
185
186
187
    this->smaller_leaf_histogram_array_[feature_index].FindBestThreshold(
      this->smaller_leaf_splits_->sum_gradients(),
      this->smaller_leaf_splits_->sum_hessians(),
      GetGlobalDataCountInLeaf(this->smaller_leaf_splits_->LeafIndex()),
Guolin Ke's avatar
Guolin Ke committed
188
189
      this->smaller_leaf_splits_->min_constraint(),
      this->smaller_leaf_splits_->max_constraint(),
Guolin Ke's avatar
Guolin Ke committed
190
      &smaller_split);
Guolin Ke's avatar
Guolin Ke committed
191
192
193
    smaller_split.feature = real_feature_index;
    if (smaller_split > smaller_bests_per_thread[tid]) {
      smaller_bests_per_thread[tid] = smaller_split;
Guolin Ke's avatar
Guolin Ke committed
194
    }
Guolin Ke's avatar
Guolin Ke committed
195
196

    // only root leaf
197
    if (this->larger_leaf_splits_ == nullptr || this->larger_leaf_splits_->LeafIndex() < 0) continue;
Guolin Ke's avatar
Guolin Ke committed
198
199

    // construct histgroms for large leaf, we init larger leaf as the parent, so we can just subtract the smaller leaf's histograms
200
201
    this->larger_leaf_histogram_array_[feature_index].Subtract(
      this->smaller_leaf_histogram_array_[feature_index]);
Guolin Ke's avatar
Guolin Ke committed
202
    SplitInfo larger_split;
Guolin Ke's avatar
Guolin Ke committed
203
    // find best threshold for larger child
204
205
206
207
    this->larger_leaf_histogram_array_[feature_index].FindBestThreshold(
      this->larger_leaf_splits_->sum_gradients(),
      this->larger_leaf_splits_->sum_hessians(),
      GetGlobalDataCountInLeaf(this->larger_leaf_splits_->LeafIndex()),
Guolin Ke's avatar
Guolin Ke committed
208
209
      this->larger_leaf_splits_->min_constraint(),
      this->larger_leaf_splits_->max_constraint(),
Guolin Ke's avatar
Guolin Ke committed
210
      &larger_split);
Guolin Ke's avatar
Guolin Ke committed
211
212
213
    larger_split.feature = real_feature_index;
    if (larger_split > larger_bests_per_thread[tid]) {
      larger_bests_per_thread[tid] = larger_split;
Guolin Ke's avatar
Guolin Ke committed
214
    }
215
    OMP_LOOP_EX_END();
Guolin Ke's avatar
Guolin Ke committed
216
  }
217
  OMP_THROW_EX();
Guolin Ke's avatar
Guolin Ke committed
218

Guolin Ke's avatar
Guolin Ke committed
219
220
221
  auto smaller_best_idx = ArrayArgs<SplitInfo>::ArgMax(smaller_bests_per_thread);
  int leaf = this->smaller_leaf_splits_->LeafIndex();
  this->best_split_per_leaf_[leaf] = smaller_bests_per_thread[smaller_best_idx];
Guolin Ke's avatar
Guolin Ke committed
222

Guolin Ke's avatar
Guolin Ke committed
223
224
225
226
227
  if (this->larger_leaf_splits_ != nullptr &&  this->larger_leaf_splits_->LeafIndex() >= 0) {
    leaf = this->larger_leaf_splits_->LeafIndex();
    auto larger_best_idx = ArrayArgs<SplitInfo>::ArgMax(larger_bests_per_thread);
    this->best_split_per_leaf_[leaf] = larger_bests_per_thread[larger_best_idx];
  }
Guolin Ke's avatar
Guolin Ke committed
228

Guolin Ke's avatar
Guolin Ke committed
229
230
  SplitInfo smaller_best_split, larger_best_split;
  smaller_best_split = this->best_split_per_leaf_[this->smaller_leaf_splits_->LeafIndex()];
Guolin Ke's avatar
Guolin Ke committed
231
  // find local best split for larger leaf
232
  if (this->larger_leaf_splits_->LeafIndex() >= 0) {
Guolin Ke's avatar
Guolin Ke committed
233
    larger_best_split = this->best_split_per_leaf_[this->larger_leaf_splits_->LeafIndex()];
Guolin Ke's avatar
Guolin Ke committed
234
235
236
  }

  // sync global best info
Guolin Ke's avatar
Guolin Ke committed
237
  SyncUpGlobalBestSplit(input_buffer_.data(), input_buffer_.data(), &smaller_best_split, &larger_best_split, this->config_->max_cat_threshold);
Guolin Ke's avatar
Guolin Ke committed
238
239

  // set best split
Guolin Ke's avatar
Guolin Ke committed
240
  this->best_split_per_leaf_[this->smaller_leaf_splits_->LeafIndex()] = smaller_best_split;
241
  if (this->larger_leaf_splits_->LeafIndex() >= 0) {
Guolin Ke's avatar
Guolin Ke committed
242
    this->best_split_per_leaf_[this->larger_leaf_splits_->LeafIndex()] = larger_best_split;
Guolin Ke's avatar
Guolin Ke committed
243
244
245
  }
}

246
247
248
249
template <typename TREELEARNER_T>
void DataParallelTreeLearner<TREELEARNER_T>::Split(Tree* tree, int best_Leaf, int* left_leaf, int* right_leaf) {
  TREELEARNER_T::Split(tree, best_Leaf, left_leaf, right_leaf);
  const SplitInfo& best_split_info = this->best_split_per_leaf_[best_Leaf];
Guolin Ke's avatar
Guolin Ke committed
250
251
252
253
254
  // need update global number of data in leaf
  global_data_count_in_leaf_[*left_leaf] = best_split_info.left_count;
  global_data_count_in_leaf_[*right_leaf] = best_split_info.right_count;
}

255
256
257
// instantiate template classes, otherwise linker cannot find the code
template class DataParallelTreeLearner<GPUTreeLearner>;
template class DataParallelTreeLearner<SerialTreeLearner>;
Guolin Ke's avatar
Guolin Ke committed
258
259

}  // namespace LightGBM