Fix three DDP correctness issues: PMI guardrails, data padding, unuse… by edopica · Pull Request #2694 · FluxML/Flux.jl

edopica · 2026-07-22T13:02:41Z

Fix three DDP deadlock/crash scenarios in DistributedUtils

This PR addresses three independent correctness issues in Flux's distributed
training path (MPI backend), originally introduced in the distributed data parallel PR (#2464). Each bug causes either a deadlock or a hard crash
when running DDP with non‑trivial models or launch configurations.

Additionally, this PR enables MPI distributed tests to run automatically on Linux CI runners.

1. PMI version mismatch guardrails

Problem: Julia's default MPICH_jll speaks PMI2. When launched under a
PMIx launcher (e.g. srun without --mpi=pmi2) or under OpenMPI's mpirun,
MPI.Init() aborts with an uninformative message, making diagnosis difficult.
Even inside a Slurm allocation with mpiexecjl, the missing PMI2 variables
produce a confusing hang.

Fix: Before calling MPI.Init(), check for environment markers:

PMIX_RANK → error with PMIx mismatch message
OMPI_COMM_WORLD_RANK → error with OpenMPI mismatch message
SLURM_JOB_ID but no PMI2_* → warn about missing --mpi=pmi2

A force=true keyword allows advanced users to bypass the check (e.g. when
using a custom MPI build). The logic is entirely inside FluxMPIExt and
runs before any MPI symbols are touched.

2. Padding in DistributedDataContainer to prevent allreduce deadlocks

Problem: When n_samples % n_workers ≠ 0, DistributedDataContainer
assigned fewer batches to the last rank. That rank would exit the training
loop early while the other ranks blocked inside an allreduce call, causing a
permanent deadlock.

Fix: Pad the index sequence so every rank receives the same number of
elements. Surplus slots are filled by wrapping around from the start
([1, 2, …, n_samples, 1, 2, …]). The data‑loading logic is unchanged;
only the index array passed to DistributedDataContainer is padded.

3. `resolve_unused_parameters!!` for conditional computation graphs

Problem: A model with conditional branches (e.g. if/else selecting
different layers) can produce nothing gradients on ranks that did not execute
a particular parameter. When DistributedOptimizer later calls allreduce
over all gradients, the mismatched set of buffers deadlocks — different ranks
participate in different collectives.

Fix: New public function resolve_unused_parameters!!(backend, gs, model)
replaces every nothing gradient with a zero‑filled array matching the
parameter's shape, element type, and device. It walks nested structures
(NamedTuples, etc.) and works with arbitrary tree depths thanks to Functors'
fmap.

Tests added / updated

File	Tests
`test/ext_distributed/data_distributedtest.jl`	Renamed from `data.jl` to run with the test suite. 3 test groups: evenly-divisible partition, non-divisible (padded) partition with correct padded sum, and the original 10‑element case with expected padded sum.
`test/ext_distributed/unused_parameters_distributedtest.jl` (new)	4 test groups (20 tests): replaces‑nothing‑with‑zeros, nested NamedTuples, no‑op when all present, preserves exact shapes and element types (Float32/Float64). Runs on both MPI and NCCL backends.
`test/ext_distributed/reduce_distributedtest.jl`	Fixed pre-existing upstream breakage in `reduce_distributedtest.jl` (added missing imports required for the test to run at all).
Pre-existing Orphaned Tests	Renamed `common.jl`, `optimizer.jl`, and `synchronized.jl` to `_distributedtest.jl` suffix so they are picked up by the test runner. Also fixed pre-existing upstream breakage (missing imports).

All tests pass with mpiexecjl -n 2 on CPU/MPI.

Backward compatibility

All changes are additive or replace a broken code path with a working one.
resolve_unused_parameters!! is a new, optional public API — existing
training loops that don't use conditional computation are unaffected.
The PMI check can be bypassed with force=true.

CarloLucibello · 2026-07-31T04:03:11Z

+           l, grad = Zygote.withgradient(loss, model)
+
+           # Resolve nothing gradients to avoid deadlocks
+           grad_resolved = DistributedUtils.resolve_unused_parameters!!(backend, grad[1], model)


Users shouldn't do this manually.
We should call resolve_unused_parameters!! internally, as part of Flux.withgradient or Optimisers.update. This could be part of a future PR, but for the time being I would avoid adding anything to the guide.

CarloLucibello · 2026-07-31T04:09:02Z

+backend_string = ARGS[1]
+
+if backend_string == "mpi"
+    import MPI
+    const backend_type = MPIBackend
+elseif backend_string == "nccl"
+    import MPI, NCCL, CUDA
+    const backend_type = NCCLBackend
+else
+    error("unsupported backend: $backend_string")
+end
+
+const dev = Flux.cpu


this boilerplate code could be added to
https://github.com/FluxML/Flux.jl/blob/master/test/test_module.jl

instead of being part of each test file

CarloLucibello · 2026-07-31T04:10:01Z

+using Optimisers
+
+
+backend_string = ARGS[1]


instead of reading ARGS this should be set by ENV variables in
https://github.com/FluxML/Flux.jl/blob/master/test/runtests.jl
as anything else

CarloLucibello · 2026-07-31T04:12:10Z

before merging this PR, CI should be activated for the distributed tests, currently it is off

…d parameters - Add PMIx/PMI2 mismatch detection in MPI backend initialization checks before MPI.Init(), with force kwarg bypass. Produces a clear error instead of a hard abort. - Implement data padding in DistributedDataContainer to ensure even distribution across workers, preventing load imbalance. - Add resolve_unused_parameters!! to replace nothing gradients with zero-filled arrays, preventing allreduce deadlocks in conditional-computation models. - Add missing imports in reduce_distributedtest.jl for CuArray and ROCArray, making it runnable under NCCL. - Rename orphaned distributed tests (common.jl, optimizer.jl, synchronized.jl, data.jl) to _distributedtest.jl suffix so they are discovered by the test runner. - Convert all distributed test files to single-arg ARGS pattern consistent with how the test runner launches them. - Add unused_parameters_distributedtest.jl test suite. - Add documentation in docs/src/guide/gpu.md for resolve_unused_parameters!!. - Add a new comprehensive distributed training guide in docs/src/guide/distributed.md. - Enable distributed MPI tests to run on Linux CI via test/runtests.jl.

- Flux's DistributedUtils.initialize already handles device isolation automatically - Explicit device assignment can conflict with NCCL backend initialization

@test

…uite default) - test/runtests.jl: revert FLUX_TEST_DISTRIBUTED_MPI default to "false" (exactly upstream master); main suite no longer runs ext_distributed inside ParallelTestRunner, which passes no ARGS (BoundsError in backend_string = ARGS[1]) - .github/workflows/distributed_ci.yml: dedicated 2-rank fast + 4-rank edge MPI jobs, system OpenMPI + MPIPreferences, launcher invoked directly with --project=test so MPI resolves (root-project weakdep is NOT loadable via using; verified empirically) - test/ext_distributed/runtests.jl: watchdog timeout (120 s default, FLUX_TEST_DISTRIBUTED_TIMEOUT) + real exit-code propagation (proc.exitcode == 0) replacing the silent Test.@test true - test/Project.toml: MPI as explicit test dependency

The PMIx/OMPI env checks assumed MPICH_jll: with system OpenMPI (PMIx), PMIX_RANK is the correct protocol and the guardrail aborted every OpenMPI launch (CI MPI jobs failed). Guardrail now checks MPI.MPI_LIBRARY and skips PMIx/PMI2 mismatch errors and the Slurm PMI2 warning for OpenMPI-based libraries; messages name the loaded library. Verified: OpenMPI_jll 2/4-rank launcher runs pass, MPICH_jll path unchanged, MPICH + PMIX_RANK still errors.

CarloLucibello reviewed Jul 31, 2026

View reviewed changes

edopica force-pushed the ddp/upstream-pr branch from df95f60 to 504e816 Compare August 17, 2026 11:30

edopica added 2 commits August 17, 2026 13:33

docs: remove explicit CUDA.device! from DDP guide

bc40df2

- Flux's DistributedUtils.initialize already handles device isolation automatically - Explicit device assignment can conflict with NCCL backend initialization

edopica force-pushed the ddp/upstream-pr branch from 504e816 to bc40df2 Compare August 17, 2026 11:34

edopica added 2 commits August 17, 2026 23:11

Navigation Menu

Uh oh!

Uh oh!

Fix three DDP correctness issues: PMI guardrails, data padding, unuse… - #2694

Fix three DDP correctness issues: PMI guardrails, data padding, unuse…#2694
edopica wants to merge 4 commits into
FluxML:masterfrom
edopica:ddp/upstream-pr

edopica commented Jul 22, 2026

Uh oh!

CarloLucibello Jul 31, 2026

Uh oh!

CarloLucibello Jul 31, 2026

Uh oh!

CarloLucibello Jul 31, 2026

Uh oh!

CarloLucibello commented Jul 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Uh oh!

Conversation

edopica commented Jul 22, 2026

Fix three DDP deadlock/crash scenarios in DistributedUtils

1. PMI version mismatch guardrails

2. Padding in DistributedDataContainer to prevent allreduce deadlocks

3. resolve_unused_parameters!! for conditional computation graphs

Tests added / updated

Backward compatibility

Uh oh!

CarloLucibello Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

CarloLucibello Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

CarloLucibello Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

CarloLucibello commented Jul 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

3. `resolve_unused_parameters!!` for conditional computation graphs