support disaggregated weight update#4638
Conversation
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.
| 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))) |
lvhan028
added
enhancement
RunningLeon
left a comment
RunningLeon
left a comment
There was a problem hiding this comment.
Can we add monkey patched ut for the new apis?
|
Do we need to support TP=1 which might use UniExecutor or would we force distributed_executor_backend='ray'? |
I think we can bypass UniExecutor. |
Better keep it for debug. |
OK. Then we'd better define the developing rules at UniExcutor side |
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