Exposing opt-in double buffering for MG batched KMeans#2323
Conversation
| int64_t init_size; | ||
|
|
||
| /** | ||
| * Whether host-resident multi-GPU KMeans should prefetch the next streaming |
There was a problem hiding this comment.
| * Whether host-resident multi-GPU KMeans should prefetch the next streaming | |
| * Use prefetching to data latencies in multi-GPU batched k-means when trained on host-resident input data``` |
| kmeans_params.batch_centroids = params.batch_centroids; | ||
| kmeans_params.init_size = params.init_size; | ||
| kmeans_params.streaming_batch_size = params.streaming_batch_size; | ||
| kmeans_params.streaming_batch_prefetch = params.streaming_batch_prefetch; |
There was a problem hiding this comment.
I really don't like the use of the term "streaming" here. It carries altogether too much baggage with it. Can we please just remove it here? I think we need to consider removing it from "streaming_batch_size" as well. Streaming to me means a k-means that runs through a streaming mechanism, meaning it can only go forward as an iterator and can restart from where it left off when more data is received. There are variants of k-means that satisfy this, but the batched variant we've implemented does not.
There was a problem hiding this comment.
The problem with batch_size is that we have two other parameters -- batch_centroids and batch_samples. So batch_size can be confusing. Perhaps we should reconsider the names of those two parameters then? Or maybe think of some other prefix for batch_size?
There was a problem hiding this comment.
What are the meanings of batch samples vs batch size?
There was a problem hiding this comment.
the tile sizes for the distance computation (local reductions are done on this tile) -- i.e. batch_samples * batch_centroids distances are computed at once. So I was saying that we are already doing what the first part of the flash-kmeans paper mentions (at least that was my understanding)
There was a problem hiding this comment.
/**
* batch_samples and batch_centroids are used to tile 1NN computation which is
* useful to optimize/control the memory footprint
* Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0
* then don't tile the centroids
*
* NB: These parameters are unrelated to streaming_batch_size, which controls how many
* samples to transfer from host to device per batch when processing out-of-core
* data.
*/
int batch_samples = 1 << 15;
/**
* if 0 then batch_centroids = n_clusters
*/
int batch_centroids = 0;
| int64_t streaming_batch_size = 0; | ||
|
|
||
| /** | ||
| * Whether host-resident multi-GPU KMeans should prefetch the next streaming |
There was a problem hiding this comment.
Please reword to not use words like "whether". The user can imply that, we don't need to say it. A more concise and clear way to state this would be to jump right into what the feasture accomplishes for the user:
"Hide latencies for host-resident inputs in multi-gpu K-means by prefetching the next batch while the current batch is processing. This is done asychronously using different CUDA streams, enabling overlap of host to device (let's not use H2D in public facing documentation, because we can't expect users to know what that means) transfer with the computation by allocating a second device batch buffer on each NCCL rank. In practice this can moderate speedups of up to "XX" over processing batches in serial".
| 1, | ||
| cuvs::cluster::kmeans::params::Array, | ||
| 20, | ||
| true}, |
There was a problem hiding this comment.
Should definitely improve the coverage here- what ahppens when there's only 1 batch? what happens when the batch size is not a multiple of the total number of vectors?
|
|
||
| Default: 0 (process all data at once). | ||
| streaming_batch_prefetch : bool | ||
| Whether host-resident multi-GPU KMeans should use a second device |
There was a problem hiding this comment.
Please follow suggestions above about the wording of descriptions.
Partially answers #2292