[KMeans] Configure Workspace for Large Batch Sizes - #2433
Conversation
tarang-jain
added
bug
| if (weight_ptr != nullptr) { | ||
| batch_staging_bytes += static_cast<size_t>(device_buffer_samples) * sizeof(DataT); | ||
| } | ||
| if (batch_staging_bytes > raft::resource::get_workspace_free_bytes(handle)) { |
There was a problem hiding this comment.
Why not just always use the large workspace? I think doing this conditionally creates an additional challenge for user debugging that we could avoid if we just use the same workspace resources all the time. Will let @achirkin comment here too.
There was a problem hiding this comment.
Thats a good point, especially since we do expect batches to be quite large ( > 20 GB or so per batch).
There was a problem hiding this comment.
The user may set up the large workspace to use a slower memory than the normal workspace (e.g. managed memory vs device pool - a setup we recommend and also set in benchmarks).
Therefore, please view this as an optimization: if access to arrays allocated via this resource is the bottleneck, we should keep it; otherwise, it's ok to use the large workspace by default.
In this case, we're talking about batching, so my understanding is having small enough batches is a normal behavior, whereas the switch to the large workspace is an edge case to make the algorithm not fail if there's not enough memory.
There was a problem hiding this comment.
I updated the PR desc to show the throughput with a managed and pooled workspace respectively. I am in favor of keeping this PR as is: use the large MR only when the required bytes exceed the workspace.
…o mem-resource
4 participants
RAFT limits the regular workspaces to one-fourth the available GPU memory. If a batch size larger than that is used, fall back to the large workspace to avoid OOM.
Merge this after #2484
Experiment to check managed vs pool MR to check if this selectively choosing the workspace makes sense. The idea is that the workspace MR is mostly a pool MR and the large MR is a managed MR. We show the difference in throughput for data transfer from host to pool and managed workspace.
H100 + Intel(R) Xeon(R) Platinum 8480CL