Adds foundations for GPU MPI support via NCCL

[romanlee]

Add the foundations for GPU MPI support (via NCCL) to ReactantMPIExt. Mainly, this amounts to adding logic and state to properly set up NCCL. This logic is concentrated in ext/ReactantMPIExt/Nccl.jl, and specifically in the function ReactantMPIExt.initialize!(MPI.COMM_WORLD) which:

• explicitly initializes an XLA gpu client with exactly one allowed device per MPI rank and then sets the NCCL device accordingly.
• sets up the NCCL communicator. The address of the communicator ultimately gets passed to EnzymeJAX as a pass option on LowerEnzymeXLAPass.

We also add backend gpu tests for MPI Allreduce, which will be supported when EnzymeAD/Enzyme-JAX#2699 lands

NOTE: Ultimately we might want to move NCCL communicator creation to EnzymeJAX, but this is what we agreed on for a first pass.

[github-actions]

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

JuliaFormatter

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 263 in d6e4ba0

MPI.Send(send_buf, dest, tag, comm)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 267 in d6e4ba0

MPI.Recv!(recv_buf, src, tag, comm)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 271 in d6e4ba0

@jit sendrecv!(comm, rank, send_buf, recv_buf, tag)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 275 in d6e4ba0

end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 276 to 278 in d6e4ba0

	elseif RUN_GPU_MPI_TESTS
	@info "Skipping GPU MPI Send / Recv! tests; Not implemented"
	end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 280 to 290 in d6e4ba0

	if RUN_CPU_MPI_TESTS
	@testset "Isend / Irecv! / Wait" begin
	comm = MPI.COMM_WORLD
	rank = MPI.Comm_rank(comm)
	for T in datatypes
	# NOTE: currently don't allow a request to cross the compile boundary
	# debugging tip: if this fails, can use pair Send with Irecv! + Wait, or Recv! with
	# Isend + Wait to isolate the issue
	send_buf = ConcreteRArray(ones(T, 5))
	recv_buf = ConcreteRArray(zeros(T, 5))

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 292 to 311 in d6e4ba0

	function isendirecvwait(send_buf, recv_buf, rank, tag, comm)
	if rank == 0
	dest = 1
	req = MPI.Isend(send_buf, dest, tag, comm)
	MPI.Wait(req)
	return nothing
	elseif rank == 1
	src = 0
	req = MPI.Irecv!(recv_buf, src, tag, comm)
	MPI.Wait(req)
	return nothing
	end
	end
	@jit isendirecvwait(send_buf, recv_buf, rank, tag, comm)
	rank == 1 && @test recv_buf == send_buf
	end
	end
	elseif RUN_GPU_MPI_TESTS
	@info "Skipping GPU MPI Isend / Irecv! / Wait tests; Not implemented"
	end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 313 to 342 in d6e4ba0

	if RUN_CPU_MPI_TESTS
	@testset "Isend / Irecv! / Waitall" begin
	comm = MPI.COMM_WORLD
	rank = MPI.Comm_rank(comm)
	tag = 42
	for T in datatypes
	# NOTE: currently don't allow a request to cross the compile boundary
	function waitall(send_buf, recv_buf)
	reqs = Reactant.TracedRNumber[]
	if rank == 0
	dest = 1
	src = 1
	req = MPI.Irecv!(recv_buf, src, tag - 1, comm)
	push!(reqs, req)
	req = MPI.Isend(send_buf, dest, tag + 1, comm)
	push!(reqs, req)
	elseif rank == 1
	dest = 0
	src = 0
	req = MPI.Isend(send_buf, dest, tag - 1, comm)
	push!(reqs, req)
	req = MPI.Irecv!(recv_buf, src, tag + 1, comm)
	push!(reqs, req)
	end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 344 to 346 in d6e4ba0

	reqs = vcat(reqs...)
	return MPI.Waitall(reqs)
	end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 348 to 349 in d6e4ba0

	send_buf = ConcreteRArray(ones(T, 5))
	recv_buf = ConcreteRArray(zeros(T, 5))

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 351 in d6e4ba0

@jit waitall(send_buf, recv_buf)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 353 to 361 in d6e4ba0

	# debug
	# if rank==0
	# println("\ncode_hlo optimize=false:\n",
	# @code_hlo optimize=false waitall(send_buf, recv_buf))
	# println("\ncode_hlo optimize=\"lower-enzymexla-mpi{backend=cpu}\":\n",
	# @code_hlo optimize="lower-enzymexla-mpi{backend=cpu}" waitall(send_buf, recv_buf))
	# println("\ncode_hlo:\n",
	# @code_hlo waitall(send_buf, recv_buf))
	# end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 363 in d6e4ba0

@test recv_buf == send_buf

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 365 in d6e4ba0

end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 366 to 368 in d6e4ba0

	elseif RUN_GPU_MPI_TESTS
	@info "Skipping GPU MPI Isend / Irecv! / Waitall tests; Not implemented"
	end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 370 to 390 in d6e4ba0

	if RUN_CPU_MPI_TESTS
	@testset "Bcast!" begin
	comm = MPI.COMM_WORLD
	rank = MPI.Comm_rank(comm)
	root = 0
	for T in datatypes
	@testset "Type: $T" begin
	# just the root have the real values, others have zeros
	if rank == root
	x = ones(T, 5)
	else
	x = zeros(T, 5)
	end
	# try block catches any invalid combinations we missed above, depending on
	# mpi implem
	expected = try
	ConcreteRArray(MPI.Bcast!(x, root, comm))
	catch
	continue
	end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 392 in d6e4ba0

@test expected == @jit MPI.Bcast!(ConcreteRArray(x), root, comm)

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Line 395 in d6e4ba0

end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/integration/mpi.jl

Lines 396 to 398 in d6e4ba0

	elseif RUN_GPU_MPI_TESTS
	@info "Skipping GPU MPI Bcast! tests; Not implemented"
	end

---

[JuliaFormatter] _{reported by reviewdog 🐶}

Reactant.jl/test/runtests.jl

Line 143 in d6e4ba0

)

[avik-pal]

XLA:CUDA already ships with NCCL, why do we need to load these via an extension?

[romanlee]

Hm yeah I suppose we could do it that way. Would that be better?

[romanlee]

Yeah actually on reflection that is definitely better. I will refactor this.

[romanlee]

@avik-pal let me know how things look to you now, see 4f9dab6.

[romanlee]

Note also that in the future, we'll probably move nccl communicator creation to enzymejax, but setting it up somewhat clunkily in reactant was what @wsmoses and I decided on for a first pass.

[mofeing]

@avik-pal we agreed to get this done like @romanlee to get some MPI runs working soon. I agree we should better use the XLA-built NCCL (and since we do that, we can maybe write our own custom calls directly in Enzyme-JAX), but we can fix it and improve it incrementally. That's what also the new comm dialect is about.

@romanlee I feel like this should better live in its own NCCL extension. You're using a lot of MPI to set it up, but the user is gonna know if it's going to execute on a GPU cluster with NCCL. Ultimately, the NCCL backend for MPI should be used for converting user-code, but the set up should be done with NCCL and not MPI. It's true that we are gonna use XLA-built NCCL, which can lead to some conflicts, and would like to know your opinion.

EDIT: I see that you had your code in a NCCL extension first, and I think we should discuss it.

[mofeing]

just curious, is preallocation giving problems? also, if you don't preallocate, then does XLA_REACTANT_GPU_MEM_FRACTION does anything?

[romanlee]

Glad you pointed this out. This shouldn't be necessary so I'll remove it, was a hydra thing.

Incidentally, I assume MEM_FRACTION sets the budget regardless of whether it's preallocated or allocated on demand? That's what other usage would suggest, eg just above in the same file

[mofeing]

mmm gpu is CUDA but also ROCm, which can right now gives problems because we don't yet support RCCL (which will be trivial to support once we merge this or other features)

[romanlee]

Thanks, see 82c9567

[romanlee]

Wait, thinking a bit more about this, are you sure this is right? Doesn't "gpu" sometimes mean "cuda"? How do we make sure to catch that case properly?

Also, as written I have "auto" routing to set_default_backend("cpu"), which probably isn't right either. Because "auto" defaults to "cuda" or "rocm" with higher priority than "cpu", which we probably want to respect.

If you could clarify the intended usage of "REACTANT_BACKEND_GROUP" here that would be very helpful.

[romanlee]

Alright, I sorted this out. See 14a4e76.

Feel free to double check my logic, but it should be all good

[mofeing]

So... you shouldn't call Base.get_extension inside regular lib code. Instead, you should declare it but left undefined, and define it in the appropriate extension.

But since whole ReactantMPIExt/Nccl.jl file doesn't depend on anything inside the ReactantMPIExt module, maybe better move to its own file inside the src/ folder?

[romanlee]

Ah thanks, this didn't feel quite right to me either. See 54a04fc. Re set_nccl_device!, I think it's appropriate as is. Let me know if you disagree.

And re moving Nccl.jl to src/, I think it makes sense to leave it in the ReactantMPIExt module even if it doesn't strictly need to be there, since it is always used together with the mpi extension. But if that's not the right style/you feel strongly I'm happy to move it

[mofeing]

maybe moving the try/finally to the runtests.jl file would make it more legible? wdyt?

[romanlee]

I don't think that's really possible bc we can't access the stuff in the finally block in runtests.jl? Anyway, that feels like it would be slightly less readable to me, so I'd prefer to keep it where it is.

That said, I did make some minor changes to make it more readable here a0ebd04

[mofeing]

mind prefixing sth like REACTANT_NCCL_ here? the reasoning being that there is overlap with MPI and both are on the same module

[mofeing]

if you move to its own module, that would be another solution.

[romanlee]

Good idea. I renamed to DEFAULT_COMM_* to NCCL_COMM_*. And DEFAULT_XLA_DEVICE got separately refactored away.

[romanlee]

I feel like this should better live in its own NCCL extension. You're using a lot of MPI to set it up, but the user is gonna know if it's going to execute on a GPU cluster with NCCL. Ultimately, the NCCL backend for MPI should be used for converting user-code, but the set up should be done with NCCL and not MPI. It's true that we are gonna use XLA-built NCCL, which can lead to some conflicts, and would like to know your opinion.

@mofeing I actually think it's better to not have a NCCL extension for a few reasons: (1) we avoid incompatibility issues between nccl.jl's and xla's packaged nccl libs, (2) since what we're doing with nccl is relatively straightfoward (just setting up a communicator) it's relatively easy to just create the wrappers ourselves and that allows us to avoid another dependency, and (3) ultimately we'll move communicator creation out of Reactant anyway, at which point we won't need the extension, so it's fine to do something lightweight for now. Curious what @avik-pal's opinion is, but I assume based on his original comment that he was also thinking better to avoid an extension unless necessary.

Re using a lot of MPI to set it up, that's just how you set up NCCL. See https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#creating-a-communicator.

[romanlee]

After thinking more about it, I decided to refactor the way the NCCL is initialized. See 15662fc.

Specifically, previously we selected the NCCL device after XLA client creation. This was fragile because it left all devices addressable. This refactor adds a function ReactantMPIExt.initialize!, which explicitly initializes the XLA client with exactly one allowed GPU per local MPI rank and then sets the NCCL device accordingly.