Federated plugin for histogram.

plugin/federated/CMakeLists.txt

@@ -37,7 +37,8 @@ target_link_libraries(federated_client INTERFACE federated_proto)
 
 # Rabit engine for Federated Learning.
 target_sources(
- objxgboost PRIVATE federated_tracker.cc federated_comm.cc federated_coll.cc federated_plugin.cc
+ objxgboost PRIVATE
+ federated_plugin.cc federated_hist.cc federated_tracker.cc federated_comm.cc federated_coll.cc
 )
 if(USE_CUDA)
  target_sources(objxgboost PRIVATE federated_comm.cu federated_coll.cu)

plugin/federated/federated_coll.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2023, XGBoost contributors
+ * Copyright 2023-2024, XGBoost contributors
  */
 #include "federated_coll.h"
 
@@ -8,11 +8,15 @@
 
 #include <algorithm> // for copy_n
 
-#include "../../src/collective/allgather.h"
-#include "../../src/common/common.h" // for AssertGPUSupport
 #include "federated_comm.h" // for FederatedComm
 #include "xgboost/collective/result.h" // for Result
 
+#if !defined(XGBOOST_USE_CUDA)
+
+#include "../../src/common/common.h" // for AssertGPUSupport
+
+#endif // !defined(XGBOOST_USE_CUDA)
+
 namespace xgboost::collective {
 namespace {
 [[nodiscard]] Result GetGRPCResult(std::string const &name, grpc::Status const &status) {

plugin/federated/federated_hist.cc

@@ -0,0 +1,159 @@
+/**
+ * Copyright 2024, XGBoost contributors
+ */
+#include "federated_hist.h"
+
+#include "../../src/collective/allgather.h" // for AllgatherV
+#include "../../src/collective/communicator-inl.h" // for GetRank
+#include "../../src/tree/hist/histogram.h" // for SubtractHistParallel, BuildSampleHistograms
+
+namespace xgboost::tree {
+namespace {
+// Copy the bins into a dense matrix.
+auto CopyBinsToDense(Context const *ctx, GHistIndexMatrix const &gidx) {
+ auto n_samples = gidx.Size();
+ auto n_features = gidx.Features();
+ std::vector<bst_bin_t> bins(n_samples * n_features);
+ auto bins_view = linalg::MakeTensorView(ctx, bins, n_samples, n_features);
+ common::ParallelFor(n_samples, ctx->Threads(), [&](auto ridx) {
+ for (bst_feature_t fidx = 0; fidx < n_features; fidx++) {
+ bins_view(ridx, fidx) = gidx.GetGindex(ridx, fidx);
+ }
+ });
+ return bins;
+}
+} // namespace
+
+template <bool any_missing>
+void FederataedHistPolicy::DoBuildLocalHistograms(
+ common::BlockedSpace2d const &space, GHistIndexMatrix const &gidx,
+ std::vector<bst_node_t> const &nodes_to_build,
+ common::RowSetCollection const &row_set_collection, common::Span<GradientPair const> gpair_h,
+ bool force_read_by_column, common::ParallelGHistBuilder *p_buffer) {
+ if (is_col_split_) {
+ // Copy the gidx information to the secure worker for encrypted histogram
+ // computation. This is copied as we don't want the plugin to handle the bin
+ // compression, which is quite internal of XGBoost.
+
+ // FIXME: this can be done during reset.
+ if (!is_gidx_initialized_) {
+ auto bins = CopyBinsToDense(ctx_, gidx);
+ auto cuts = gidx.Cuts().Ptrs();
+ plugin_->Reset(cuts, bins);
+ is_gidx_initialized_ = true;
+ }
+
+ // Share the row set collection without copy.
+ std::vector<std::uint64_t const *> ptrs(nodes_to_build.size());
+ std::vector<std::size_t> sizes(nodes_to_build.size());
+ std::vector<bst_node_t> nodes(nodes_to_build.size());
+ for (std::size_t i = 0; i < nodes_to_build.size(); ++i) {
+ auto nidx = nodes_to_build[i];
+ ptrs[i] = row_set_collection[nidx].begin();
+ sizes[i] = row_set_collection[nidx].Size();
+ nodes[i] = nidx;
+ }
+ hist_data_ = this->plugin_->BuildEncryptedHistVert(ptrs, sizes, nodes);
+ } else {
+ BuildSampleHistograms<any_missing>(this->ctx_->Threads(), space, gidx, nodes_to_build,
+ row_set_collection, gpair_h, force_read_by_column, p_buffer);
+ }
+}
+
+template void FederataedHistPolicy::DoBuildLocalHistograms<true>(
+ common::BlockedSpace2d const &space, GHistIndexMatrix const &gidx,
+ std::vector<bst_node_t> const &nodes_to_build,
+ common::RowSetCollection const &row_set_collection, common::Span<GradientPair const> gpair_h,
+ bool force_read_by_column, common::ParallelGHistBuilder *buffer);
+template void FederataedHistPolicy::DoBuildLocalHistograms<false>(
+ common::BlockedSpace2d const &space, GHistIndexMatrix const &gidx,
+ std::vector<bst_node_t> const &nodes_to_build,
+ common::RowSetCollection const &row_set_collection, common::Span<GradientPair const> gpair_h,
+ bool force_read_by_column, common::ParallelGHistBuilder *buffer);
+
+namespace {
+// The label owner needs to gather the result from all workers.
+void GatherWorkerHist(common::Span<double> hist_aggr, std::int32_t n_workers,
+ std::vector<bst_node_t> const &nodes_to_build, bst_bin_t n_total_bins,
+ tree::BoundedHistCollection *p_hist) {
+ bst_idx_t worker_size = hist_aggr.size() / n_workers;
+ bst_node_t n_nodes = nodes_to_build.size();
+ auto &hist = *p_hist;
+ // for each worker
+ for (auto widx = 0; widx < n_workers; ++widx) {
+ auto worker_hist = hist_aggr.subspan(widx * worker_size, worker_size);
+ // for each node
+ for (bst_node_t nidx_in_set = 0; nidx_in_set < n_nodes; ++nidx_in_set) {
+ auto hist_size = n_total_bins * kHist2F64; // Histogram size for one node.
+ auto hist_src = worker_hist.subspan(hist_size * nidx_in_set, hist_size);
+ auto hist_src_g = common::RestoreType<GradientPairPrecise>(hist_src);
+ auto hist_dst = hist[nodes_to_build[nidx_in_set]];
+ CHECK_EQ(hist_src_g.size(), hist_dst.size());
+ common::IncrementHist(hist_dst, hist_src_g, 0, hist_dst.size());
+ }
+ }
+}
+} // namespace
+
+void FederataedHistPolicy::DoSyncHistogram(common::BlockedSpace2d const &space,
+ std::vector<bst_node_t> const &nodes_to_build,
+ std::vector<bst_node_t> const &nodes_to_trick,
+ common::ParallelGHistBuilder *p_buffer,
+ tree::BoundedHistCollection *p_hist) {
+ auto n_total_bins = p_buffer->TotalBins();
+ CHECK(!nodes_to_build.empty());
+
+ auto &hist = *p_hist;
+ if (is_col_split_) {
+ // Under secure vertical mode, we perform allgather to get the global histogram. Note
+ // that only the label owner (rank == 0) needs the global histogram
+
+ // Perform AllGather
+ HostDeviceVector<std::int8_t> hist_entries;
+ std::vector<std::int64_t> recv_segments;
+ collective::SafeColl(
+ collective::AllgatherV(ctx_, linalg::MakeVec(hist_data_), &recv_segments, &hist_entries));
+
+ // Call the plugin here to get the resulting histogram. Histogram from all workers are
+ // gathered to the label owner.
+ common::Span<double> hist_aggr =
+ plugin_->SyncEncryptedHistVert(common::RestoreType<std::uint8_t>(hist_entries.HostSpan()));
+
+ // Update histogram for the label owner
+ if (collective::GetRank() == 0) {
+ std::int32_t n_workers = collective::GetWorldSize();
+ CHECK_EQ(hist_aggr.size() % n_workers, 0);
+ // Initialize histogram. For the normal case, this is done by the parallel hist
+ // buffer. We should try to unify the code paths.
+ for (auto nidx : nodes_to_build) {
+ auto hist_dst = hist[nidx];
+ std::fill_n(hist_dst.data(), hist_dst.size(), GradientPairPrecise{});
+ }
+ GatherWorkerHist(hist_aggr, n_workers, nodes_to_build, n_total_bins, p_hist);
+ }
+ } else {
+ common::ParallelFor2d(space, this->ctx_->Threads(), [&](std::size_t node, common::Range1d r) {
+ // Merging histograms from each thread.
+ p_buffer->ReduceHist(node, r.begin(), r.end());
+ });
+ // Encrtyped mode, we need to call the plugin to perform encryption and decryption.
+ auto first_nidx = nodes_to_build.front();
+ std::size_t n = n_total_bins * nodes_to_build.size() * kHist2F64;
+ auto src_hist = common::Span{reinterpret_cast<double const *>(hist[first_nidx].data()), n};
+ auto hist_buf = plugin_->BuildEncryptedHistHori(src_hist);
+
+ // allgather
+ HostDeviceVector<std::int8_t> hist_entries;
+ std::vector<std::int64_t> recv_segments;
+ auto rc =
+ collective::AllgatherV(ctx_, linalg::MakeVec(hist_buf), &recv_segments, &hist_entries);
+ collective::SafeColl(rc);
+
+ auto hist_aggr =
+ plugin_->SyncEncryptedHistHori(common::RestoreType<std::uint8_t>(hist_entries.HostSpan()));
+ // Assign the aggregated histogram back to the local histogram
+ auto hist_dst = reinterpret_cast<double *>(hist[first_nidx].data());
+ std::copy_n(hist_aggr.data(), hist_aggr.size(), hist_dst);
+ }
+}
+} // namespace xgboost::tree

plugin/federated/federated_hist.h

@@ -0,0 +1,58 @@
+/**
+ * Copyright 2024, XGBoost contributors
+ */
+#pragma once
+#include <cstdint> // for int32_t
+#include <vector> // for vector
+
+#include "../../src/collective/comm_group.h" // for GlobalCommGroup
+#include "../../src/common/hist_util.h" // for ParallelGHistBuilder
+#include "../../src/common/row_set.h" // for RowSetCollection
+#include "../../src/common/threading_utils.h" // for BlockedSpace2d
+#include "../../src/data/gradient_index.h" // for GHistIndexMatrix
+#include "../../src/tree/hist/hist_cache.h" // for BoundedHistCollection
+#include "federated_comm.h" // for FederatedComm
+#include "xgboost/base.h" // for GradientPair
+#include "xgboost/context.h" // for Context
+#include "xgboost/span.h" // for Span
+#include "xgboost/tree_model.h" // for RegTree
+
+namespace xgboost::tree {
+/**
+ * @brief Federated histogram build policy
+ */
+class FederataedHistPolicy {
+ // fixme: duplicated code
+ bool is_col_split_{false};
+ bool is_distributed_{false};
+ decltype(std::declval<collective::FederatedComm>().EncryptionPlugin()) plugin_;
+ xgboost::common::Span<std::uint8_t> hist_data_;
+ // Only initialize the aggregation context once
+ bool is_gidx_initialized_{false};
+ Context const* ctx_;
+
+ public:
+ void DoReset(Context const *ctx, bool is_distributed, bool is_col_split) {
+ this->is_distributed_ = is_distributed;
+ CHECK(is_distributed);
+ this->ctx_ = ctx;
+ this->is_col_split_ = is_col_split;
+ auto const &comm = collective::GlobalCommGroup()->Ctx(ctx, DeviceOrd::CPU());
+ auto const &fed = dynamic_cast<collective::FederatedComm const &>(comm);
+ plugin_ = fed.EncryptionPlugin();
+ CHECK(is_distributed_) << "Unreachable. Single node training can not be federated.";
+ }
+
+ template <bool any_missing>
+ void DoBuildLocalHistograms(common::BlockedSpace2d const &space, GHistIndexMatrix const &gidx,
+ std::vector<bst_node_t> const &nodes_to_build,
+ common::RowSetCollection const &row_set_collection,
+ common::Span<GradientPair const> gpair_h, bool force_read_by_column,
+ common::ParallelGHistBuilder *p_buffer);
+
+ void DoSyncHistogram(common::BlockedSpace2d const &space,
+ std::vector<bst_node_t> const &nodes_to_build,
+ std::vector<bst_node_t> const &nodes_to_trick,
+ common::ParallelGHistBuilder *p_buffer, tree::BoundedHistCollection *p_hist);
+};
+} // namespace xgboost::tree

plugin/federated/federated_plugin.h

@@ -12,6 +12,12 @@
  * - Build histogram for vertical federated learning.
  * - Build histogram for horizontal federated learning.
  *
+ * Since we don't require the plugin to have network capability, the synchronization is
+ * performed in XGBoost. As a result, the build procedure is divided into four steps,
+ * first we need to build a local histogram, then encrypt it with the plugin. Afterward,
+ * the control returns to XBGoost, which is responsible for synchronization. Lastly, the
+ * plugin will recieve the synchronization result and return the decrypted histogram.
+ *
  * See below function prototypes for details. All prototypes are for C functions that are
  * suitable for `dlopen`.
  */

src/data/gradient_index.h

@@ -7,8 +7,8 @@
 
 #include <algorithm> // for min
 #include <atomic> // for atomic
-#include <cinttypes> // for uint32_t
 #include <cstddef> // for size_t
+#include <cstdint> // for uint32_t
 #include <memory> // for make_unique
 #include <vector>
 

src/tree/hist/histogram.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2023 by XGBoost Contributors
+ * Copyright 2023-2024, XGBoost Contributors
  */
 #include "histogram.h"
 
@@ -10,7 +10,7 @@
 
 #include "../../common/transform_iterator.h" // for MakeIndexTransformIter
 #include "expand_entry.h" // for MultiExpandEntry, CPUExpandEntry
-#include "xgboost/logging.h" // for CHECK_NE
+#include "xgboost/logging.h" // for CHECK_EQ
 #include "xgboost/span.h" // for Span
 #include "xgboost/tree_model.h" // for RegTree