mirror of https://github.com/razor-ai/soup.git
137 lines
4.2 KiB
Python
137 lines
4.2 KiB
Python
"""MLX backend utilities — detection + hardware info + batch size estimation.
|
||
|
||
MLX is Apple's ML framework for Apple Silicon (M1-M4 chips). This module
|
||
provides feature detection and helpers so the rest of Soup can opportunistically
|
||
enable MLX training paths without hard-depending on the ``mlx`` package.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import platform
|
||
from typing import Any, Optional
|
||
|
||
|
||
def detect_mlx() -> bool:
|
||
"""Return True if the ``mlx`` package is importable on this machine.
|
||
|
||
This does **not** check for Apple Silicon hardware — use ``get_mlx_info``
|
||
for a full detection report.
|
||
"""
|
||
try:
|
||
import mlx # noqa: F401
|
||
import mlx.core # noqa: F401
|
||
except ImportError:
|
||
return False
|
||
return True
|
||
|
||
|
||
def is_apple_silicon() -> bool:
|
||
"""Return True if the current machine is an Apple Silicon Mac."""
|
||
return platform.system() == "Darwin" and platform.machine() in ("arm64", "aarch64")
|
||
|
||
|
||
def get_mlx_version() -> Optional[str]:
|
||
"""Return the MLX version string, or None if not installed."""
|
||
try:
|
||
import mlx
|
||
except ImportError:
|
||
return None
|
||
return getattr(mlx, "__version__", "unknown")
|
||
|
||
|
||
def get_chip_info() -> dict[str, str]:
|
||
"""Best-effort chip detection for Apple Silicon Macs."""
|
||
info: dict[str, str] = {
|
||
"platform": platform.system(),
|
||
"machine": platform.machine(),
|
||
"processor": platform.processor() or "unknown",
|
||
}
|
||
if not is_apple_silicon():
|
||
return info
|
||
|
||
try:
|
||
import subprocess # noqa: S404 — used with list args only
|
||
result = subprocess.run( # noqa: S603, S607
|
||
["sysctl", "-n", "machdep.cpu.brand_string"],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=2,
|
||
check=False,
|
||
)
|
||
if result.returncode == 0:
|
||
info["chip"] = result.stdout.strip()
|
||
except (OSError, subprocess.SubprocessError):
|
||
pass
|
||
|
||
return info
|
||
|
||
|
||
def get_unified_memory_bytes() -> Optional[int]:
|
||
"""Return total unified memory in bytes, or None if unavailable."""
|
||
if not is_apple_silicon():
|
||
return None
|
||
try:
|
||
import subprocess # noqa: S404
|
||
result = subprocess.run( # noqa: S603, S607
|
||
["sysctl", "-n", "hw.memsize"],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=2,
|
||
check=False,
|
||
)
|
||
if result.returncode == 0:
|
||
return int(result.stdout.strip())
|
||
except (OSError, ValueError, subprocess.SubprocessError):
|
||
return None
|
||
return None
|
||
|
||
|
||
def get_mlx_info() -> dict[str, Any]:
|
||
"""Return a full MLX detection report suitable for ``soup doctor``."""
|
||
available = detect_mlx()
|
||
info: dict[str, Any] = {
|
||
"available": available,
|
||
"version": get_mlx_version(),
|
||
"apple_silicon": is_apple_silicon(),
|
||
"chip": get_chip_info(),
|
||
"unified_memory_bytes": get_unified_memory_bytes(),
|
||
}
|
||
return info
|
||
|
||
|
||
def estimate_mlx_batch_size(
|
||
model_params_b: float,
|
||
unified_memory_bytes: int,
|
||
max_length: int,
|
||
quantization: str = "4bit",
|
||
) -> int:
|
||
"""Rough batch-size estimator for MLX training.
|
||
|
||
Uses the same heuristic as ``utils.gpu.estimate_batch_size`` but scaled
|
||
for Apple Silicon unified memory. Returns at least 1.
|
||
"""
|
||
bytes_per_param = {"4bit": 0.5, "8bit": 1.0, "none": 2.0}.get(quantization, 2.0)
|
||
model_bytes = model_params_b * 1e9 * bytes_per_param
|
||
activation_budget = max(0, unified_memory_bytes - model_bytes * 1.6)
|
||
# Rough: 2 bytes per token × seq × hidden * 4 (q/k/v/o) + grads
|
||
tokens_cost = max_length * 2.0 * 4096 * 4 * 1.5
|
||
if tokens_cost <= 0:
|
||
return 1
|
||
batch = int(activation_budget / tokens_cost)
|
||
return max(1, min(batch, 32))
|
||
|
||
|
||
def load_mlx_model(
|
||
model_path: str, quantization: str = "4bit",
|
||
) -> tuple[Any, Any]:
|
||
"""Thin wrapper around ``mlx_lm.load`` (lazy import).
|
||
|
||
``quantization`` is informational: MLX models are typically already
|
||
quantized at build time (e.g. ``mlx-community/...-4bit``), so this
|
||
parameter is currently advisory and is not forwarded to ``mlx_lm.load``.
|
||
"""
|
||
del quantization # advisory only — MLX models are pre-quantized
|
||
from mlx_lm import load
|
||
|
||
return load(model_path)
|