Add support for Apple Silicon MPS for Torch

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 All notable changes to `ml4t-models` will be documented in this file.
 
+## Unreleased
+
+- Neural configs: `device="mps"` is supported on Apple Silicon when PyTorch MPS is available
+  (with CPU fallback), alongside existing `cpu` / `cuda` handling in `resolve_device`.
+
 ## 0.1.0a0
 
 - Added finance-native data contracts for persistent panels, ragged cross-sections, and

README.md

@@ -62,6 +62,16 @@ pip install ml4t-models[docs]         # mkdocs site build
 pip install ml4t-models[all]
 ```
 
+### Neural compute devices
+
+Torch-backed configs inherit `device` from `BaseModelConfig` (defaults to `cpu`). Supported values are:
+
+- **`cpu`** — always available.
+- **`cuda`** / **`cuda:N`** — used when `torch.cuda.is_available()`; otherwise training falls back to CPU.
+- **`mps`** — Apple Silicon GPU when `torch.backends.mps.is_available()`; otherwise CPU.
+
+Set `device="mps"` or `device="cuda:0"` on the relevant `*Config` (for example `LSTMPortfolioConfig`, `CAEConfig`, `StochasticDiscountFactorConfig`).
+
 ## Quick Start
 
 ### 1. Latent-Factor Forecast Pipeline
@@ -157,12 +167,12 @@ write_backtest_frames("artifacts/run_001", predictions=frame)
 
 These models estimate a structural representation first, then let a separate forecaster produce ex ante factor premia.
 
-| Model | Contract | Native output | Predictive step |
-|---|---|---|---|
-| `PCAModel` | `PersistentPanelBatch` | static loadings, factor returns | factor-premium forecaster + mapper |
-| `RPPCAModel` | `PersistentPanelBatch` | risk-premium-aware latent factors | factor-premium forecaster + mapper |
-| `IPCAModel` | `CrossSectionBatch` | characteristic-implied betas, factor history | factor-premium forecaster + mapper |
-| `CAEModel` | `CrossSectionBatch` | nonlinear characteristic betas, factor history | factor-premium forecaster + mapper |
+| Model        | Contract               | Native output                                  | Predictive step                    |
+| ------------ | ---------------------- | ---------------------------------------------- | ---------------------------------- |
+| `PCAModel`   | `PersistentPanelBatch` | static loadings, factor returns                | factor-premium forecaster + mapper |
+| `RPPCAModel` | `PersistentPanelBatch` | risk-premium-aware latent factors              | factor-premium forecaster + mapper |
+| `IPCAModel`  | `CrossSectionBatch`    | characteristic-implied betas, factor history   | factor-premium forecaster + mapper |
+| `CAEModel`   | `CrossSectionBatch`    | nonlinear characteristic betas, factor history | factor-premium forecaster + mapper |
 
 ### Stochastic Discount Factor
 

docs/getting-started/installation.md

@@ -31,6 +31,18 @@ This extra is required for:
 - `LSTMPortfolioModel`
 - `DeepPortfolioModel`
 
+#### Compute device
+
+Neural model configs include a `device` string (default `cpu`):
+
+| Value               | Behavior                                                               |
+| ------------------- | ---------------------------------------------------------------------- |
+| `cpu`               | CPU tensors.                                                           |
+| `cuda`, `cuda:0`, … | GPU when CUDA is available; otherwise CPU.                             |
+| `mps`               | Apple Metal (MPS) when available in your PyTorch build; otherwise CPU. |
+
+Example: `LSTMPortfolioConfig(..., device="mps")` on Apple Silicon with a suitable `torch` install.
+
 ### Cross-Library Integration
 
 Install tabular and spec helpers:

docs/reference/architecture.md

@@ -98,6 +98,16 @@ ml4t.models
 └── integration/
 ```
 
+## Neural backends and devices
+
+`torch`-based models resolve `device` from config via `ml4t.models._internal.torch_runtime.resolve_device`:
+
+- **`cpu`** — default.
+- **`cuda`** / **`cuda:N`** — when CUDA is available.
+- **`mps`** — when the PyTorch MPS backend is available (typical on Apple Silicon); otherwise CPU.
+
+Unavailable accelerators fall back to CPU so jobs stay runnable in CI or CPU-only environments.
+
 ## Boundary Rules
 
 ### Belongs Here

src/ml4t/models/_internal/torch_runtime.py

@@ -16,12 +16,25 @@ def import_torch() -> Any:
 
 
 def resolve_device(torch: Any, requested: str) -> Any:
-    if requested.startswith("cuda") and torch.cuda.is_available():
-        return torch.device(requested)
+    """Map a config device string to a ``torch.device``, with CUDA/MPS fallbacks to CPU."""
+    raw = requested.strip()
+    lower = raw.lower()
+    if lower.startswith("cuda") and torch.cuda.is_available():
+        return torch.device(raw)
+    mps_backend = getattr(torch.backends, "mps", None)
+    if lower == "mps" or lower.startswith("mps:"):
+        if mps_backend is not None and mps_backend.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
     return torch.device("cpu")
 
 
 def seed_torch(torch: Any, seed: int, device: Any) -> None:
     torch.manual_seed(seed)
-    if getattr(device, "type", "cpu") == "cuda":
+    dev_type = getattr(device, "type", "cpu")
+    if dev_type == "cuda":
         torch.cuda.manual_seed_all(seed)
+    elif dev_type == "mps":
+        mps_manual_seed = getattr(getattr(torch, "mps", None), "manual_seed", None)
+        if callable(mps_manual_seed):
+            mps_manual_seed(seed)

src/ml4t/models/configs/base.py

@@ -7,7 +7,12 @@
 
 @dataclass(frozen=True, slots=True)
 class BaseModelConfig:
-    """Common configuration for ML4T models."""
+    """Common configuration for ML4T models.
+
+    The ``device`` field selects the PyTorch accelerator where applicable:
+    ``cpu`` (default), ``cuda``/``cuda:N`` when CUDA is available, or ``mps`` when the
+    PyTorch MPS backend is available; otherwise training falls back to CPU.
+    """
 
     seed: int = 42
     device: str = "cpu"

tests/test_torch_runtime.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import pytest
+
+pytest.importorskip("torch")
+
+from ml4t.models._internal.torch_runtime import import_torch, resolve_device, seed_torch
+
+
+def test_resolve_device_cpu() -> None:
+    torch = import_torch()
+    assert resolve_device(torch, "cpu").type == "cpu"
+
+
+def test_resolve_device_mps_or_cpu_fallback() -> None:
+    torch = import_torch()
+    d = resolve_device(torch, "mps")
+    mps_backend = getattr(torch.backends, "mps", None)
+    if mps_backend is not None and mps_backend.is_available():
+        assert d.type == "mps"
+    else:
+        assert d.type == "cpu"
+
+
+def test_seed_torch_mps_does_not_raise() -> None:
+    torch = import_torch()
+    mps_backend = getattr(torch.backends, "mps", None)
+    if mps_backend is None or not mps_backend.is_available():
+        pytest.skip("MPS not available")
+    device = torch.device("mps")
+    seed_torch(torch, 123, device)