chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Speed up update_weights_from_tensor (#2695)

Browse Source
This commit is contained in:
fzyzcjy
2025-01-02 18:05:19 +08:00
committed by GitHub
parent 148254d4db
commit 9183c23eca
6 changed files with 48 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/managers/tp_worker.py
View File

@@ -30,7 +30,7 @@ from sglang.srt.managers.schedule_batch import ModelWorkerBatch, global_server_a
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import broadcast_pyobj, set_random_seed
from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed
logger = logging.getLogger(__name__)
@@ -197,7 +197,7 @@ class TpModelWorker:
def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
success, message = self.model_runner.update_weights_from_tensor(
recv_req.name, recv_req.tensor
MultiprocessingSerializer.deserialize(recv_req.serialized_named_tensors)
)
return success, message