support vllm

src/llmtuner/extras/misc.py

@@ -157,6 +157,12 @@ def get_current_device() -> torch.device:
 
 
 def get_device_count() -> int:
+    r"""
+    Gets the number of available GPU devices.
+    """
+    if not torch.cuda.is_available():
+        return 0
+
     return torch.cuda.device_count()
 
 

src/llmtuner/extras/packages.py

@@ -51,3 +51,7 @@ def is_unsloth_available():
 
 def is_uvicorn_available():
     return _is_package_available("uvicorn")
+
+
+def is_vllm_available():
+    return _is_package_available("vllm")