feat: auto-detect Apple Silicon (MPS) and keep Triton CUDA-only

Extends the --device auto resolution from openai#17 to include Apple Silicon
(MPS) so Mac users get GPU acceleration by default instead of falling
back to CPU.

Two coordinated changes make this safe:

1. opf/_common/device.py — "auto" now picks cuda > mps > cpu. Each
   fallback emits an info line on stderr so the user always knows which
   backend was selected.

2. opf/_model/model.py — the Triton-backed MoE kernels are CUDA-only
   (Triton does not target Metal). Previously the default enabled Triton
   on any non-CPU device, so trying mps crashed once the MoE layer was
   hit. Narrow the auto-enable to device.type == "cuda"; mps and cpu
   both fall back to the torch-ops path unless the user explicitly sets
   OPF_MOE_TRITON=1.

3. opf/_train/runner.py — mirror the same CUDA-only gate when setting
   OPF_MOE_TRITON=1 on behalf of the user (previously set it for any
   non-CPU device, which would silently enable Triton on mps).

4. opf/_cli/common.py — expand --device help text to list the full
   backend order (cuda > mps > cpu).

Verified on macOS (Apple Silicon, Python 3.14, torch 2.11):

- resolve_device("auto") → mps (with stderr info line)
- resolve_device("mps")  → mps
- resolve_device("cpu")  → cpu
- resolve_device("cuda") → returns cuda device (still fails loudly at
  tensor alloc when the user explicitly asks for it — unchanged)

Low-level MPS op sanity check passed for embedding, attention-like
matmul/softmax, log_softmax, topk, argsort, bincount — all the ops
the inference path relies on.

Fixes openai#21

Author: Berkkirik

opf/_cli/common.py

@@ -55,8 +55,13 @@ def add_device_arg(parser: object) -> None:
     parser.add_argument(
         "--device",
         type=str,
-        default="cuda",
-        help="Device to run on",
+        default="auto",
+        help=(
+            "Device to run on. 'auto' (default) picks the best available "
+            "backend: cuda > mps (Apple Silicon) > cpu. Pass an explicit "
+            "value like 'cuda', 'mps', or 'cpu' to override or to get a "
+            "loud error when the requested backend is unavailable."
+        ),
     )
 
 

opf/_common/device.py

@@ -0,0 +1,51 @@
+"""Device-name resolution helpers shared by CLI entrypoints."""
+
+from __future__ import annotations
+
+import sys
+
+import torch
+
+AUTO_DEVICE: str = "auto"
+
+
+def _mps_is_available() -> bool:
+    """Return True when the current PyTorch build supports Apple Metal (MPS)."""
+    backend = getattr(torch.backends, "mps", None)
+    if backend is None:
+        return False
+    is_available = getattr(backend, "is_available", None)
+    if is_available is None:
+        return False
+    try:
+        return bool(is_available())
+    except Exception:
+        return False
+
+
+def resolve_device(device_name: str) -> torch.device:
+    """Resolve a user-supplied device name into a concrete ``torch.device``.
+
+    ``"auto"`` selects the best available device in this order: CUDA (NVIDIA
+    GPU) > MPS (Apple Silicon GPU) > CPU. Any other value is passed through
+    to ``torch.device`` as-is so that explicit requests like ``"cuda"`` or
+    ``"mps"`` still fail loudly when the underlying backend is unavailable.
+    """
+    if device_name == AUTO_DEVICE:
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        if _mps_is_available():
+            print(
+                "info: no CUDA device detected; using Apple Metal (MPS).",
+                file=sys.stderr,
+                flush=True,
+            )
+            return torch.device("mps")
+        print(
+            "info: no CUDA or MPS device detected; falling back to CPU "
+            "(pass --device cuda or --device mps to override).",
+            file=sys.stderr,
+            flush=True,
+        )
+        return torch.device("cpu")
+    return torch.device(device_name)

opf/_model/model.py

@@ -750,8 +750,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         expert_indices = experts.indices
         expert_weights = expert_weights / self.experts_per_token
         experts_per_token_eff = self.experts_per_token
-        not_running_on_cpu = t.device.type != "cpu"
-        use_triton = get_env_bool("OPF_MOE_TRITON", default=not_running_on_cpu)
+        # Triton kernels are CUDA-only; auto-enable only on CUDA devices. MPS
+        # and CPU fall back to the torch-ops path unless the user explicitly
+        # opts in via OPF_MOE_TRITON=1.
+        is_cuda_device = t.device.type == "cuda"
+        use_triton = get_env_bool("OPF_MOE_TRITON", default=is_cuda_device)
         if use_triton:
             _require_triton()
 

opf/_train/runner.py

@@ -590,9 +590,10 @@ def main(argv: Sequence[str] | None = None, *, prog: str | None = None) -> int:
     checkpoint = resolve_checkpoint_path(args.checkpoint)
     device = torch.device(args.device)
 
-    # Default to Triton-backed MoE kernels on non-CPU devices unless callers
-    # explicitly opt out. CPU uses torch ops by default so Triton stays optional.
-    if device.type != "cpu":
+    # Default to Triton-backed MoE kernels on CUDA devices unless callers
+    # explicitly opt out. CPU and MPS use torch ops by default so Triton
+    # stays CUDA-only (the kernels don't run on Metal).
+    if device.type == "cuda":
         os.environ.setdefault("OPF_MOE_TRITON", "1")
 
     base_config = _load_checkpoint_config(checkpoint)