[TRAIN-1] Megatron async GRPO with dis-colocated generation deadlocks at weight sync
- 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
Assessment
This issue has not been assessed yet.