google / google/tunix

GRPO hangs indefinitely on 2 RTX PRO 6000 Blackwell GPUs

Open
#2,196 0 comments 0 reactions 1 assignee Claimed by @sizhit2 View on GitHub
type:bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.