misc: speedup load safetensors (#1319)

Co-authored-by: ispobock <ISPObaoke@163.com>
2024-09-04 04:29:53 +10:00
parent 1e495e0847
commit dc67d97693
1 changed files with 1 additions and 0 deletions
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -162,6 +162,7 @@ class ModelRunner:
        return min_per_gpu_memory

    def load_model(self):
+        torch.set_num_threads(1)
        logger.info(
            f"Load weight begin. avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
        )