Unverified Commit 66d4859a authored by Lianmin Zheng's avatar Lianmin Zheng Committed by GitHub
Browse files

Revert "Only stream output on tp rank 0" (#2130)

parent e1b63624
...@@ -134,8 +134,8 @@ class Scheduler: ...@@ -134,8 +134,8 @@ class Scheduler:
) )
else: else:
self.recv_from_tokenizer = None self.recv_from_tokenizer = None
self.send_to_tokenizer = SimpleNamespace(send_pyobj=lambda _: None) self.send_to_tokenizer = SimpleNamespace(send_pyobj=lambda x: None)
self.send_to_detokenizer = SimpleNamespace(send_pyobj=lambda _: None) self.send_to_detokenizer = SimpleNamespace(send_pyobj=lambda x: None)
# Init tokenizer # Init tokenizer
self.model_config = ModelConfig( self.model_config = ModelConfig(
...@@ -1028,7 +1028,6 @@ class Scheduler: ...@@ -1028,7 +1028,6 @@ class Scheduler:
else: else:
self.tree_cache.cache_unfinished_req(req) self.tree_cache.cache_unfinished_req(req)
if self.tp_rank == 0:
self.stream_output(batch.reqs) self.stream_output(batch.reqs)
def process_batch_result_decode(self, batch: ScheduleBatch, result): def process_batch_result_decode(self, batch: ScheduleBatch, result):
...@@ -1080,7 +1079,6 @@ class Scheduler: ...@@ -1080,7 +1079,6 @@ class Scheduler:
torch.cuda.current_stream().synchronize() torch.cuda.current_stream().synchronize()
batch.next_batch_sampling_info.sampling_info_done.set() batch.next_batch_sampling_info.sampling_info_done.set()
if self.tp_rank == 0:
self.stream_output(batch.reqs) self.stream_output(batch.reqs)
self.token_to_kv_pool.free_group_end() self.token_to_kv_pool.free_group_end()
......
...@@ -179,7 +179,7 @@ class ModelRunner: ...@@ -179,7 +179,7 @@ class ModelRunner:
if self.device == "cuda": if self.device == "cuda":
torch.cuda.set_device(self.gpu_id) torch.cuda.set_device(self.gpu_id)
backend = "nccl" backend = "nccl"
# TODO(liangan1): Just use gloo to bypass the initilization fail # ToDO(liangan1):Just use gloo to bypass the initilization fail
# Need to use xccl for xpu backend in the future # Need to use xccl for xpu backend in the future
elif self.device == "xpu": elif self.device == "xpu":
torch.xpu.set_device(self.gpu_id) torch.xpu.set_device(self.gpu_id)
......
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