[Fix] Resolve GPU Memory Leak in update_weights_from_tensor (#4446)

2025-03-17 16:54:30 +08:00
parent c614dbdf95
commit 91ba98fe50
3 changed files with 42 additions and 18 deletions
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -214,7 +214,7 @@ class TpModelWorker:
    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
        success, message = self.model_runner.update_weights_from_tensor(
            named_tensors=MultiprocessingSerializer.deserialize(
-                recv_req.serialized_named_tensors
+                recv_req.serialized_named_tensors[self.tp_rank]
            ),
            load_format=recv_req.load_format,
        )