[Core] Add `update_load_config` RPC method #19884

22quinn · 2025-06-20T01:17:05Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

This is to make #18745 feature-complete. Without an load config update API, there's no way to load real weights inplace.
With this PR now we can perform the following:

Initialize engine with dummy weights
Update load config
Load real weights inplace

Test Plan

See #19640 for E2E example

Test Result

See #19640

(Optional) Documentation Update

Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>

github-actions · 2025-06-20T01:17:13Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

gemini-code-assist

Summary of Changes

Hello @22quinn, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new RPC method to dynamically update model loading configurations within the GPU and TPU worker components. This enhancement is crucial for enabling the ability to initialize the engine with placeholder weights and subsequently load real model weights in-place, significantly improving the flexibility and efficiency of model management.

Highlights

New RPC Method for Configuration Update: I've introduced an update_load_config method within both GPUModelRunner and TPUModelRunner. This method leverages dataclasses.replace to allow dynamic modification of the load_config attribute, enabling in-place updates to model loading parameters.
API Exposure for Dynamic Loading: To expose this new functionality, I've added corresponding update_load_config proxy methods to GPUWorker and TPUWorker. These methods delegate calls to their respective model_runner instances, making the dynamic configuration update accessible from the worker level.
Enabling In-Place Weight Loading: This pull request is a critical step towards completing the feature outlined in #18745. It provides the necessary API to update load configurations, which is essential for the workflow of initializing an engine with dummy weights and then loading real weights in-place without restarting the engine.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces an update_load_config method to both GPU and TPU model runners and their respective workers. This change is crucial for enabling a workflow where an engine can be initialized with dummy weights, its loading configuration updated, and then real weights loaded in-place, as described in the PR. The implementation leverages dataclasses.replace for updating the load_config attribute, which is a clean and appropriate approach.

The main area for improvement is the addition of docstrings to the new public methods to enhance code clarity and maintainability. Otherwise, the changes are well-focused and appear correct.

gemini-code-assist · 2025-06-20T01:18:07Z

vllm/v1/worker/gpu_model_runner.py

+    def update_load_config(self, **kwargs) -> None:
+        self.load_config = dataclasses.replace(self.load_config, **kwargs)


The new public method update_load_config should have a docstring explaining its purpose, how **kwargs are used to update the load_config, and any side effects or intended usage patterns (e.g., calling load_model afterwards).

According to PEP 257, a docstring is a string literal that occurs as the first statement in a module, function, class, or method definition. Such a docstring becomes the __doc__ special attribute of that object.

For example:

def update_load_config(self, **kwargs) -> None: """Updates the loading configuration for the model. This method allows modifying the `self.load_config` attribute by providing new values for its fields as keyword arguments. It's typically used before a subsequent call to `load_model` to apply the new configuration. Args: **kwargs: Keyword arguments corresponding to fields in `self.load_config` to be updated. """ self.load_config = dataclasses.replace(self.load_config, **kwargs)

gemini-code-assist · 2025-06-20T01:18:08Z

vllm/v1/worker/gpu_worker.py

+    def update_load_config(self, **kwargs) -> None:
+        self.model_runner.update_load_config(**kwargs)


Please add a docstring to the new public method update_load_config. It should explain that this method delegates the update to the underlying model_runner and what its purpose is in the context of the worker.

For example:

def update_load_config(self, **kwargs) -> None: """Updates the loading configuration in the model runner. This method forwards the provided keyword arguments to the `update_load_config` method of the `self.model_runner`. Args: **kwargs: Keyword arguments to update the load configuration. """ self.model_runner.update_load_config(**kwargs)

gemini-code-assist · 2025-06-20T01:18:08Z

vllm/v1/worker/tpu_model_runner.py

+    def update_load_config(self, **kwargs) -> None:
+        self.load_config = dataclasses.replace(self.load_config, **kwargs)


Similar to the GPU model runner, this new public method update_load_config should include a docstring. It should describe its function of updating self.load_config using **kwargs and its role in the model loading process.

For example:

def update_load_config(self, **kwargs) -> None: """Updates the loading configuration for the model on TPU. This method allows modifying the `self.load_config` attribute by providing new values for its fields as keyword arguments. This is typically used to change loading parameters before re-loading model weights. Args: **kwargs: Keyword arguments corresponding to fields in `self.load_config` to be updated. """ self.load_config = dataclasses.replace(self.load_config, **kwargs)

gemini-code-assist · 2025-06-20T01:18:08Z

vllm/v1/worker/tpu_worker.py

+    def update_load_config(self, **kwargs) -> None:
+        self.model_runner.update_load_config(**kwargs)


A docstring should be added to this new public method update_load_config. It should clarify that the call is delegated to the model_runner and its purpose within the TPU worker's lifecycle.

For example:

def update_load_config(self, **kwargs) -> None: """Updates the loading configuration in the TPU model runner. This method forwards the provided keyword arguments to the `update_load_config` method of the `self.model_runner`. Args: **kwargs: Keyword arguments to update the load configuration. """ self.model_runner.update_load_config(**kwargs)

22quinn · 2025-07-14T07:10:01Z

close in favor of #20095

Add update_load_config RPC

9b56a6f

Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>

22quinn requested review from WoosukKwon, robertgshaw2-redhat, njhill, ywang96, comaniac and alexm-redhat as code owners June 20, 2025 01:17

gemini-code-assist bot reviewed Jun 20, 2025

View reviewed changes

mergify bot added v1 tpu labels Jun 20, 2025

gemini-code-assist bot reviewed Jun 20, 2025

View reviewed changes

22quinn requested review from yeqcharlotte and houseroad June 20, 2025 01:18

22quinn mentioned this pull request Jun 20, 2025

[Doc] Add inplace weights loading example #19640

Merged

4 tasks

22quinn requested a review from youkaichao June 20, 2025 01:21

22quinn mentioned this pull request Jun 20, 2025

[RFC]: In-place weights loading and model swapping #19886

Open

1 task

22quinn closed this Jul 14, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[Core] Add `update_load_config` RPC method #19884

[Core] Add `update_load_config` RPC method #19884

Uh oh!

22quinn commented Jun 20, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jun 20, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jun 20, 2025

Uh oh!

gemini-code-assist bot Jun 20, 2025

Uh oh!

gemini-code-assist bot Jun 20, 2025

Uh oh!

gemini-code-assist bot Jun 20, 2025

Uh oh!

22quinn commented Jul 14, 2025

Uh oh!

Uh oh!

		def update_load_config(self, **kwargs) -> None:
		self.load_config = dataclasses.replace(self.load_config, **kwargs)

		def update_load_config(self, **kwargs) -> None:
		self.model_runner.update_load_config(**kwargs)

Uh oh!

[Core] Add update_load_config RPC method #19884

[Core] Add update_load_config RPC method #19884

Uh oh!

Conversation

22quinn commented Jun 20, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jun 20, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

22quinn commented Jul 14, 2025

Uh oh!

Uh oh!

[Core] Add `update_load_config` RPC method #19884

[Core] Add `update_load_config` RPC method #19884

22quinn commented Jun 20, 2025 •

edited by github-actions bot

Loading