Skip to content

feat: Multi turn async #506

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 34 commits into from
Jun 26, 2025
Merged

feat: Multi turn async #506

merged 34 commits into from
Jun 26, 2025

Conversation

parthchadha
Copy link
Contributor

@parthchadha parthchadha commented Jun 11, 2025 •
edited
Loading

What does this PR do ?

Enables multi-turn rollouts with async engine.

Changes:

  1. Worker Groups:
    Allow consumers to handle async generators directly rather than blocking on result retrieval.

  2. vllm:
    VllmGeneration.generate_async():
    Returns AsyncGenerator now and coordinates multiple workers using queue-based async coordination.
    Yields samples as they become available from any worker, not waiting for entire batches

VllmGenerationWorker.generate_async()
Now yields (original_index, BatchedDataDict) as individual samples complete

  1. rollouts:
    run_async_multi_turn_rollout(): async rollout that processes each sample independently. Calls async_generate_response_for_sample_turn() to generate response for a single sample.

Issues

List issues that this PR closes (syntax):

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

parthchadha and others added 6 commits June 10, 2025 23:34
@parthchadha
WIP changes to make multi turn fully async
ca3d2ae 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
More changes to get multi gpu working
a62c5e7 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
Bring back commented test code
7b13d8b 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
Enable async rollouts for grpo
e7c8f5b 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@yuki-97
feat: support round-robin and fix multiturn (#515)
c334c6f 
Signed-off-by: Yuki Huang <yukih@nvidia.com>
@parthchadha
Fix failing unit test; add max_seqlen test for async
15b1064 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha parthchadha added the CI:L0 Run doctests and unit tests label Jun 17, 2025
@parthchadha parthchadha changed the title Multi turn async feat: Multi turn async Jun 17, 2025
parthchadha and others added 12 commits June 17, 2025 13:52
@parthchadha
WIP changes to make multi turn fully async
aee236a 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
More changes to get multi gpu working
efe3635 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
Bring back commented test code
cdca27a 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
Enable async rollouts for grpo
10580fa 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@yuki-97 @parthchadha
feat: support round-robin and fix multiturn (#515)
fa2990b 
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
Fix failing unit test; add max_seqlen test for async
21234c2 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@ZhiyuLi-Nvidia @terrykong @parthchadha
feat: token_mult_prob_error sample visualization if above a threshold (…
8b4124b 
…#389)

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Signed-off-by: ZhiyuLi-Nvidia <zhiyul@NVIDIA.com>
Co-authored-by: Terry Kong <terrycurtiskong@gmail.com>
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@yuki-97 @parthchadha
feat: support non-colocated sync vllm (#489)
bc798dc 
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@SahilJain314 @parthchadha
feat: Moving everything to 'Policy' and lm_policy for Megatron (remov…
93185f5 
…ing 'hf') (#511)

Signed-off-by: Sahil Jain <sahilj@nvidia.com>
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@joyang-nv @parthchadha
feat: add context parallel. (#450)
4b1331d 
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@SahilJain314 @yfw
@terrykong @parthchadha
feat: Add Megatron-LM based training (#517)
8985d65 
Signed-off-by: Sahil Jain <sahilj@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@terrykong @SahilJain314 @parthchadha
feat: async ray monitoring now tracks system memory (#349)
67056bd 
Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Sahil Jain <48468750+SahilJain314@users.noreply.github.com>
Co-authored-by: Sahil Jain <48468750+SahilJain314@users.noreply.github.com>
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha parthchadha removed the CI:L0 Run doctests and unit tests label Jun 17, 2025
@parthchadha parthchadha marked this pull request as ready for review June 17, 2025 20:56
@parthchadha parthchadha added the CI:L0 Run doctests and unit tests label Jun 17, 2025
@parthchadha parthchadha marked this pull request as draft June 17, 2025 20:57
@parthchadha
Fix failing large model unit test due to async api change
da1d714 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha
Merge remote-tracking branch 'origin/pchadha/multi-turn-fully-async' …
d7e69c1 
…into pchadha/multi-turn-fully-async
@parthchadha
Fix vllm test generation with async apis
943c465 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha parthchadha requested a review from SahilJain314 June 18, 2025 20:37
SahilJain314
SahilJain314 reviewed Jun 18, 2025
View reviewed changes
Copy link
Contributor

@SahilJain314 SahilJain314 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

reviewing in chunks

terrykong
terrykong reviewed Jun 18, 2025
View reviewed changes
@github-actions github-actions bot added the documentation Improvements or additions to documentation label Jun 24, 2025
@parthchadha parthchadha added CI:L0 Run doctests and unit tests and removed CI:L1 Run doctests, unit tests, and functional tests labels Jun 25, 2025
terrykong
terrykong reviewed Jun 25, 2025
View reviewed changes
terrykong
terrykong previously approved these changes Jun 25, 2025
View reviewed changes
@parthchadha
Fix failing vllm colocated unit test
e600069 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
@parthchadha parthchadha added CI:L0 Run doctests and unit tests and removed CI:L0 Run doctests and unit tests labels Jun 25, 2025
@parthchadha parthchadha enabled auto-merge June 25, 2025 19:24
@parthchadha parthchadha added this pull request to the merge queue Jun 25, 2025
@github-merge-queue github-merge-queue bot removed this pull request from the merge queue due to no response for status checks Jun 26, 2025
@terrykong terrykong added this pull request to the merge queue Jun 26, 2025
Merged via the queue into main with commit 41791c7 Jun 26, 2025
21 of 23 checks passed
@terrykong terrykong deleted the pchadha/multi-turn-fully-async branch June 26, 2025 10:08
@parthchadha parthchadha mentioned this pull request Jul 1, 2025
Open
therealnaveenkamal pushed a commit to therealnaveenkamal/RL that referenced this pull request Jul 7, 2025
@parthchadha @yuki-97
@ZhiyuLi-Nvidia @terrykong @SahilJain314 @joyang-nv @yfw @therealnaveenkamal
feat: Multi turn async (NVIDIA-NeMo#506)
66f22e5 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Signed-off-by: ZhiyuLi-Nvidia <zhiyul@NVIDIA.com>
Signed-off-by: Sahil Jain <sahilj@nvidia.com>
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Sahil Jain <48468750+SahilJain314@users.noreply.github.com>
Co-authored-by: yuki <48991475+yuki-666@users.noreply.github.com>
Co-authored-by: ZhiyuLi-Nvidia <zhiyul@NVIDIA.com>
Co-authored-by: Terry Kong <terrycurtiskong@gmail.com>
Co-authored-by: Sahil Jain <48468750+SahilJain314@users.noreply.github.com>
Co-authored-by: Jonas Yang CN <joyang@nvidia.com>
Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Co-authored-by: Yuki Huang <yukih@nvidia.com>
YzjiaoNvd pushed a commit to YzjiaoNvd/NeMo-RL that referenced this pull request Jul 14, 2025
@parthchadha @yuki-97
@ZhiyuLi-Nvidia @terrykong @SahilJain314 @joyang-nv @yfw
feat: Multi turn async (NVIDIA-NeMo#506)
0741cbc 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Signed-off-by: ZhiyuLi-Nvidia <zhiyul@NVIDIA.com>
Signed-off-by: Sahil Jain <sahilj@nvidia.com>
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Sahil Jain <48468750+SahilJain314@users.noreply.github.com>
Co-authored-by: yuki <48991475+yuki-666@users.noreply.github.com>
Co-authored-by: ZhiyuLi-Nvidia <zhiyul@NVIDIA.com>
Co-authored-by: Terry Kong <terrycurtiskong@gmail.com>
Co-authored-by: Sahil Jain <48468750+SahilJain314@users.noreply.github.com>
Co-authored-by: Jonas Yang CN <joyang@nvidia.com>
Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Co-authored-by: Yuki Huang <yukih@nvidia.com>
KiddoZhu pushed a commit that referenced this pull request Jul 28, 2025
@parthchadha @yuki-97
@ZhiyuLi-Nvidia @terrykong @SahilJain314 @joyang-nv @yfw @KiddoZhu
feat: Multi turn async (#506)
a0042a1 
Signed-off-by: Parth Chadha <pchadha@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Signed-off-by: ZhiyuLi-Nvidia <zhiyul@NVIDIA.com>
Signed-off-by: Sahil Jain <sahilj@nvidia.com>
Signed-off-by: Jonas yang <joyang@nvidia.com>
Signed-off-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Sahil Jain <48468750+SahilJain314@users.noreply.github.com>
Co-authored-by: yuki <48991475+yuki-666@users.noreply.github.com>
Co-authored-by: ZhiyuLi-Nvidia <zhiyul@NVIDIA.com>
Co-authored-by: Terry Kong <terrycurtiskong@gmail.com>
Co-authored-by: Sahil Jain <48468750+SahilJain314@users.noreply.github.com>
Co-authored-by: Jonas Yang CN <joyang@nvidia.com>
Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Co-authored-by: Yuki Huang <yukih@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@SahilJain314 SahilJain314 SahilJain314 left review comments

@yuki-97 yuki-97 yuki-97 left review comments

@terrykong terrykong terrykong approved these changes

Assignees
No one assigned
Labels
CI:L0 Run doctests and unit tests documentation Improvements or additions to documentation
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
6 participants
@parthchadha @yuki-97 @terrykong @SahilJain314 @ZhiyuLi-Nvidia @joyang-nv