GRPO hangs indefinitely on 2 RTX PRO 6000 Blackwell GPUs

[dirmeier]

Expected Behavior

One GRPO training step over a 2-GPU mesh (fsdp=2, tp=1) completes successfully.

Actual Behavior

It hangs indefinitely on the first rollout, in _generate_and_compute_advantage. I left one run for 57 minutes. Rerunning reproduces the issue every time.

• With NCCL_P2P_DISABLE=1 it runs to completion.
• It works flawlessly on 1 GPU.

I have py-spy stacks, NCCL debug logs and nsys profiles if any of that is useful.

Steps to Reproduce the Problem

1. gemma4_e2b orbax checkpoint + qwix LoRA adapter.
2. RLCluster with rollout_engine="vanilla", actor and rollout on the same
   Mesh(np.array(jax.devices()[:2]).reshape(2, 1), ("fsdp", "tp")).
3. One step: GRPOLearner(..., GRPOConfig(num_generations=2, beta=0.0, num_iterations=1)).train([{"prompts": np.array(["...", "..."])}]).

Same code with jax.devices()[:1] completes normally.

Environment

• OS: Ubuntu 26.04 LTS, kernel 7.0.0-30-generic
• Project Version: tunix 0.1.7, qwix 0.1.8, jax/jaxlib 0.11.1, NCCL 2.31.2
• GPU: 4x RTX PRO 6000 Blackwell Max-Q (sm_120), driver 595.84, CUDA 13.2.
  Single process, local devices, no jax.distributed.initialize()

Checklist

• I have searched the existing issues for a similar bug report.
• I have provided all the required information in the "Environment" section.
• I have provided a minimal, reproducible example. — my repro needs a 17 GB
  Gemma 4 E2B checkpoint; happy to cut it down if that would help.

Would you like to help us fix it?

Yes :)