ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

replace torch MPS with MLX

Browse source
This commit is contained in:
Roland Tannous 2026-02-11 16:04:35 +00:00
commit 63c583c54f
3 changed files with 86 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

89
studio/backend/tests/test_utils.py
View file

@ -2,12 +2,12 @@
Tests for utils/hardware and utils/utils — device detection, GPU memory, error formatting.
These tests are designed to pass on ANY platform:
• NVIDIA GPU (CUDA backend)
• Apple Silicon (MPS backend)
• NVIDIA GPU (CUDA backend, requires torch)
• Apple Silicon (MLX backend, requires mlx)
• CPU-only (no GPU at all)
No ML framework is imported at the top level.
Tests that need torch internals for mocking are skipped when torch is unavailable.
Tests that need torch/mlx internals for mocking are skipped when unavailable.
Run with:
cd studio/backend
@ -18,14 +18,21 @@ from unittest.mock import patch, MagicMock
import pytest
# --- Conditional torch import ---
# --- Conditional framework imports ---
try:
import torch
HAS_TORCH = True
except ImportError:
HAS_TORCH = False
try:
import mlx.core as mx
HAS_MLX = True
except ImportError:
HAS_MLX = False
needs_torch = pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not installed")
needs_mlx = pytest.mark.skipif(not HAS_MLX, reason="MLX not installed")
from utils.hardware import (
get_device,
@ -44,11 +51,10 @@ from utils.utils import format_error_message
def _actual_device() -> str:
"""Return the real device string for the current machine."""
if HAS_TORCH:
if torch.cuda.is_available():
return "cuda"
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return "mps"
if HAS_TORCH and torch.cuda.is_available():
return "cuda"
if is_apple_silicon() and HAS_MLX:
return "mlx"
return "cpu"
@ -64,39 +70,36 @@ class TestGetDevice:
"""Tests for get_device() — should agree with the real hardware."""
def setup_method(self):
"""Save DEVICE before each test."""
self._saved_device = _hw_module.DEVICE
def teardown_method(self):
"""Restore DEVICE after each test so mocked tests don't poison later ones."""
_hw_module.DEVICE = self._saved_device
def test_returns_valid_device_type(self):
result = get_device()
assert result in (DeviceType.CUDA, DeviceType.MPS, DeviceType.CPU)
assert result in (DeviceType.CUDA, DeviceType.MLX, DeviceType.CPU)
def test_matches_actual_hardware(self):
assert get_device().value == _actual_device()
# --- Mocked paths (require torch for patching) ---
# --- Mocked paths ---
@needs_torch
def test_returns_cuda_when_cuda_available(self):
with patch("torch.cuda.is_available", return_value=True):
with patch("utils.hardware.hardware._has_torch", return_value=True), \
patch("torch.cuda.is_available", return_value=True):
assert _reset_and_detect() == DeviceType.CUDA
@needs_torch
def test_returns_mps_when_only_mps_available(self):
mock_mps = MagicMock()
mock_mps.is_available.return_value = True
with patch("torch.cuda.is_available", return_value=False), \
patch.object(torch.backends, "mps", mock_mps, create=True):
assert _reset_and_detect() == DeviceType.MPS
def test_returns_mlx_when_on_apple_silicon_with_mlx(self):
with patch("utils.hardware.hardware._has_torch", return_value=False), \
patch("utils.hardware.hardware.is_apple_silicon", return_value=True), \
patch("utils.hardware.hardware._has_mlx", return_value=True):
assert _reset_and_detect() == DeviceType.MLX
@needs_torch
def test_returns_cpu_when_nothing_available(self):
with patch("torch.cuda.is_available", return_value=False), \
patch("builtins.hasattr", side_effect=lambda obj, name: False if name == "mps" else hasattr(obj, name)):
with patch("utils.hardware.hardware._has_torch", return_value=False), \
patch("utils.hardware.hardware.is_apple_silicon", return_value=False), \
patch("utils.hardware.hardware._has_mlx", return_value=False):
assert _reset_and_detect() == DeviceType.CPU
@ -144,15 +147,10 @@ class TestClearGpuCache:
mock_empty.assert_called_once()
mock_ipc.assert_called_once()
@needs_torch
def test_calls_mps_cache_when_mps(self):
mock_mps_module = MagicMock()
mock_mps_module.empty_cache = MagicMock()
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS), \
patch.object(torch, "mps", mock_mps_module, create=True):
def test_mlx_does_not_raise(self):
"""MLX cache clear is a no-op — should just succeed."""
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX):
clear_gpu_cache()
mock_mps_module.empty_cache.assert_called_once()
def test_noop_on_cpu(self):
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU):
@ -215,29 +213,24 @@ class TestGetGpuMemoryInfo:
assert abs(result["free_gb"] - 12.0) < 0.01
assert abs(result["utilization_pct"] - 25.0) < 0.1
# --- MPS-specific mocked test ---
# --- MLX-specific mocked test ---
@needs_torch
def test_mps_path_returns_correct_fields(self):
@needs_mlx
def test_mlx_path_returns_correct_fields(self):
mock_psutil_mem = MagicMock()
mock_psutil_mem.total = 32 * (1024 ** 3)
mock_psutil_mem.total = 32 * (1024 ** 3) # 32 GB unified
mock_psutil = MagicMock()
mock_psutil.virtual_memory.return_value = mock_psutil_mem
mock_mps = MagicMock()
mock_mps.current_allocated_memory.return_value = 2 * (1024 ** 3)
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS), \
patch.dict("sys.modules", {"psutil": mock_psutil}), \
patch.object(torch, "mps", mock_mps, create=True):
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX), \
patch.dict("sys.modules", {"psutil": mock_psutil}):
result = get_gpu_memory_info()
assert result["available"] is True
assert result["backend"] == "mps"
assert result["backend"] == "mlx"
assert "Apple Silicon" in result["device_name"]
assert abs(result["total_gb"] - 32.0) < 0.01
assert abs(result["allocated_gb"] - 2.0) < 0.01
# --- CPU-only path ---
@ -329,11 +322,11 @@ class TestFormatErrorMessage:
assert "big/model" not in msg
assert "model" in msg
# --- OOM on MPS ---
# --- OOM on MLX ---
def test_mps_oom(self):
err = Exception("MPS backend out of memory")
with patch("utils.hardware.get_device", return_value=DeviceType.MPS):
def test_mlx_oom(self):
err = Exception("MLX backend out of memory")
with patch("utils.hardware.get_device", return_value=DeviceType.MLX):
msg = format_error_message(err, "unsloth/huge-model")
assert "Apple Silicon" in msg

73
studio/backend/utils/hardware/hardware.py
View file

@ -25,7 +25,7 @@ logger = logging.getLogger(__name__)
class DeviceType(str, Enum):
"""Supported compute backends. Inherits from str so it serializes cleanly in JSON."""
CUDA = "cuda"
MPS = "mps"
MLX = "mlx"
CPU = "cpu"
@ -39,7 +39,24 @@ DEVICE: Optional[DeviceType] = None
def is_apple_silicon() -> bool:
"""Check if running on Apple Silicon hardware (pure platform check, no ML imports)."""
return platform.system() == "Darwin" and platform.machine() == "arm64"
pass
def _has_torch() -> bool:
"""Check if PyTorch is importable."""
try:
import torch
return True
except ImportError:
return False
def _has_mlx() -> bool:
"""Check if MLX is importable."""
try:
import mlx.core
return True
except ImportError:
return False
def detect_hardware() -> DeviceType:
@ -51,34 +68,31 @@ def detect_hardware() -> DeviceType:
Detection order:
1. CUDA (NVIDIA GPU, requires torch)
2. MPS (Apple Silicon via PyTorch MPS backend)
2. MLX (Apple Silicon via MLX framework)
3. CPU (fallback)
"""
global DEVICE
# --- Try PyTorch first (covers CUDA and MPS) ---
try:
# --- CUDA: try PyTorch ---
if _has_torch():
import torch
if torch.cuda.is_available():
DEVICE = DeviceType.CUDA
device_name = torch.cuda.get_device_properties(0).name
logger.info(f"Hardware detected: CUDA — {device_name}")
return DEVICE
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
DEVICE = DeviceType.MPS
chip = platform.processor() or platform.machine()
logger.info(f"Hardware detected: MPS — Apple Silicon ({chip})")
return DEVICE
except ImportError:
logger.warning("PyTorch not installed — falling back to CPU")
# --- MLX: Apple Silicon ---
if is_apple_silicon() and _has_mlx():
DEVICE = DeviceType.MLX
chip = platform.processor() or platform.machine()
logger.info(f"Hardware detected: MLX — Apple Silicon ({chip})")
return DEVICE
# --- Fallback ---
DEVICE = DeviceType.CPU
logger.info("Hardware detected: CPU (no GPU backend available)")
return DEVICE
pass
# ========== Convenience helpers ==========
@ -92,7 +106,6 @@ def get_device() -> DeviceType:
if DEVICE is None:
detect_hardware()
return DEVICE
pass
def clear_gpu_cache():
@ -109,17 +122,16 @@ def clear_gpu_cache():
import torch
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
elif device == DeviceType.MPS:
import torch
if hasattr(torch.mps, "empty_cache"):
torch.mps.empty_cache()
pass
elif device == DeviceType.MLX:
# MLX manages memory automatically; no explicit cache clear needed.
# mlx.core has no empty_cache equivalent — gc.collect() above is enough.
pass
def get_gpu_memory_info() -> Dict[str, Any]:
"""
Get GPU memory information.
Supports CUDA (NVIDIA), MPS (Apple Silicon), and CPU-only environments.
Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only environments.
"""
device = get_device()
@ -149,13 +161,16 @@ def get_gpu_memory_info() -> Dict[str, Any]:
logger.error(f"Error getting CUDA GPU info: {e}")
return {"available": False, "backend": device.value, "error": str(e)}
# ---- MPS path (Apple Silicon) ----
if device == DeviceType.MPS:
# ---- MLX path (Apple Silicon) ----
if device == DeviceType.MLX:
try:
import torch
import mlx.core as mx
import psutil
allocated = torch.mps.current_allocated_memory() if hasattr(torch.mps, "current_allocated_memory") else 0
# MLX uses unified memory — report system memory as the pool
total = psutil.virtual_memory().total
# MLX doesn't expose per-process GPU allocation; report 0 as allocated
allocated = 0
return {
"available": True,
@ -164,17 +179,16 @@ def get_gpu_memory_info() -> Dict[str, Any]:
"device_name": f"Apple Silicon ({platform.processor() or platform.machine()})",
"total_gb": total / (1024**3),
"allocated_gb": allocated / (1024**3),
"reserved_gb": 0, # MPS doesn't have a separate reserved pool
"reserved_gb": 0,
"free_gb": (total - allocated) / (1024**3),
"utilization_pct": (allocated / total) * 100 if total else 0,
}
except Exception as e:
logger.error(f"Error getting MPS GPU info: {e}")
logger.error(f"Error getting MLX GPU info: {e}")
return {"available": False, "backend": device.value, "error": str(e)}
# ---- CPU-only ----
return {"available": False, "backend": "cpu"}
pass
def log_gpu_memory(context: str):
@ -192,4 +206,3 @@ def log_gpu_memory(context: str):
)
else:
logger.info(f"GPU Memory [{context}]: No GPU available (CPU-only)")
pass

4
studio/backend/utils/utils.py
View file

@ -95,10 +95,10 @@ def format_error_message(error: Exception, model_name: str) -> str:
if "invalid user token" in error_str:
return "Invalid HF token. Please check your token and try again."
if "memory" in error_str or "cuda" in error_str or "mps" in error_str or "out of memory" in error_str:
if "memory" in error_str or "cuda" in error_str or "mlx" in error_str or "out of memory" in error_str:
from utils.hardware import get_device
device = get_device()
device_label = {"cuda": "GPU", "mps": "Apple Silicon GPU", "cpu": "system"}.get(device.value, "GPU")
device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get(device.value, "GPU")
return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory."
# Generic fallback