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

Add update_weights_from_tensor (#2631)

Browse Source
This commit is contained in:
fzyzcjy
2024-12-29 05:30:27 +08:00
committed by GitHub
parent 7863e4368a
commit fd28640dc5
10 changed files with 120 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/managers/tp_worker_overlap_thread.py
View File

@@ -28,6 +28,7 @@ from sglang.srt.managers.io_struct import (
InitWeightsUpdateGroupReqInput,
UpdateWeightFromDiskReqInput,
UpdateWeightsFromDistributedReqInput,
UpdateWeightsFromTensorReqInput,
)
from sglang.srt.managers.schedule_batch import ModelWorkerBatch
from sglang.srt.managers.tp_worker import TpModelWorker
@@ -225,6 +226,10 @@ class TpModelWorkerClient:
success, message = self.worker.update_weights_from_distributed(recv_req)
return success, message
def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
success, message = self.worker.update_weights_from_tensor(recv_req)
return success, message
def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
return self.worker.get_weights_by_name(recv_req)