fix: use fp16 dtype for sm75 (#1136)

sgl-project · Aug 17, 2024 · 9208591 · 9208591
1 parent 5d0d40d
commit 9208591
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
@@ -148,6 +148,11 @@ def load_model(self):
             f"[gpu={self.gpu_id}] Load weight begin. "
             f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
         )
+        if torch.cuda.get_device_capability()[0] < 8:
+            logger.info(
+                "Compute capability below sm80 use float16 due to lack of bfloat16 support."
+            )
+            self.server_args.dtype = "float16"
 
         monkey_patch_vllm_dummy_weight_loader()
         device_config = DeviceConfig()