Introducing `reduce` / `all_reduce` in RapidsMPF

[madsbk]

RapidsMPF needs two kinds of reduction operations:

1. Fixed-size reduction
2. Dynamic-size reduction over packed data

These two cases have different constraints and therefore belong at different layers of the stack.

Fixed-size reduction

Fixed-size reductions operate on buffers with a known size and layout (scalars, POD structs, fixed arrays).
Both MPI and UCX+UCC provide native reduce / allreduce for this case.

Because of this, fixed-size reduce / all_reduce should be part of the Communicator API:

• MPI backend → MPI_Allreduce / MPI_Reduce
• UCX/UCC backend → UCC ALLREDUCE collectives

This gives us efficient, backend-optimized collectives for simple fixed-size data.

However, UCXX currently does not expose UCC support, so enabling UCC-based collectives may require significant work (maybe a lot of work, @pentschev?). As an intermediate step, the UCXX communicator in RapidsMPF could implement reduce using the existing send/recv primitives.

Dynamic-size reduction (packed data)

Packed data varies in both size and structure, and no backend (MPI or UCX/UCC) supports all-reduce for variable-length payloads.

Supporting this case requires a custom protocol implemented above the communicator:

• exchange sizes
• reserve/allocate buffers (with spilling if needed)
• apply a user-provided reduction operator

For this reason, dynamic-size reductions should not be part of the low-level Communicator, but instead implemented in higher-level RapidsMPF logic.

Questions

• Do we agree that fixed-size reduction is the highest priority?
• Do we need user-defined reduction operators for fixed-size reduction, or is it sufficient to provide basic built-in operators (ADD, MUL, MIN, MAX) on fundamental datatypes such as int and float?

[wence-]

Note that collective communications in (both) MPI and UCC are matched up by program order on a particular communicator.

Since we cannot guarantee that tasks on different ranks execute in the same order, we need, at least the non-blocking collectives. And then, further, a different (collectively agreed) communicator per collective.

We will also need a way to poll these collectives without blocking, so that tasks on different ranks waiting for a collective can suspend so that everyone makes forward progress.

I think fixed-size is most needed (for variable size, which will probably be during groupby, we can either do allgathers or shuffles today).

I think we do want user-defined reduction operators, otherwise we have to encode all of the datatypes (at least) that cudf supports. e.g. we might want to run a reduction on a cudf::decimal datatype.

I think if we use MPI/UCC under the hood, providing access to the builtin ops is probably worthwhile.

[pentschev]

I think there are some more considerations we must make.

1. Do we foresee requiring other collectives?
2. UCC builds on top of UCX, however, it does not expose UCX internals (like workers, endpoints, etc.) nor can it be initialized from an existing UCX "cluster", therefore we will require to duplicate the UCX cluster when creating a UCC team. This may have implications on scalability.
   1. AFAICT, UCC is entirely collectives based, so we couldn't for example ditch the current UCX(X) communicator entirely in favor of UCC to prevent duplication the UCX cluster.
   2. Supporting UCC in UCXX seems like a project on its own, there's very little/nothing we could reuse from UCXX for UCC.
3. UCC uses the UCS API directly which provides absolutely no ABI guarantees (see ABI stability within patch releases openucx/ucx#10688), not even across commits, this has practical implications on versioning and packaging. While today we can support multiple UCX versions, with UCC we would be tied to whatever version of UCX that UCC is built for without the capacity for going to newer/older versions even for testing (unless we rebuild from source, which is often impractical especially when supporting other teams/customers).
4. Do we really want UCC (CPU-driven) or NCCL (GPU-driven)? This has a practical implication on performance, so it probably depends on what/how we use collectives for.

Questions

• Do we agree that fixed-size reduction is the highest priority?
• Do we need user-defined reduction operators for fixed-size reduction, or is it sufficient to provide basic built-in operators (ADD, MUL, MIN, MAX) on fundamental datatypes such as int and float?

I haven't been involved in discussions which led to the conclusions we now need collectives, so I can't agree or tell what we need. When we started RapidsMPF discussions ~9-12 months ago we decided not to use collectives entirely because they wouldn't fit the model we were intending for RapidsMPF, and Lawrence's comment above seems to confirm we still don't want/can't use "traditional" collectives because of program order even where we want these collectives. It would be good to have more context on what are the exact use cases for collectives in RapidsMPF and how this conclusion was reached.

[wence-]

When we started RapidsMPF discussions ~9-12 months ago we decided not to use collectives entirely because they wouldn't fit the model we were intending for RapidsMPF

This is not true, we started with async shuffle because that was the most pressing need.

It would be good to have more context on what are the exact use cases for collectives in RapidsMPF and how this conclusion was reached.

Since there is, by design, no central coordinator in streaming execution, any time we make an adaptive decision based on local data about (say) whether to do shuffle vs. broadcast join every rank needs to agree on the decision => allreduce.

Here are some (other) examples:

• when splitting up input parquet files across ranks if we have a skip_rows parameter to the read, rank-k needs to know how many rows ranks [0, ..., k-1] will read to determine its local skip_rows parameter => equivalent to MPI_Scan.
• Having decided to do a shuffle join, I want to pre-filter one of the sides approximately using a bloom filter. This can be done by constructing local (per rank) bloom filters and then doing an Allreduce with BITAND of the local bloom filters to get a global bloom filter.
• If we want rebalancing across ranks, then everyone probably has to participate to figure out where the new (rebalanced) split points want to be, before a neighbourwise exchange.

[pentschev]

Thanks for the clarifications, Lawrence.

• when splitting up input parquet files across ranks if we have a skip_rows parameter to the read, rank-k needs to know how many rows ranks [0, ..., k-1] will read to determine its local skip_rows parameter => equivalent to MPI_Scan.
• Having decided to do a shuffle join, I want to pre-filter one of the sides approximately using a bloom filter. This can be done by constructing local (per rank) bloom filters and then doing an Allreduce with BITAND of the local bloom filters to get a global bloom filter.
• If we want rebalancing across ranks, then everyone probably has to participate to figure out where the new (rebalanced) split points want to be, before a neighbourwise exchange.

So it seems that we're expecting at a minimum to require MPI_Scan, Allreduce and Allgather, is that correct? IOW, it seems like we need a "full-fledged" collectives framework, as in, unless we might need other collectives in some still to be determined future, it probably doesn't make sense for us to partially implementing collectives as proposed in the description:

As an intermediate step, the UCXX communicator in RapidsMPF could implement reduce using the existing send/recv primitives.

With that said, I think we need answers to the following questions, that maybe we already have and I'm just uninformed:

1. What is the estimated timeline to need each of the collectives operations? It seems fixed-size reduce is needed immediately, but are others too?
   1. Would a single fixed-size Allreduce collective satisfy our collectives needs for the near-future (3-6 months)?
2. We have an AllGather implementation that, AFAIU, is more robust than MPI's, as it is non-blocking and thus not strictly dependent on program order. For other collectives, do we need collectives that operate regardless of program order, or are we generally ok with the block-the-world collectives?

[madsbk]

First iteration could be a non-blocking collective reduce operator that uses the existing Communicator API. I think the performance would be sufficient for our use case?

[wence-]

So it seems that we're expecting at a minimum to require MPI_Scan, Allreduce and Allgather, is that correct?

Scan and reduce, not allgather.

[wence-]

First iteration could be a non-blocking collective reduce operator that uses the existing Communicator API.

What approach do you want to take for multiple "out of order" collectives to run simultaneously?

[madsbk]

User provides a tag like we do with shuffler and all gather?

[pentschev]

I can give it a shot at the reduce tomorrow.

[nirandaperera]

@wence-

Note that collective communications in (both) MPI and UCC are matched up by program order on a particular communicator.

This is not true. UCC collectives are tagged. https://openucx.github.io/ucc/api/v1.4.4/html/group___u_c_c___c_o_l_l_e_c_t_i_v_e_s.html#structucc__coll__args

[nirandaperera]

If we want to take an iterative approach, I propose that we start off Allreduce, by piggybacking on Allgather impl. Then we can finalize the front-end API for reductions.
I have some questions @madsbk @wence- @pentschev

• How are we planning to handle types?
• Take in the reduction operation as lambdas? or as enums?
• Also, I think there are various reduction types, eg associative (sum), "partially" associative (mean), etc. So, it would be good to restrict ourselves to associative reductions for the moment (sum, min, max, boolean ops)

[pentschev]

@wence-

Note that collective communications in (both) MPI and UCC are matched up by program order on a particular communicator.

This is not true. UCC collectives are tagged. https://openucx.github.io/ucc/api/v1.4.4/html/group___u_c_c___c_o_l_l_e_c_t_i_v_e_s.html#structucc__coll__args

Niranda is right, I didn't know about it either, but it seems like UCC can match solely based on tags and thus not enforce a program order requirement, which is different from MPI and NCCL (the latter also must ensure program order).

Although UCC would still require duplicating the communicator currently, it seems well suited for the programming model we want/need for RapidsMPF. So perhaps short-term (immediately follow this route) or mid-term (implement our own Allreduce for now, and then follow this route after) could be useful to instead of implementing our own collectives just use UCC even if directly (without a C++ wrapper/API like UCXX), then work with the UCC devs to request exposing underlying UCP handles (like ucp_worker_h and ucp_ep_h) to also support P2P comms we need (e.g., eventually add support for UCC in UCXX, exposing a C++ collectives API, then have UCXX retrieve UCP handles from UCC and re-expose the P2P communicators via already existing tagSend/tagRecv calls, for example). WDYT?

[nirandaperera]

then work with the UCC devs to request exposing underlying UCP handles (like ucp_worker_h and ucp_ep_h) to also support P2P comms we need

@pentschev I actually asked this question some time back (I requested other way around). But from using UCC previously, what I gathered was, it requires the whole UCP tag space, because the are partitioning the 64 bits in the tag to support tagged collective (like we do, in a way). So, this one reason why it needs to create its own context. May be we can initiate a conversation with them.

But I'm all for reusing an existing impl, rather than reinventing the wheel 😇

What would be the overhead of managing two sets of UCX workers and contexts?

[pentschev]

What would be the overhead of managing two sets of UCX workers and contexts?

The overhead is something that UCX probably cannot (and most likely will not, even if they could, as that would be enormous amounts of work) deal with because it's tied to the underlying OS and hardware. For example, when using sockets (e.g., TCP) that would require two pairs of endpoints, which results in insufficient number of file descriptors at scale. Similarly, InfiniBand uses queue pairs to exchange information between devices, which is also a limited hardware resource, and similarly would require duplicated use of that resource, which at scale will become insufficient -- this primarily applies to rc, dc overcomes that limitation to certain extent when UCX can manage all queue pairs, if we now have to split queue pairs among UCXX and UCC then it will probably have an even higher impact in performance with dc.

For short-term while we're still not working at high scale duplication could work, but if we go over a certain number of processes this will become problematic. Previously we've seen issues with as few as 64 GPUs/processes for sockets, and I believe hundreds of processes for InfiniBand queue pairs, and this was without duplicating resource consumption, we would probably have to cut those numbers in half if we do so.

I actually asked this question some time back (I requested other way around). But from using UCC previously, what I gathered was, it requires the whole UCP tag space, because the are partitioning the 64 bits in the tag to support tagged collective (like we do, in a way). So, this one reason why it needs to create its own context. May be we can initiate a conversation with them.

Thanks for sharing that conversation, it's indeed insightful. If we decide to pursue the UCC route we will definitely need to talk to them to explore the options. I know some of those guys and if we choose to try it out I'd be happy to get the conversation started.