NVIDIA-NeMo / NVIDIA-NeMo/RL

[TRAIN-2] DTensor async GRPO with dis-colocated generation hangs when tensor_parallel=2

Open
#2,718 3 comments 0 reactions 1 assignee Claimed by @ZhiyuLi-Nvidia View on GitHub
bug p0 Speed
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Problem

Dis-colocated async GRPO on the DTensor backend hangs during training when `tensor_parallel` is set to 2. Setting `tensor_parallel=1` unblocks training. This points to a synchronization issue that is specific to TP > 1.

## Steps to reproduce

1. Configure dis-colocated async GRPO with the DTensor backend and `tensor_parallel: 2`.
2. Start training and observe a hang with no error output.
3. Set `tensor_parallel: 1` and confirm training proceeds normally.

## Suggested fix

Audit synchronization barriers in the DTensor async rollout path for behavior that differs by tensor parallelism count.

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.