Reduce thread contention in column split tests. (#10658)

---------

Co-authored-by: Philip Hyunsu Cho <chohyu01@cs.washington.edu>
This commit is contained in:
Jiaming Yuan
2024-08-01 18:36:46 +08:00
committed by GitHub
parent 778751a1bb
commit 77c844cef7
8 changed files with 101 additions and 61 deletions

View File

@@ -1,12 +1,9 @@
/**
* Copyright 2017-2023 by XGBoost contributors
* Copyright 2017-2024, XGBoost contributors
*/
#include <gtest/gtest.h>
#include <xgboost/predictor.h>
#include <cstdint>
#include <thread>
#include "../../../src/collective/communicator-inl.h"
#include "../../../src/data/adapter.h"
#include "../../../src/data/proxy_dmatrix.h"

View File

@@ -511,15 +511,20 @@ void VerifyIterationRangeColumnSplit(bool use_gpu, Json const &ranged_model,
if (use_gpu) {
ctx = MakeCUDACtx(common::AllVisibleGPUs() == 1 ? 0 : rank);
}
auto n_threads = collective::GetWorkerLocalThreads(world_size);
ctx.UpdateAllowUnknown(
Args{{"nthread", std::to_string(n_threads)}, {"device", ctx.DeviceName()}});
auto dmat = RandomDataGenerator(rows, cols, 0).GenerateDMatrix(true, true, classes);
std::shared_ptr<DMatrix> Xy{dmat->SliceCol(world_size, rank)};
std::unique_ptr<Learner> learner{Learner::Create({Xy})};
learner->SetParam("device", ctx.DeviceName());
auto args = Args{{"device", ctx.DeviceName()}, {"nthread", std::to_string(ctx.Threads())}};
learner->SetParams(args);
learner->LoadModel(ranged_model);
std::unique_ptr<Learner> sliced{Learner::Create({Xy})};
sliced->SetParam("device", ctx.DeviceName());
sliced->SetParams(args);
sliced->LoadModel(sliced_model);
HostDeviceVector<float> out_predt_sliced;