[ckpt] feat: add CheckpointEngineManager #5031
Conversation
wuxibin89
requested review from
PeterSH6,
chenhaiq,
eric-haibin-lin,
tongyx361,
vermouth1992,
zhaochenyang20 and
zw0610
as code owners
wuxibin89
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the CheckpointEngineManager to streamline weight synchronization between trainer and rollout replicas, which is a significant architectural improvement. It refactors existing checkpoint engine tests to use this new manager and introduces new test cases. The changes also include updates to configuration structures and utility functions to support asynchronous operations and ensure correctness during weight transfers.
wuxibin89
requested review from
FightingZhen,
ISEEKYAN,
ji-huazhong and
tardis-key
as code owners
wuxibin89
force-pushed
the
wuxibin/checkpoint_engine_manager
branch
from
9452913 to
79fa28e
Compare
| self, | ||
| backend: str, | ||
| trainer: RayWorkerGroup, | ||
| replicas: list[RolloutReplica], |
There was a problem hiding this comment.
Suggesting avoid initializing from workers and let AgentLoop register itself instead, hiding replicas from CheckpointEngineManager. Otherwise if we use P2P backend to support elastic rollout, we need to go through all the labors to pass replicas back to CheckpointEngineManager even if we don't need to rebuild process group.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a CheckpointEngineManager to centralize and streamline weight synchronization between trainer and rollout replicas. This is a significant architectural improvement that enhances the management of distributed model updates. The changes involve refactoring configuration structures, abstracting checkpoint engine operations, and updating various worker implementations to integrate with the new manager. This new abstraction is well-designed and improves the clarity and maintainability of the weight synchronization logic. However, there are a few areas that need attention to ensure full functional correctness and test coverage.
| self._init_agent_loop_workers() | ||
|
|
||
| # Initially we're in sleep mode. | ||
| if self.config.actor_rollout_ref.rollout.free_cache_engine: |
There was a problem hiding this comment.
checkpoint engine 支持后，不需要再用sleep/wakeup 去切换了，是吗？
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
There was a problem hiding this comment.
Is there any relevant documentation for the checkpoint engine? Once this feature is supported, will it no longer impede the subsequent Elastic rollout?
4 participants
What does this PR do?
#4280 refactor vllm breaking
one-step-off-policyandfully-async. This PR introduce CheckpointEngineManager to coordinate weight synchronization between trainer and rollout replicas.Next PR, refactor
one-step-off-policyandfully-asyncwith CheckpointEngineManager.design doc: https://github.com/volcengine/verl/tree/main/verl/checkpoint_engine