[RFC]: Add runtime weight update API

[lyuqin-scale]

Motivation.

In online RL training, vLLM can significantly accelerate the rollout stage. To achieve this, we need weight sync from main training process to vLLM worker process, and then call the existing API in vLLM to update the weights by
model_runner.model.load_weights
An example of such implementation can be found in OpenRLHF, https://github.com/OpenLLMAI/OpenRLHF/blob/main/openrlhf/trainer/ray/vllm_worker_wrap.py

However, user has to monkey patch vLLM worker to introduce such behavior. It would be great if vLLM naturally supports weight sync at runtime.

Proposed Change.

1. Add a NCCL-based weight sync process group during vLLM initialization, so that main process can dist.broadcast weight to vLLM worker process later
2. Expose a weight sync API, for example:
   def update_weight(self, name, dtype, shape)

then in master process, user can achieve weight sync via the following (modified from OpenRLHF):

for name, param in model.named_parameters():
    # Fire all vllm engines for broadcast
    if torch.distributed.get_rank() == 0:
        shape = param.shape if self.strategy.args.zero_stage != 3 else param.ds_shape
        refs = [
            engine.update_weight.remote(name, dtype=param.dtype, shape=shape, empty_cache=count == num_params)
            for engine in self.vllm_engines
        ]

        torch.distributed.broadcast(param.data, 0, group=self._model_update_group)
        ray.get(refs)

Feedback Period.

No response

CC List.

No response

Any Other Things.

No response

[youkaichao]

thanks for the information! can you describe, the processes involved, which tensor lives in which device and which process, and what is the desired transfer?

also, cc @hijkzzz from #5477

[lyuqin-scale]

@youkaichao
process: vLLM worker process(es) on GPU 0 and 1, main training process on GPU 2
tensors: HF weights in main training process on GPU2, to be dist.broadcast to temp tensors of same size to vLLM workers on GPU 0 and 1, then within vLLM worker process:
model_runner.model.load_weights(weights=[(name, weight)])
where weight is the temp tensor of one of the weights broadcasted from main process to vLLM worker process

the transfer is preferred to be via NCCL

[hijkzzz]

More importantly, we need to support establishing an NCCL group between DeepSpeed and vLLM engines.

[hijkzzz]

Update:
This API should support LoRA weight updates as much as possible

See:
https://github.com/OpenLLMAI/OpenRLHF/pull/335/files

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[vwxyzjn]

Hi there, has there been an update on this issue?

[ZSL98]

Hi, has there been any progress in this issue?

[youkaichao]

will add it recently.

[ZSL98]

will add it recently.

Appreciate it! Weight updating is indeed important in RL frameworks. It would be convenient if the model_runner on each device can directly update their parameters by names.
BTW, weight offloading is also an essential technique to reduce peak GPU memory usage in RL frameworks. It may also add flexibility if similar APIs such as model_runner.model.offload_weights(name: str]) is provided (https://github.com/volcengine/verl/blob/main/verl/third_party/vllm/vllm_v_0_6_3/worker.py#L283).