[feat] update bucketed weights from distributed (#13824)

Co-authored-by: Stefan He <hebiaobuaa@gmail.com>
2025-11-26 15:30:45 -08:00
parent 082b54c689
commit 5155016b56
5 changed files with 124 additions and 15 deletions
@@ -153,7 +153,11 @@ class BaseTpWorker(ABC):
        self, recv_req: UpdateWeightsFromDistributedReqInput
    ):
        success, message = self.model_runner.update_weights_from_distributed(
-            recv_req.names, recv_req.dtypes, recv_req.shapes, recv_req.group_name
+            recv_req.names,
+            recv_req.dtypes,
+            recv_req.shapes,
+            recv_req.group_name,
+            recv_req.load_format,
        )
        return success, message