Skip to content

Ensure MG to have the same number of allreduce calls in mean_stddev for sparse matrix to avoid hanging #6141

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 4 commits into from
Dec 6, 2024
Merged

Ensure MG to have the same number of allreduce calls in mean_stddev for sparse matrix to avoid hanging #6141

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a500a95
ensure all GPUs having same allreduces to avoid hanging
lijinf2 Nov 22, 2024
9ff5320
avoid skipping mean substraction and square root
lijinf2 Dec 3, 2024
ab0cab2
add test case of shifted scaled dataset for sparse standardization
lijinf2 Dec 4, 2024
90d7df5
refine PR
lijinf2 Dec 4, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 24 additions & 17 deletions cpp/src/glm/qn/mg/standardization.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,16 @@ void mean(const raft::handle_t& handle,
auto stream = handle.get_stream();
auto& comm = handle.get_comms();

if (X.nnz == 0) {
SimpleVec<T> meanVec(mean_vector, D);
meanVec.fill(0., stream);

// call allreduces on zeroes to sync with other GPUs to avoid hanging
comm.allreduce(mean_vector, mean_vector, D, raft::comms::op_t::SUM, stream);
comm.sync_stream(stream);
return;
}

int chunk_size = 500000; // split matrix by rows for better numeric precision
rmm::device_uvector<I> buff_row_ids(chunk_size + 1, stream);

Expand Down Expand Up @@ -200,27 +210,24 @@ void mean_stddev(const raft::handle_t& handle,
auto stream = handle.get_stream();
int D = X.n;

if (X.nnz == 0) {
SimpleVec<T> meanVec(mean_vector, D);
meanVec.fill(0., stream);

SimpleVec<T> stddevVec(stddev_vector, D);
stddevVec.fill(0., stream);
return;
}

mean(handle, X, n_samples, mean_vector);

// calculate stdev.S
rmm::device_uvector<T> X_values_squared(X.nnz, stream);
raft::copy(X_values_squared.data(), X.values, X.nnz, stream);
auto square_op = [] __device__(const T a) { return a * a; };
raft::linalg::unaryOp(X_values_squared.data(), X_values_squared.data(), X.nnz, square_op, stream);

auto X_squared =
SimpleSparseMat<T, I>(X_values_squared.data(), X.cols, X.row_ids, X.nnz, X.m, X.n);

mean(handle, X_squared, n_samples, stddev_vector);
if (X.nnz == 0) {
mean(handle, X, n_samples, stddev_vector);
} else {
rmm::device_uvector<T> X_values_squared(X.nnz, stream);
raft::copy(X_values_squared.data(), X.values, X.nnz, stream);
auto square_op = [] __device__(const T a) { return a * a; };
raft::linalg::unaryOp(
X_values_squared.data(), X_values_squared.data(), X.nnz, square_op, stream);

auto X_squared =
SimpleSparseMat<T, I>(X_values_squared.data(), X.cols, X.row_ids, X.nnz, X.m, X.n);

mean(handle, X_squared, n_samples, stddev_vector);
}

T weight = n_samples / T(n_samples - 1);
auto submean_no_neg_op = [weight] __device__(const T a, const T b) -> T {
Expand Down
Loading