Unverified Commit 42a45df0 authored by XinyuanTong's avatar XinyuanTong Committed by GitHub
Browse files

[Fix] `self.worker` assignment in `TpModelWorker` and refactor references (#4788)


Signed-off-by: default avatarXinyuan Tong <justinning0323@outlook.com>
parent 04eb6062
...@@ -1703,18 +1703,12 @@ class Scheduler( ...@@ -1703,18 +1703,12 @@ class Scheduler(
def save_remote_model(self, params): def save_remote_model(self, params):
url = params["url"] url = params["url"]
if isinstance(self.tp_worker, TpModelWorkerClient): worker = self.tp_worker.worker
worker = self.tp_worker.worker
else:
worker = self.tp_worker
worker.model_runner.save_remote_model(url) worker.model_runner.save_remote_model(url)
def save_sharded_model(self, params): def save_sharded_model(self, params):
if isinstance(self.tp_worker, TpModelWorkerClient): worker = self.tp_worker.worker
worker = self.tp_worker.worker
else:
worker = self.tp_worker
worker.model_runner.save_sharded_model( worker.model_runner.save_sharded_model(
path=params["path"], path=params["path"],
......
...@@ -53,6 +53,8 @@ class TpModelWorker: ...@@ -53,6 +53,8 @@ class TpModelWorker:
req_to_token_pool: Optional[ReqToTokenPool] = None, req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[TokenToKVPoolAllocator] = None, token_to_kv_pool_allocator: Optional[TokenToKVPoolAllocator] = None,
): ):
self.worker = self
# Parse args # Parse args
self.tp_rank = tp_rank self.tp_rank = tp_rank
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment