NVIDIA-NeMo / NVIDIA-NeMo/RL

[TRAIN-1] Megatron async GRPO with dis-colocated generation deadlocks at weight sync

Open
#2,717 1 comment 0 reactions 1 assignee Assigned to @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

Async GRPO with `colocate_generation: false` on the Megatron backend hangs permanently at the first weight sync (step 1 to 2). After the first training step completes, the trainer and the async rollout collector each wait on the other. The pipeline stalls with idle GPUs and never recovers. The same step runs fine with colocated generation and on the DTensor backend.

## Steps to reproduce

1. Configure async GRPO with `colocate_generation: false` using the Megatron backend.
2. Start training.
3. Observe the pipeline stall after step 1 with no error output and idle GPUs.

## Suggested fix

- Investigate the weight handoff coordination between `MegatronPolicyWorker` and `VllmAsyncGenerationWorker`.
- Add a timeout and explicit error on weight-sync deadlock so the failure surfaces rather than hanging silently.

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.