[Outdated] Scaling workspace resources

[achirkin]

Brief

Add another workspace memory resource that does not have the explicit memory limit. That is, after the change we have the following:

1. rmm::mr::get_current_device_resource() is default for all allocations, as before. It is used for the allocations with unlimited lifetime, e.g. returned to the user.
2. raft::get_workspace_resource() is for temporary allocations and forced to have fixed size, as before. However, it becomes smaller and should be used only for allocations, which do not scale with problem size. It defaults to a thin layer on top of the current_device_resource.
3. raft::get_large_workspace_resource() (new) is for temporary allocations, which can scale with the problem size. Unlike workspace_resource, its size is not fixed. By default, it points to the current_device_resource, but the user can set it to something backed by the host memory (e.g. managed memory) to avoid OOM exceptions when there's not enough device memory left.

Problem

We have a list of issues/preference/requirements, some of which contradict others

1. We rely on RMM to handle all allocations and we often use rmm::mr::pool_memory_resource for performance reasons (to avoid lots of cudaMalloc calls in the loops)
2. Historically, we've used managed memory allocators as a workaround to avoid OOM errors or improve speed (by increasing batch sizes).
3. However, the design goal is to avoid setting allocators on our own and to give the full control to the user (hence the workaround in 2 was removed).
4. We introduced the workspace resource earlier to allow querying the available memory reliably and maximize the batch sizes accordingly (see also issue #1310). Without this, some of our batched algorithms either fail with OOM or severely underperform due to small batch sizes.
5. However, we cannot just put all of RAFT temporary allocations into the limited workspace_resource, because some of them scale with the problem size and would inevitably fail with OOM at some point.
6. Setting the workspace resource to the managed memory is not advisable as well for performance reasons: we have lots of small allocations in performance critical sections, so we need a pool, but a pool in the managed memory inevitably outgrows the device memory and makes the whole program slow.

Solution

I propose to split the workspace memory into two:

1. small, fixed-size workspace for small, frequent allocations
2. large workspace for the allocations that scale with the problem size

Notes:

• We still leave the full control over the allocator types to the user.
• Neither of the workspace resource should have unlimited lifetime / returned to the user. As a result, if the user sets the managed memory as the large workspace resource, the memory is guaranteed to be released after the function call.
• We have the option to use the slow managed memory without a pool for large allocations, while still using a fast pool for small allocations.
• We have more flexible control over which allocations are "large" and which are "small", so hopefully using the managed memory is not so bad for performance.

[tfeher]

Thanks Artem for proposing this solution. On one hand, it is nice to have a secondary workspace allocator to handle large allocations. I need to still think about this.

An alternative solution would be to keep a single workspace allocator, and that would provide the large allocator as a fall-back when allocating from the fast (but smaller) pool fails.

[tfeher]

The OOM errors we have seen with CAGRA were related to workspace pool grabbing all this place. What about limiting to a much smaller workspace size? (E.g. faiss has 1.5 GiB limit).

[achirkin]

That is an option, but so far I think it's not necessary. I also think it can hurt performance a little by reducing the batch size in places like ivf_pq::search or ivf_pq::extend.

With the current proposal, ann-bench executable (as a user of raft) set these resources:

• default - pool on top of device memory
• limited workspace - shares the same pool with default
• large workspace - managed memory (without pooling)

Hence the dataset/user allocations do not conflict for the same memory with the workspace (as they both use the same pool). At the same time, large temporary allocations (such as the cagra graph on device) use the managed memory and free it as soon as the algorithm finishes.

[achirkin]

Thanks for joining the conversation, Tamas. I've updated the description with my rationale since you reviewed the PR.
I think, the alternative solution you propose is a viable solution short-term to avoid OOMs, but it can be taxing on the performance: imagine a large allocation at the beginning of the algorithm (e.g. training set in kmeans) takes all of the device/workspace memory; then many subsequent small allocations are forced to use a tiny remaining free fraction of the memory via managed memory interface with a large oversubscription rate. This could lead to a terrible performance while the device memory is "wasted" on an allocation that may be not often used.

[achirkin]

Benchmarking update: there's a limited evidence that the update improves performance of cagra::build: I've got ~6% speedup with default parameters on DEEP-100M (which is anyway very slow and takes a lot of memory due to large default graph degrees and low ivf-pq compression).

[tfeher]

Thanks Artem for the update! It is a nice idea to have an extra memory resource that we can use for potentially host mem backed large temporary allocations. This can be useful for systems with improved H2D interconnect, such as Grace Hopper.

[harrism]

If you know you want all allocations within certain size ranges to be allocated with specific resources, you should have a look at binning_memory_resource. Basically you could have one or more fixed_size_resources for really small or common stuff, and a pool_mr (or cuda_async_mr) for larger stuff.

The difference is that choice of MR would be automatic and based on size, not on usage. The advantage of that is that implementers of RAFT functions wouldn't have to think about which MR to use so is potentially less bug prone.

But if you need to explicitly choose MRs based on non-size-based logic, then you might want to store separate resources instead. If the logic you need is common across all functions, then you might want to encode that logic into a custom memory resource class that decides which of a group of upstream resources to allocate from. Encoding the logic in the MR class again reduces potential bugs by not having to duplicate the logic everywhere.

[cjnolet]

It'll be nice when we can remove the need to pass around memory resources as arguments and we can just pull everything from the handle.

[achirkin]

Yes, but I'd suggest we do this during porting to cuVS to not make things even more complicated for users (if there are any).

[achirkin]

Opened #2322 dropping the changes to neighbor methods, which are moved to cuVS. Keeping this PR open, so that we can copy those neighbor changes when cuVS is ready for them.

[cjnolet]

@achirkin is this ready to be closed now that you've started a new PR for this?

[achirkin]

@cjnolet If you don't mind, I'd like to keep it open until we open cuVS PR with the corresponding neighbor changes.