xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-09 05:47:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLHF] Fix torch.dtype not serializable in example (#22158)

Browse Source 
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>
This commit is contained in:
22quinn 2025-08-03 19:43:33 -07:00 committed by GitHub
parent e27d25a0dc
commit 845420ac2c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
examples/offline_inference/rlhf.py
View File

@ -126,7 +126,10 @@ for name, p in train_model.named_parameters():
# Synchronize the updated weights to the inference engine. # Synchronize the updated weights to the inference engine.
for name, p in train_model.named_parameters(): for name, p in train_model.named_parameters():
handle = llm.collective_rpc.remote("update_weight", args=(name, p.dtype, p.shape)) dtype_name = str(p.dtype).split(".")[-1]
handle = llm.collective_rpc.remote(
"update_weight", args=(name, dtype_name, p.shape)
)
model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream()) model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
ray.get(handle) ray.get(handle)

3
examples/offline_inference/rlhf_utils.py
View File

@ -45,7 +45,8 @@ class WorkerExtension:
self.device, self.device,
) )
def update_weight(self, name, dtype, shape): def update_weight(self, name, dtype_name, shape):
dtype = getattr(torch, dtype_name)
weight = torch.empty(shape, dtype=dtype, device="cuda") weight = torch.empty(shape, dtype=dtype, device="cuda")
self.model_update_group.broadcast( self.model_update_group.broadcast(
weight, src=0, stream=torch.cuda.current_stream() weight, src=0, stream=torch.cuda.current_stream()