introduce device_context to simplify code.

Author: ykaitao

unsloth/device_type.py

@@ -20,6 +20,10 @@
     "DEVICE_COUNT",
     "ALLOW_PREQUANTIZED_MODELS",
     "ALLOW_BITSANDBYTES",
+    "DeviceContext",
+    "device_context",
+    "clean_gpu_cache",
+    "get_current_device",
 ]
 
 import torch
@@ -125,3 +129,56 @@ def get_device_count():
                 Params4bit
             ):
                 ALLOW_PREQUANTIZED_MODELS = False
+
+
+class DeviceContext:
+    """Encapsulates device-specific operations for XPU/HIP/CUDA."""
+
+    def __init__(self, device_type: str = DEVICE_TYPE) -> None:
+        if device_type not in ("cuda", "hip", "xpu"):
+            raise ValueError(f"Unsloth: Unsupported device type: {device_type}")
+        self.device_type = device_type
+        # Cache the torch module for this device
+        self.torch_module = torch.xpu if device_type == "xpu" else torch.cuda
+
+    def get_stats(self) -> tuple[str, str, float]:
+        """Return (name, stats_snippet, max_memory_gb)."""
+        gpu_stats = self.torch_module.get_device_properties(0)
+        max_mem = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
+
+        # Device name
+        name = gpu_stats.name + ". " if gpu_stats.name else self._get_default_name()
+
+        # Toolkit snippet
+        snippet = self._get_toolkit_snippet(gpu_stats)
+
+        return name, snippet, max_mem
+
+    def _get_default_name(self) -> str:
+        """Get default device name when props.name is empty."""
+        names = {"xpu": "Intel XPU", "cuda": "NVIDIA GPU", "hip": "AMD GPU"}
+        return names[self.device_type] + " Device. "
+
+    def _get_toolkit_snippet(self, props) -> str:
+        """Get toolkit version snippet."""
+        if self.device_type == "cuda":
+            return f"CUDA: {props.major}.{props.minor}. CUDA Toolkit: {torch.version.cuda}."
+        elif self.device_type == "hip":
+            return f"ROCm Toolkit: {torch.version.hip}."
+        else:  # xpu
+            return f"Intel Toolkit: {torch.version.xpu}."
+
+
+# Singleton instance
+device_context = DeviceContext()
+
+
+# Module-level functions for backward compatibility
+def clean_gpu_cache() -> None:
+    """Clear GPU cache for current device type."""
+    device_context.torch_module.empty_cache()
+
+
+def get_current_device() -> int:
+    """Get current device index."""
+    return device_context.torch_module.current_device()

unsloth/models/llama.py

@@ -52,6 +52,9 @@
     DEVICE_TYPE_TORCH,
     DEVICE_COUNT,
     ALLOW_PREQUANTIZED_MODELS,
+    device_context,
+    clean_gpu_cache,
+    get_current_device,
 )
 
 transformers_version = Version(transformers_version)
@@ -119,13 +122,6 @@
     xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None
 )
 
-if DEVICE_TYPE == "xpu":
-    clean_gpu_cache = torch.xpu.empty_cache
-    get_current_device = torch.xpu.current_device
-else:
-    clean_gpu_cache = torch.cuda.empty_cache
-    get_current_device = torch.cuda.current_device
-
 
 def original_apply_qkv(self, X):
     Q = self.q_proj(X)
@@ -2188,43 +2184,12 @@ def from_pretrained(
             model_patcher = FastLlamaModel
         SUPPORTS_BFLOAT16 = is_bfloat16_supported()
 
-        if DEVICE_TYPE == "cuda":
-            gpu_stats = torch.cuda.get_device_properties(0)
-            gpu_stats_name = (
-                gpu_stats.name + ". " if gpu_stats.name != "" else "NVIDIA GPU Device. "
-            )
-            gpu_version = torch.version.cuda
-            gpu_stats_snippet = f"CUDA: {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit: {gpu_version}."
-            try:
-                vllm_version = f" vLLM: {importlib_version('vllm')}."
-            except:
-                vllm_version = ""
-        elif DEVICE_TYPE == "hip":
-            gpu_stats = torch.cuda.get_device_properties(0)
-            gpu_stats_name = (
-                gpu_stats.name + ". " if gpu_stats.name != "" else "AMD GPU Device. "
-            )
-            gpu_version = torch.version.hip
-            gpu_stats_snippet = f"ROCm Toolkit: {gpu_version}."
-            try:
-                vllm_version = f" vLLM: {importlib_version('vllm')}."
-            except:
-                vllm_version = ""
-        elif DEVICE_TYPE == "xpu":
-            gpu_stats = torch.xpu.get_device_properties(0)
-            gpu_stats_name = (
-                gpu_stats.name + ". " if gpu_stats.name != "" else "Intel XPU Device. "
-            )
-            gpu_version = torch.version.xpu
-            gpu_stats_snippet = f"Intel Toolkit: {gpu_version}."
-            try:
-                vllm_version = f" vLLM: {importlib_version('vllm')}."
-            except:
-                vllm_version = ""
-        else:
-            raise ValueError(f"Unsloth: Unsupported device type: {DEVICE_TYPE}")
+        gpu_stats_name, gpu_stats_snippet, max_memory = _device_ctx.get_stats()
 
-        max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
+        try:
+            vllm_version = f" vLLM: {importlib_version('vllm')}."
+        except:
+            vllm_version = ""
 
         statistics = (
             f"==((====))==  Unsloth {__version__}: Fast {model_patcher.__name__[4:-5]} patching. Transformers: {transformers_version}.{vllm_version}\n"