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

Signed-off-by: 22quinn <[email protected]>

Author: 22quinn

examples/offline_inference/rlhf.py

@@ -126,7 +126,10 @@ def __init__(self, *args, **kwargs):
 
 # Synchronize the updated weights to the inference engine.
 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())
     ray.get(handle)
 

examples/offline_inference/rlhf_utils.py

@@ -45,7 +45,8 @@ def init_weight_update_group(
             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")
         self.model_update_group.broadcast(
             weight, src=0, stream=torch.cuda.current_stream()