Optimize transfer/compute overlap in out-of-core KMeans - #2538
Optimize transfer/compute overlap in out-of-core KMeans#2538viclafargue wants to merge 7 commits into
Conversation
viclafargue
added
improvement
|
/ok to test 4b8a5c5 |
|
/ok to test 9dd96bf |
| * @param[in] num_streams Number of CUDA streams in each device's pool | ||
| * @return cuvsError_t | ||
| */ | ||
| CUVS_EXPORT cuvsError_t cuvsMultiGpuResourcesSetStreamPool(cuvsResources_t res, |
There was a problem hiding this comment.
again -- I am not entirely opposed to this because I see how it makes things convenient. But why are we providing an API for this specifically? SNMG resources or raft resources with nccl comms are not special in the way stream pools should be set, right? In general I think such APIs to directly modify raft resources, should be in raft and not cuvs. We can add them here and then they have to stick around until we are allowed to break the ABI.
| auto data_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>( | ||
| handle, X.data_handle(), n_samples, n_features, device_buffer_samples, stream); | ||
| auto batch_mr = data_on_device ? raft::resource::get_workspace_resource_ref(handle) | ||
| : raft::resource::get_large_workspace_resource_ref(handle); |
There was a problem hiding this comment.
if the data is already device accessible, why are we even using the workspace?
There was a problem hiding this comment.
Furthermore, if the large workspace is managed, the transfer speed is slightly slower (I found 55 GB/s versus 48 GB/s). I agree that the batchsize is not bounded here (it can be as large as the dataset) but I would argue for falling back to the large workspace only at the breaking point where allocating from the regular workspace is not possible. That calculation can get complicated to account for whether or not weights are present, so I'll tag @achirkin for some ideas.
There was a problem hiding this comment.
Actually since this is unbounded (only bounded by dataset size) lets just stick to large_mr @viclafargue. Since compute will overlap, its not a big deal.
| }; | ||
|
|
||
| bool input_pipeline_started = false; | ||
| auto start_input_pipeline = [&] { |
There was a problem hiding this comment.
Can we make this function part of the batch loader class, so starting the pipeline explicitly becomes part of the batch loader? It makes it more readable. Overall there has been an increase in these lambdas all over the codebase (not the most readable I think)
| if (!batches_.empty()) { | ||
| RAFT_CUDA_TRY_NO_THROW(cudaStreamSynchronize(raft::resource::get_cuda_stream(*res_))); | ||
| } | ||
| RAFT_CUDA_TRY_NO_THROW(cudaStreamSynchronize(copy_stream_)); |
There was a problem hiding this comment.
Please use raft APIs for these. I tend to find these, but these should not need to be pointed out anymore. You can make rules for agents to use raft APIs wherever possible.
| * Unlike minClusterAndDistanceCompute, this path does not calculate cluster labels. | ||
| */ | ||
| template <typename DataT, typename IndexT> | ||
| void cluster_cost( |
There was a problem hiding this comment.
why do we have this new function? Its not used anywhere, right?
3 participants
Summary
This PR improves data-transfer and compute overlap for host-resident out-of-core KMeans using a cyclic two-buffer pipeline.
Implementation
Benchmark under similar configuration
10 GiB pinned-host FP32 dataset (10,485,760 × 256), 2,560 clusters, three iterations, ten 1 GiB out-of-core batches, and 131,072-sample assignment tiles.
mainProfile
Main branch :
This PR :
This PR (multi-GPU) :
