Remove unnecessary fetch operations in external memory. (#10342)

This commit is contained in:
Jiaming Yuan
2024-05-31 13:16:40 +08:00
committed by GitHub
parent c2e3d4f3cd
commit d2d01d977a
12 changed files with 180 additions and 71 deletions

View File

@@ -1,8 +1,7 @@
/**
* Copyright 2019-2023, XGBoost contributors
* Copyright 2019-2024, XGBoost contributors
*/
#include <memory>
#include <utility>
#include "ellpack_page.cuh"
#include "ellpack_page.h" // for EllpackPage

View File

@@ -6,7 +6,6 @@
#include <any> // for any, any_cast
#include <memory>
#include <string>
#include <type_traits> // for invoke_result_t
#include <utility>

View File

@@ -56,10 +56,10 @@ SparsePageDMatrix::SparsePageDMatrix(DataIterHandle iter_handle, DMatrixHandle p
auto iter = DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext>{
iter_, reset_, next_};
uint32_t n_batches = 0;
size_t n_features = 0;
size_t n_samples = 0;
size_t nnz = 0;
std::uint32_t n_batches = 0;
bst_feature_t n_features = 0;
bst_idx_t n_samples = 0;
bst_idx_t nnz = 0;
auto num_rows = [&]() {
bool type_error {false};
@@ -72,7 +72,7 @@ SparsePageDMatrix::SparsePageDMatrix(DataIterHandle iter_handle, DMatrixHandle p
};
auto num_cols = [&]() {
bool type_error {false};
size_t n_features = HostAdapterDispatch(
bst_feature_t n_features = HostAdapterDispatch(
proxy, [](auto const &value) { return value.NumCols(); }, &type_error);
if (type_error) {
n_features = detail::NFeaturesDevice(proxy);
@@ -121,10 +121,9 @@ void SparsePageDMatrix::InitializeSparsePage(Context const *ctx) {
this->n_batches_, cache_info_.at(id));
}
BatchSet<SparsePage> SparsePageDMatrix::GetRowBatchesImpl(Context const* ctx) {
BatchSet<SparsePage> SparsePageDMatrix::GetRowBatchesImpl(Context const *ctx) {
this->InitializeSparsePage(ctx);
auto begin_iter = BatchIterator<SparsePage>(sparse_page_source_);
return BatchSet<SparsePage>(BatchIterator<SparsePage>(begin_iter));
return BatchSet{BatchIterator<SparsePage>{this->sparse_page_source_}};
}
BatchSet<SparsePage> SparsePageDMatrix::GetRowBatches() {
@@ -143,8 +142,7 @@ BatchSet<CSCPage> SparsePageDMatrix::GetColumnBatches(Context const *ctx) {
} else {
column_source_->Reset();
}
auto begin_iter = BatchIterator<CSCPage>(column_source_);
return BatchSet<CSCPage>(BatchIterator<CSCPage>(begin_iter));
return BatchSet{BatchIterator<CSCPage>{this->column_source_}};
}
BatchSet<SortedCSCPage> SparsePageDMatrix::GetSortedColumnBatches(Context const *ctx) {
@@ -158,8 +156,7 @@ BatchSet<SortedCSCPage> SparsePageDMatrix::GetSortedColumnBatches(Context const
} else {
sorted_column_source_->Reset();
}
auto begin_iter = BatchIterator<SortedCSCPage>(sorted_column_source_);
return BatchSet<SortedCSCPage>(BatchIterator<SortedCSCPage>(begin_iter));
return BatchSet{BatchIterator<SortedCSCPage>{this->sorted_column_source_}};
}
BatchSet<GHistIndexMatrix> SparsePageDMatrix::GetGradientIndex(Context const *ctx,
@@ -169,8 +166,8 @@ BatchSet<GHistIndexMatrix> SparsePageDMatrix::GetGradientIndex(Context const *ct
}
detail::CheckEmpty(batch_param_, param);
auto id = MakeCache(this, ".gradient_index.page", cache_prefix_, &cache_info_);
this->InitializeSparsePage(ctx);
if (!cache_info_.at(id)->written || detail::RegenGHist(batch_param_, param)) {
this->InitializeSparsePage(ctx);
cache_info_.erase(id);
MakeCache(this, ".gradient_index.page", cache_prefix_, &cache_info_);
LOG(INFO) << "Generating new Gradient Index.";
@@ -190,15 +187,13 @@ BatchSet<GHistIndexMatrix> SparsePageDMatrix::GetGradientIndex(Context const *ct
CHECK(ghist_index_source_);
ghist_index_source_->Reset();
}
auto begin_iter = BatchIterator<GHistIndexMatrix>(ghist_index_source_);
return BatchSet<GHistIndexMatrix>(BatchIterator<GHistIndexMatrix>(begin_iter));
return BatchSet{BatchIterator<GHistIndexMatrix>{this->ghist_index_source_}};
}
#if !defined(XGBOOST_USE_CUDA)
BatchSet<EllpackPage> SparsePageDMatrix::GetEllpackBatches(Context const *, const BatchParam &) {
common::AssertGPUSupport();
auto begin_iter = BatchIterator<EllpackPage>(ellpack_page_source_);
return BatchSet<EllpackPage>(BatchIterator<EllpackPage>(begin_iter));
return BatchSet{BatchIterator<EllpackPage>{this->ellpack_page_source_}};
}
#endif // !defined(XGBOOST_USE_CUDA)
} // namespace xgboost::data

View File

@@ -1,5 +1,5 @@
/**
* Copyright 2021-2023 by XGBoost contributors
* Copyright 2021-2024, XGBoost contributors
*/
#include <memory> // for unique_ptr
@@ -21,8 +21,8 @@ BatchSet<EllpackPage> SparsePageDMatrix::GetEllpackBatches(Context const* ctx,
detail::CheckEmpty(batch_param_, param);
auto id = MakeCache(this, ".ellpack.page", cache_prefix_, &cache_info_);
size_t row_stride = 0;
this->InitializeSparsePage(ctx);
if (!cache_info_.at(id)->written || detail::RegenGHist(batch_param_, param)) {
this->InitializeSparsePage(ctx);
// reinitialize the cache
cache_info_.erase(id);
MakeCache(this, ".ellpack.page", cache_prefix_, &cache_info_);
@@ -52,7 +52,6 @@ BatchSet<EllpackPage> SparsePageDMatrix::GetEllpackBatches(Context const* ctx,
ellpack_page_source_->Reset();
}
auto begin_iter = BatchIterator<EllpackPage>(ellpack_page_source_);
return BatchSet<EllpackPage>(BatchIterator<EllpackPage>(begin_iter));
return BatchSet{BatchIterator<EllpackPage>{this->ellpack_page_source_}};
}
} // namespace xgboost::data

View File

@@ -1,5 +1,5 @@
/**
* Copyright 2015-2023, XGBoost Contributors
* Copyright 2015-2024, XGBoost Contributors
* \file sparse_page_dmatrix.h
* \brief External-memory version of DMatrix.
* \author Tianqi Chen
@@ -7,12 +7,10 @@
#ifndef XGBOOST_DATA_SPARSE_PAGE_DMATRIX_H_
#define XGBOOST_DATA_SPARSE_PAGE_DMATRIX_H_
#include <algorithm>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "ellpack_page_source.h"
#include "gradient_index_page_source.h"
@@ -22,7 +20,7 @@
namespace xgboost::data {
/**
* \brief DMatrix used for external memory.
* @brief DMatrix used for external memory.
*
* The external memory is created for controlling memory usage by splitting up data into
* multiple batches. However that doesn't mean we will actually process exactly 1 batch
@@ -51,8 +49,13 @@ namespace xgboost::data {
* want to change the generated page like Ellpack, pass parameter into `GetBatches` to
* re-generate them instead of trying to modify the pages in-place.
*
* A possible optimization is dropping the sparse page once dependent pages like ellpack
* are constructed and cached.
* The overall chain of responsibility of external memory DMatrix:
*
* User defined iterator (in Python/C/R) -> Proxy DMatrix -> Sparse page Source ->
* Other sources (Like Ellpack) -> Sparse Page DMatrix -> Caller
*
* A possible optimization is skipping the sparse page source for `hist` based algorithms
* similar to the Quantile DMatrix.
*/
class SparsePageDMatrix : public DMatrix {
MetaInfo info_;
@@ -67,7 +70,7 @@ class SparsePageDMatrix : public DMatrix {
float missing_;
Context fmat_ctx_;
std::string cache_prefix_;
uint32_t n_batches_{0};
std::uint32_t n_batches_{0};
// sparse page is the source to other page types, we make a special member function.
void InitializeSparsePage(Context const *ctx);
// Non-virtual version that can be used in constructor
@@ -93,11 +96,11 @@ class SparsePageDMatrix : public DMatrix {
}
}
MetaInfo &Info() override;
const MetaInfo &Info() const override;
Context const *Ctx() const override { return &fmat_ctx_; }
[[nodiscard]] MetaInfo &Info() override;
[[nodiscard]] const MetaInfo &Info() const override;
[[nodiscard]] Context const *Ctx() const override { return &fmat_ctx_; }
// The only DMatrix implementation that returns false.
bool SingleColBlock() const override { return false; }
[[nodiscard]] bool SingleColBlock() const override { return false; }
DMatrix *Slice(common::Span<int32_t const>) override {
LOG(FATAL) << "Slicing DMatrix is not supported for external memory.";
return nullptr;
@@ -107,6 +110,20 @@ class SparsePageDMatrix : public DMatrix {
return nullptr;
}
[[nodiscard]] bool EllpackExists() const override {
return static_cast<bool>(ellpack_page_source_);
}
[[nodiscard]] bool GHistIndexExists() const override {
return static_cast<bool>(ghist_index_source_);
}
[[nodiscard]] bool SparsePageExists() const override {
return static_cast<bool>(sparse_page_source_);
}
// For testing, getter for the number of fetches for sparse page source.
[[nodiscard]] auto SparsePageFetchCount() const {
return this->sparse_page_source_->FetchCount();
}
private:
BatchSet<SparsePage> GetRowBatches() override;
BatchSet<CSCPage> GetColumnBatches(Context const *ctx) override;
@@ -118,24 +135,24 @@ class SparsePageDMatrix : public DMatrix {
return BatchSet<ExtSparsePage>(BatchIterator<ExtSparsePage>(nullptr));
}
private:
// source data pointers.
std::shared_ptr<SparsePageSource> sparse_page_source_;
std::shared_ptr<EllpackPageSource> ellpack_page_source_;
std::shared_ptr<CSCPageSource> column_source_;
std::shared_ptr<SortedCSCPageSource> sorted_column_source_;
std::shared_ptr<GradientIndexPageSource> ghist_index_source_;
bool EllpackExists() const override { return static_cast<bool>(ellpack_page_source_); }
bool GHistIndexExists() const override { return static_cast<bool>(ghist_index_source_); }
bool SparsePageExists() const override { return static_cast<bool>(sparse_page_source_); }
};
inline std::string MakeId(std::string prefix, SparsePageDMatrix *ptr) {
[[nodiscard]] inline std::string MakeId(std::string prefix, SparsePageDMatrix *ptr) {
std::stringstream ss;
ss << ptr;
return prefix + "-" + ss.str();
}
/**
* @brief Make cache if it doesn't exist yet.
*/
inline std::string MakeCache(SparsePageDMatrix *ptr, std::string format, std::string prefix,
std::map<std::string, std::shared_ptr<Cache>> *out) {
auto &cache_info = *out;

View File

@@ -8,7 +8,7 @@
#include <algorithm> // for min
#include <atomic> // for atomic
#include <cstdio> // for remove
#include <future> // for async
#include <future> // for future
#include <memory> // for unique_ptr
#include <mutex> // for mutex
#include <numeric> // for partial_sum
@@ -55,7 +55,7 @@ struct Cache {
offset.push_back(0);
}
static std::string ShardName(std::string name, std::string format) {
[[nodiscard]] static std::string ShardName(std::string name, std::string format) {
CHECK_EQ(format.front(), '.');
return name + format;
}
@@ -174,7 +174,7 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
ExceHandler exce_;
common::Monitor monitor_;
bool ReadCache() {
[[nodiscard]] bool ReadCache() {
CHECK(!at_end_);
if (!cache_info_->written) {
return false;
@@ -216,7 +216,6 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
return page;
});
}
CHECK_EQ(std::count_if(ring_->cbegin(), ring_->cend(), [](auto const& f) { return f.valid(); }),
n_prefetch_batches)
<< "Sparse DMatrix assumes forward iteration.";
@@ -279,9 +278,9 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
}
}
[[nodiscard]] uint32_t Iter() const { return count_; }
[[nodiscard]] std::uint32_t Iter() const { return count_; }
const S &operator*() const override {
[[nodiscard]] S const& operator*() const override {
CHECK(page_);
return *page_;
}
@@ -311,22 +310,29 @@ inline void DevicePush(DMatrixProxy*, float, SparsePage*) { common::AssertGPUSup
#endif
class SparsePageSource : public SparsePageSourceImpl<SparsePage> {
// This is the source from the user.
// This is the source iterator from the user.
DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext> iter_;
DMatrixProxy* proxy_;
std::size_t base_row_id_{0};
bst_idx_t fetch_cnt_{0}; // Used for sanity check.
void Fetch() final {
fetch_cnt_++;
page_ = std::make_shared<SparsePage>();
// The first round of reading, this is responsible for initialization.
if (!this->ReadCache()) {
bool type_error { false };
bool type_error{false};
CHECK(proxy_);
HostAdapterDispatch(proxy_, [&](auto const &adapter_batch) {
page_->Push(adapter_batch, this->missing_, this->nthreads_);
}, &type_error);
HostAdapterDispatch(
proxy_,
[&](auto const& adapter_batch) {
page_->Push(adapter_batch, this->missing_, this->nthreads_);
},
&type_error);
if (type_error) {
DevicePush(proxy_, missing_, page_.get());
}
page_->SetBaseRowId(base_row_id_);
base_row_id_ += page_->Size();
n_batches_++;
@@ -351,11 +357,13 @@ class SparsePageSource : public SparsePageSourceImpl<SparsePage> {
SparsePageSource& operator++() final {
TryLockGuard guard{single_threaded_};
count_++;
if (cache_info_->written) {
at_end_ = (count_ == n_batches_);
} else {
at_end_ = !iter_.Next();
}
CHECK_LE(count_, n_batches_);
if (at_end_) {
CHECK_EQ(cache_info_->offset.size(), n_batches_ + 1);
@@ -381,6 +389,8 @@ class SparsePageSource : public SparsePageSourceImpl<SparsePage> {
TryLockGuard guard{single_threaded_};
base_row_id_ = 0;
}
[[nodiscard]] auto FetchCount() const { return fetch_cnt_; }
};
// A mixin for advancing the iterator.
@@ -394,11 +404,11 @@ class PageSourceIncMixIn : public SparsePageSourceImpl<S> {
bool sync_{true};
public:
PageSourceIncMixIn(float missing, int nthreads, bst_feature_t n_features, uint32_t n_batches,
PageSourceIncMixIn(float missing, int nthreads, bst_feature_t n_features, std::uint32_t n_batches,
std::shared_ptr<Cache> cache, bool sync)
: Super::SparsePageSourceImpl{missing, nthreads, n_features, n_batches, cache}, sync_{sync} {}
PageSourceIncMixIn& operator++() final {
[[nodiscard]] PageSourceIncMixIn& operator++() final {
TryLockGuard guard{this->single_threaded_};
if (sync_) {
++(*source_);
@@ -422,6 +432,13 @@ class PageSourceIncMixIn : public SparsePageSourceImpl<S> {
}
return *this;
}
void Reset() final {
if (sync_) {
this->source_->Reset();
}
Super::Reset();
}
};
class CSCPageSource : public PageSourceIncMixIn<CSCPage> {

View File

@@ -7,7 +7,6 @@
#include <algorithm> // for max, copy, transform
#include <cstddef> // for size_t
#include <cstdint> // for uint32_t, int32_t
#include <exception> // for exception
#include <memory> // for allocator, unique_ptr, make_unique, shared_ptr
#include <ostream> // for operator<<, basic_ostream, char_traits
#include <utility> // for move
@@ -20,7 +19,6 @@
#include "../common/random.h" // for ColumnSampler
#include "../common/threading_utils.h" // for ParallelFor
#include "../common/timer.h" // for Monitor
#include "../common/transform_iterator.h" // for IndexTransformIter
#include "../data/gradient_index.h" // for GHistIndexMatrix
#include "common_row_partitioner.h" // for CommonRowPartitioner
#include "dmlc/registry.h" // for DMLC_REGISTRY_FILE_TAG