diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc index 0d77092fd96f..25b54a37e7de 100644 --- a/src/tree/updater_quantile_hist.cc +++ b/src/tree/updater_quantile_hist.cc @@ -26,8 +26,8 @@ #include "hist/evaluate_splits.h" // for HistEvaluator, HistMultiEvaluator, UpdatePre... #include "hist/expand_entry.h" // for MultiExpandEntry, CPUExpandEntry #include "hist/hist_cache.h" // for BoundedHistCollection -#include "hist/histogram.h" // for MultiHistogramBuilder #include "hist/hist_param.h" // for HistMakerTrainParam +#include "hist/histogram.h" // for MultiHistogramBuilder #include "hist/sampler.h" // for SampleGradient #include "param.h" // for TrainParam, GradStats #include "xgboost/base.h" // for Args, GradientPairPrecise, GradientPair, Gra... @@ -152,15 +152,23 @@ class MultiTargetHistBuilder { p_last_fmat_ = p_fmat; bst_bin_t n_total_bins = 0; - partitioner_.clear(); + size_t page_idx = 0; for (auto const &page : p_fmat->GetBatches(ctx_, HistBatch(param_))) { if (n_total_bins == 0) { n_total_bins = page.cut.TotalBins(); } else { CHECK_EQ(n_total_bins, page.cut.TotalBins()); } - partitioner_.emplace_back(ctx_, page.Size(), page.base_rowid, p_fmat->Info().IsColumnSplit()); + if (page_idx < partitioner_.size()) { + partitioner_[page_idx].Reset(ctx_, page.Size(), page.base_rowid, + p_fmat->Info().IsColumnSplit()); + } else { + partitioner_.emplace_back(ctx_, page.Size(), page.base_rowid, + p_fmat->Info().IsColumnSplit()); + } + page_idx++; } + partitioner_.resize(page_idx); bst_target_t n_targets = p_tree->NumTargets(); histogram_builder_ = std::make_unique(); @@ -364,6 +372,7 @@ class HistUpdater { } page_idx++; } + partitioner_.resize(page_idx); histogram_builder_->Reset(ctx_, n_total_bins, 1, HistBatch(param_), collective::IsDistributed(), fmat->Info().IsColumnSplit(), hist_param_); evaluator_ = std::make_unique(ctx_, this->param_, fmat->Info(), col_sampler_); diff --git a/tests/cpp/tree/test_quantile_hist.cc b/tests/cpp/tree/test_quantile_hist.cc index 362c6c782e09..d8e1e2c016ee 100644 --- a/tests/cpp/tree/test_quantile_hist.cc +++ b/tests/cpp/tree/test_quantile_hist.cc @@ -3,19 +3,24 @@ */ #include #include +#include #include +#include #include // for size_t +#include +#include #include #include #include "../../../src/tree/common_row_partitioner.h" #include "../../../src/tree/hist/expand_entry.h" // for MultiExpandEntry, CPUExpandEntry -#include "../collective/test_worker.h" // for TestDistributedGlobal +#include "../collective/test_worker.h" // for TestDistributedGlobal #include "../helpers.h" #include "test_column_split.h" // for TestColumnSplit #include "test_partitioner.h" #include "xgboost/data.h" +#include "xgboost/task.h" namespace xgboost::tree { namespace { @@ -246,4 +251,94 @@ INSTANTIATE_TEST_SUITE_P(ColumnSplit, TestHistColumnSplit, ::testing::ValuesIn([ } return params; }())); + +namespace { +void FillGradients(linalg::Matrix* gpair) { + auto h = gpair->HostView(); + for (std::size_t row = 0; row < h.Shape(0); ++row) { + for (std::size_t target = 0; target < h.Shape(1); ++target) { + h(row, target) = GradientPair{1.0f, 0.0f}; + } + } +} + +// Verify partitioner doesn't write past buffer end when doing +// update on small dataset after large one. +void TestPartitionerOverrun(bst_target_t n_targets) { + constexpr bst_idx_t kNBig = 1 << 16, kNSmall = 1024; + constexpr int kCols = 3; + + Context ctx; + ctx.InitAllowUnknown(Args{{"nthread", "1"}}); + + ObjInfo task{ObjInfo::kRegression, true, true}; + auto updater = + std::unique_ptr{TreeUpdater::Create("grow_quantile_histmaker", &ctx, &task)}; + + TrainParam param; + param.InitAllowUnknown(Args{{"max_depth", "1"}, + {"max_bin", "32"}, + {"lambda", "0"}, + {"gamma", "0"}, + {"min_child_weight", "0"}}); + updater->Configure(Args{}); + + auto const n_targets_size = static_cast(n_targets); + + auto dmat_large = + RandomDataGenerator{kNBig, kCols, 0.0f}.Seed(0).Batches(8).GenerateSparsePageDMatrix( + "part_resize_big_first", true); + + std::size_t shape_large[2]{dmat_large->Info().num_row_, n_targets_size}; + linalg::Matrix gpair_large(shape_large, ctx.Device()); + FillGradients(&gpair_large); + + RegTree tree_large{n_targets, static_cast(kCols)}; + std::vector trees_large{&tree_large}; + std::vector> position_large(1); + common::Span> pos_large{position_large.data(), 1}; + updater->Update(¶m, &gpair_large, dmat_large.get(), pos_large, trees_large); + + auto dmat_small = + RandomDataGenerator{kNSmall, kCols, 0.0f}.Seed(1).Batches(1).GenerateSparsePageDMatrix( + "part_resize_small_second", false); + + std::vector> position_small(1); + auto& pos = position_small.front(); + pos.Resize(kNBig); // Allocate large + pos.Resize(kNSmall); // Shrink logical size, capacity remains large + + auto& hv = pos.HostVector(); + std::size_t cap = hv.capacity(); + ASSERT_GE(cap, static_cast(kNBig)); + + std::size_t tail_elems = cap - hv.size(); + ASSERT_GT(tail_elems, 0u) << "Expected reserved tail storage"; + std::vector tail_before(tail_elems); + std::memcpy(tail_before.data(), hv.data() + hv.size(), tail_elems * sizeof(bst_node_t)); + + std::size_t shape_small[2]{dmat_small->Info().num_row_, n_targets_size}; + linalg::Matrix gpair_small(shape_small, ctx.Device()); + FillGradients(&gpair_small); + + RegTree tree_small{n_targets, static_cast(kCols)}; + std::vector trees_small{&tree_small}; + common::Span> pos_small{position_small.data(), 1}; + updater->Update(¶m, &gpair_small, dmat_small.get(), pos_small, trees_small); + + // Verify no buffer overrun: tail bytes should be unchanged + ASSERT_EQ(hv.capacity(), cap) << "Test precondition violated: capacity changed"; + std::vector tail_after(tail_elems); + std::memcpy(tail_after.data(), hv.data() + hv.size(), tail_elems * sizeof(bst_node_t)); + + EXPECT_EQ(tail_before, tail_after) + << "Buffer overrun detected: writes past kNSmall when updating small " + "single-batch DMatrix after large multi-batch one. " + "Likely stale partitioner writing to buffer."; +} +} // anonymous namespace + +TEST(QuantileHist, HistUpdaterPartitionerOverrun) { TestPartitionerOverrun(1); } + +TEST(QuantileHist, MultiTargetHistBuilderPartitionerOverrun) { TestPartitionerOverrun(3); } } // namespace xgboost::tree