improve logging & fix vllm version

2024-05-27 14:32:05 -07:00
parent 09de730dee
commit 565d727409
2 changed files with 2 additions and 1 deletions
--- a/python/sglang/srt/managers/router/model_runner.py
+++ b/python/sglang/srt/managers/router/model_runner.py
@@ -235,6 +235,7 @@ class ModelRunner:
        }

        # Init torch distributed
+        logger.info(f"[rank={self.tp_rank}] Set cuda device.")
        torch.cuda.set_device(self.tp_rank)
        logger.info(f"[rank={self.tp_rank}] Init torch begin. Avail mem={get_available_gpu_memory(self.tp_rank):.2f} GB")
        torch.distributed.init_process_group(