[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.
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.