A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[RLHF] Fix torch.dtype not serializable in example (#22158)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>
22quinn committed
Aug 3, 2025 at 19:43 UTC
845420ac2c2bc27ae0f96c25430b4f1cd20063cc
2 files changed
+6
-2
examples/offline_inference/rlhf.py
+4
-1
@@ -126,7 +126,10 @@ for name, p in train_model.named_parameters():
126
127
# Synchronize the updated weights to the inference engine.
128
for name, p in train_model.named_parameters():
129
- handle = llm.collective_rpc.remote("update_weight", args=(name, p.dtype, p.shape))
129
+ dtype_name = str(p.dtype).split(".")[-1]
130
+ handle = llm.collective_rpc.remote(
131
+ "update_weight", args=(name, dtype_name, p.shape)
132
+ )
133
model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
134
ray.get(handle)
135
examples/offline_inference/rlhf_utils.py
+2
-1
@@ -45,7 +45,8 @@ class WorkerExtension:
45
self.device,
46
)
47
48
- def update_weight(self, name, dtype, shape):
48
+ def update_weight(self, name, dtype_name, shape):
49
+ dtype = getattr(torch, dtype_name)
50
weight = torch.empty(shape, dtype=dtype, device="cuda")
51
self.model_update_group.broadcast(
52
weight, src=0, stream=torch.cuda.current_stream()