Only stream output on tp rank 0 (#2124)

This commit is contained in:
Lianmin Zheng
2024-11-22 15:13:44 -08:00
committed by GitHub
parent c35cd1f8c7
commit e1b63624d7
2 changed files with 7 additions and 5 deletions

View File

@@ -179,7 +179,7 @@ class ModelRunner:
if self.device == "cuda":
torch.cuda.set_device(self.gpu_id)
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
elif self.device == "xpu":
torch.xpu.set_device(self.gpu_id)