[BUG] Balanced KMeans triggers limiting_resource_adaptor.hpp:152: Exceeded memory limit errors for large datasets
Nobody has claimed this yet.
- Dominant language
- Cuda
- Stars
- 854
- Forks
- 236
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 62
Description
Describe the bug
The Balanced KMeans implementation uses a RAFT get_workspace() resource that is used to allocated arrays on the order of minibatch_size within the build_fine_clusters() function (passed as device_memory) which then allocates mc_trainset_buf [mesocluster_size_max x dim] which is on the order of dataset size / n_clusters, i.e., orders of magnitude larger than minibatch size which is ~1GB. This will trigger a limiting_resource_adaptor.hpp:152: Exceeded memory limit exception, because the default allocation limit is set to total device memory / 4.
To avoid this problem for large datasets (~ device memory size), the user must increase the number of (mesoscale) clusters. However, while increasing the number of clusters commensurate with the dataset size is generally advisable, I believe that we should not artificially limit the allocation size when the user explicitly uses managed memory. This means even if we do not generally remove the resource limiter on the workspace resource, we should at least remove it specifically for the mc_trainset_buf allocation since there is no expectation that it should be on the order of minibatch size which is otherwise used to estimate the expected workspace resource needs.
Steps/Code to reproduce bug
The issue can be reproduced with the test script posted in this issue.
Expected behavior
I would expect to not run into a device resource allocator before device memory is sufficiently exhausted and I would expect to not encounter any OOM or resource limiter issues when using a managed memory allocator.
Environment details (please complete the following information):
- Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)]
- Method of RAFT install: [conda, Docker, or from source]
- If method of install is [Docker], provide
docker pull&docker runcommands used
- If method of install is [Docker], provide
Additional context
Add any other context about the problem here.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with cpp/src/cluster/detail/kmeans_balanced.cuh, especially build_fine_clusters() around the linked line, and trace how get_workspace() is passed as device_memory before mc_trainset_buf is allocated. Use the reproduction script from cuML issue 6204 to observe the limiter failure. Done means large managed-memory runs no longer hit the premature resource limit or an avoidable OOM.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100