diff --git a/src/common/hist_util.cu b/src/common/hist_util.cu index 5da1006ed..fe1305d4b 100644 --- a/src/common/hist_util.cu +++ b/src/common/hist_util.cu @@ -264,7 +264,6 @@ void ProcessWeightedBatch(int device, const SparsePage& page, d_temp_weights[idx] = weights[group]; }); } else { - CHECK_EQ(weights.size(), page.offset.Size() - 1); dh::LaunchN(device, temp_weights.size(), [=] __device__(size_t idx) { size_t element_idx = idx + begin; size_t ridx = thrust::upper_bound(thrust::seq, row_ptrs.begin(), diff --git a/tests/cpp/common/test_hist_util.cu b/tests/cpp/common/test_hist_util.cu index d8a75ba49..3ec49668a 100644 --- a/tests/cpp/common/test_hist_util.cu +++ b/tests/cpp/common/test_hist_util.cu @@ -209,6 +209,23 @@ TEST(HistUtil, DeviceSketchMultipleColumnsExternal) { } } +// See https://github.com/dmlc/xgboost/issues/5866. +TEST(HistUtil, DeviceSketchExternalMemoryWithWeights) { + int bin_sizes[] = {2, 16, 256, 512}; + int sizes[] = {100, 1000, 1500}; + int num_columns = 5; + dmlc::TemporaryDirectory temp; + for (auto num_rows : sizes) { + auto x = GenerateRandom(num_rows, num_columns); + auto dmat = GetExternalMemoryDMatrixFromData(x, num_rows, num_columns, 100, temp); + dmat->Info().weights_.HostVector() = GenerateRandomWeights(num_rows); + for (auto num_bins : bin_sizes) { + auto cuts = DeviceSketch(0, dmat.get(), num_bins); + ValidateCuts(cuts, dmat.get(), num_bins); + } + } +} + template void ValidateBatchedCuts(Adapter adapter, int num_bins, int num_columns, int num_rows, DMatrix* dmat) {