Skip to content

[rollout] feat: partial rollout part1 - rollout batch dynamic adjustment#2101

Closed
chenhaiq wants to merge 4 commits intoverl-project:mainfrom
chenhaiq:dynamic_batch
Closed

[rollout] feat: partial rollout part1 - rollout batch dynamic adjustment#2101
chenhaiq wants to merge 4 commits intoverl-project:mainfrom
chenhaiq:dynamic_batch

Conversation

@chenhaiq
Copy link
Collaborator

Partial rollout need the ability to postpone unfinished samples to later batches, and prefetch samples from later batches.

This PR make DataProto adjustable between batches, and provides an extension point for who want more advanced policy to adjust the batch. The main idea is to change DataProto to a linked list, so that the chat scheduler can access data from later batch and move samples between them.

For example:
postpone unfinished samples to next batch is the most straight idea to achieve partial rollout,
however, someone may think moving those unfinished samples to the last batch in the epoch is a better idea.

The new LinkedDataProto class provides there methods that can be override in subclass for customization:

  method purpose description 
LinkedDataProto.prefetch_samples(count) -> DataProto prefetch samples to later batches Return 0 to indicate a fixed batch and retain the old behavior
LinkedDataProto.postpone_samples(samples: DataProto) postpone samples to later batches Return False to indicate a fixed batch and retain the old behavior
LinkedDataProto.should_stop_gen() whether chat scheduler should stop rollout in advance True to indicate should stop

Notes: this PR does not change any behavior. It will work with part2 PR in chat scheduler together to enable partial rollout
The reason to submit this PR is to receive feedback before the part2 PR is completed.

Test

tests/test_protocol_on_cpu.py

@chenhaiq
Copy link
Collaborator Author

chenhaiq commented Jul 1, 2025

move partial rollout with streaming to #2200

@chenhaiq chenhaiq closed this Jul 1, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant