support disaggregated weight update#4638
Open
irexyc wants to merge 3 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a disaggregated weight-update path to the PyTorch backend, allowing an external RL trainer to broadcast updated model weights into a running rollout engine via a dedicated NCCL process group (init / update / destroy lifecycle), and threads it through the engine, executors, model agent, and OpenAI HTTP API.
Changes:
- New OpenAI-compatible HTTP endpoints
/init_weights_update_group,/update_weights_from_distributed,/destroy_weights_update_group(PyTorch backend only) and matching pydantic request models. - New
init_weights_update_group/update_weights_from_distributed/destroy_weights_update_groupmethods plumbed fromEngine→ executor (base worker, ray executor) →BaseModelAgent, plus equivalents in the MP engine path. - Adds
init_custom_process_grouphelper inlmdeploy/utils.py(copied from XTuner) that creates a non-default-world NCCL group, including PyTorch ≥ 2.6backend_optionscompatibility.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| lmdeploy/utils.py | Adds init_custom_process_group helper to create a custom NCCL PG outside the default world. |
| lmdeploy/serve/openai/protocol.py | New request schemas for init/update/destroy disaggregated weight-update endpoints. |
| lmdeploy/serve/openai/api_server.py | Registers the three new endpoints with a PyTorch-backend-only guard. |
| lmdeploy/pytorch/engine/engine.py | Forwards the three new operations to the executor. |
| lmdeploy/pytorch/engine/executor/base_worker.py | Forwards calls to the model agent. |
| lmdeploy/pytorch/engine/executor/ray_executor.py | Implements collective_rpc dispatch and a shared _reduce_worker_status reducer. |
| lmdeploy/pytorch/engine/model_agent/agent.py | Core implementation: per-name PG dict, NCCL broadcast receive, main/draft split, finalization & graph reset. |
| lmdeploy/pytorch/engine/mp_engine/base.py | MPEngine wrappers for the three new methods. |
| lmdeploy/pytorch/engine/mp_engine/base_worker.py | MP engine worker delegates to underlying engine. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+501
to
503
|
|
||
|
|
||
|
|
Comment on lines
+1217
to
+1221
| if request.names: | ||
| named_tensors = [] | ||
| for name, dtype_str, shape in zip(request.names, request.dtypes, request.shapes): | ||
| target_dtype = getattr(torch, dtype_str) if isinstance(dtype_str, str) else dtype_str | ||
| named_tensors.append((name, torch.empty(shape, dtype=target_dtype, device=device))) |
Comment on lines
+1277
to
+1287
| group_name = request.group_name | ||
| pg = self._model_update_group.pop(group_name, None) | ||
| if pg is None: | ||
| return False, f'group {group_name!r} not initialized' | ||
| try: | ||
| dist.destroy_process_group(pg) | ||
| return True, f'Succeeded to destroy group {group_name!r}.' | ||
| except Exception as e: | ||
| msg = f'Failed to destroy weights update group {group_name!r}: {e}' | ||
| logger.exception(msg) | ||
| return False, msg |
lvhan028
added
enhancement
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Thanks for your contribution and we appreciate it a lot. The following instructions would make your pull request more healthy and more easily receiving feedbacks. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers.
Motivation
Please describe the motivation of this PR and the goal you want to achieve through this PR.
Modification
Please briefly describe what modification is made in this PR.
BC-breaking (Optional)
Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.
Use cases (Optional)
If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.
Checklist