[RL] support update_weights_from_distributed with different group and multiple weights (#7292)

2025-07-02 19:29:11 -07:00
parent 09e699bba4
commit 0626f678de
6 changed files with 73 additions and 38 deletions
@@ -259,7 +259,7 @@ class TpModelWorker:
        self, recv_req: UpdateWeightsFromDistributedReqInput
    ):
        success, message = self.model_runner.update_weights_from_distributed(
-            recv_req.name, recv_req.dtype, recv_req.shape
+            recv_req.names, recv_req.dtypes, recv_req.shapes, recv_req.group_name
        )
        return success, message