[NVIDIA GPU] [XLA_GPU_MS_COLLECTIVE] Support round-robin runtime stream assignment #22450
Conversation
terryysun
changed the title
There was a problem hiding this comment.
Thank you for working on this. Multiple stream for collectives will be great to have!
This is a big PR and I think at least the stream assignment, the runtime integration, and some util functions could be three separate PRs.
One thing I am wondering is if we could do this in the latency-hiding scheduler. The scheduler models resources and we could have one per collective stream. That way the scheduler would perform the stream assignment and compose well with it. I'm not saying that would be better but it might be worth thinking about since you seem to run into issues with the LHS (implied by some comments).
| inline constexpr int64_t kAsyncStreamTotal = | ||
| static_cast<int64_t>(AsyncStreamKind::kMemCpyP2P) + 1; | ||
| std::max(static_cast<int64_t>(AsyncStreamKind::kMemCpyP2P) + 1, (int64_t)7); |
There was a problem hiding this comment.
Can you extract this 7 as a constant with a meaningful name?
| @@ -36,6 +36,21 @@ limitations under the License. | |||
| #include "xla/side_effect_util.h" | |||
|
|
|||
| namespace xla::gpu { | |||
| namespace { | |||
| static bool is_async_collective(const HloInstruction* instruction) { | |||
There was a problem hiding this comment.
There may be something like this in collective_ops_utils.cc. If not that could be a good place to add it.
| @@ -36,6 +36,21 @@ limitations under the License. | |||
| #include "xla/side_effect_util.h" | |||
|
|
|||
| namespace xla::gpu { | |||
| namespace { | |||
There was a problem hiding this comment.
No need for anonymous space for static functions
| @@ -40,8 +40,10 @@ enum class AsyncStreamKind : int64_t { | |||
| kMemCpyP2P = 3, // Stream for MemCpyP2P | |||
| }; | |||
|
|
|||
| // Taking the maximum of max stream kind + 1 and 4 (max compute stream) + 2 (max | |||
| // collective stream) + 1; | |||
There was a problem hiding this comment.
This is not clear to me. Is kAsyncStreamTotal a static upper bound? How is that enforced when the number of collective streams can be set per flag xla_gpu_experimental_parallel_collective_overlap_limit? Am I misunderstanding this?
| std::make_unique<NcclThunkType>(thunk_info, instr, buffers); | ||
|
|
||
| const ExecutionStreamAssignment& stream_assignment = | ||
| ir_emitter_context.execution_stream_assignment(); |
There was a problem hiding this comment.
Why is this passes separately? Can we add this to the thunks when emitting them? I think we do something like that for send and recv already
| static bool is_async_collective(const HloInstruction* instruction) { | ||
| if (instruction->IsAsynchronous()) { | ||
| auto opcode = instruction->async_wrapped_opcode(); | ||
| return opcode == HloOpcode::kAllGather || opcode == HloOpcode::kAllReduce || |
There was a problem hiding this comment.
I don't think this covers all collectives. If you move this to collective_ops_utils, we can test it too
| @@ -45,7 +60,11 @@ ExecutionStreamAssignment::ExecutionStreamAssignment( | |||
| // on the entrypoint computation will be assigned `ExecutionStreamId(0)`, and | |||
| // each invocation of `async-start` will result in the target computation | |||
| // being assigned a new `ExecutionStreamId`. | |||
| ExecutionStreamId next_stream_id = ExecutionStreamId(1); | |||
| ExecutionStreamId next_compute_stream_id = ExecutionStreamId(1); | |||
| ExecutionStreamId next_collective_stream_id = | |||
There was a problem hiding this comment.
Why is the tyoe of the next collective stream id ExecutionStreamId?
| @@ -405,7 +405,19 @@ GpuThunkAotCompilationResult::LoadExecutable( | |||
| compiler->BufferSizeBytesFunction(), | |||
| /*can_share_buffer=*/nullptr)); | |||
|
|
|||
| ExecutionStreamAssignment execution_stream_assignment(hlo_module.get()); | |||
| ExecutionStreamAssignment execution_stream_assignment( | |||
There was a problem hiding this comment.
This looks like duplicate code.
| @@ -126,19 +126,43 @@ bool IsAsyncPair(const HloInstruction& from, const HloInstruction& target) { | |||
| return IsGpuAsyncStart(from) && IsGpuAsyncDone(target); | |||
| } | |||
|
|
|||
| // Util function for getting replica groups from different data structures. | |||
| static std::vector<std::vector<int64_t>> get_replica_groups( | |||
There was a problem hiding this comment.
I think there is something like this in the collective ops utils. Can you mve it there and test it separately?
| std::vector<int64_t> ids({pair.first, pair.second}); | ||
| return ids; | ||
| }); | ||
| } else { |
There was a problem hiding this comment.
Does this work for all collectives except collective-permute start? Can we test that?
With #19026, LHS can overlap appropriate async collectives without deadlock. This PR adds support at runtime where we leverage the overlapping schedule produced by LHS and perform a round-robin stream assignment for collectives.