GRPO hangs indefinitely on 2 RTX PRO 6000 Blackwell GPUs
- Dominant language
- Python
- Stars
- 2.5k
- Forks
- 345
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 240
Description
**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**
- [x] I have searched the existing issues for a similar bug report.
- [x] 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 :)
Contributor guide
Assessment
This issue has not been assessed yet.