Skip to content

Apple Silicon: MPS device routing + native MLX backend for ESMC - #410

Open
Tyronita wants to merge 8 commits into
Biohub:mainfrom
Tyronita:feat/apple-silicon
Open

Apple Silicon: MPS device routing + native MLX backend for ESMC#410
Tyronita wants to merge 8 commits into
Biohub:mainfrom
Tyronita:feat/apple-silicon

Conversation

@Tyronita

@Tyronita Tyronita commented Sep 22, 2026 •
edited
Loading

Copy link
Copy Markdown

Summary

Combines two complementary Apple Silicon contributions into a single PR:

Layer 1 — PyTorch MPS routing (credit: @eamag, PR #409)

  • esm/utils/device.py: resolve_device("auto") picks CUDA → MPS → CPU
  • get_default_model_dtype(): float32 on MPS, bfloat16 on CUDA
  • Eager attention fallback for non-CUDA devices
  • All from_pretrained calls updated to device="auto"

Layer 2 — Native MLX backend for all ESMC models

  • esm/models/esmc/mlx_model.py: full ESMC forward pass in Apple MLX — no PyTorch, zero CPU↔GPU copies
  • Supports all three published ESMC checkpoints: 300M, 600M, and 6B
    • All three share the same pre-LN transformer architecture; only depth/width differs
    • 6B: 80 layers, hidden=2560, ~12 GB bf16, ~3 GB mxfp4

Layer 3 — Unified factory

  • load_esmc(model_name, device="auto") selects the optimal backend automatically:
    • CUDA present → EsmcForMaskedLM (bfloat16)
    • MLX available on M-chip → EsmcMLX (native Metal, no PyTorch overhead)
    • MPS only → EsmcForMaskedLM (float32, via feat: add Apple Metal (MPS) support #409)
    • CPU → EsmcForMaskedLM
from esm.utils.device import load_esmc

model = load_esmc("biohub/ESMC-300M")          # auto: best backend
model = load_esmc("biohub/ESMC-600M", "mlx")   # explicit MLX
model = load_esmc("biohub/ESMC-6B",   "mlx")   # 6B on Apple Silicon (≥16 GB)
model = load_esmc("biohub/ESMC-6B",   "cuda")  # 6B on CUDA

Why two backends for Apple Silicon?

PyTorch MPS (#409) Native MLX (this PR)
Requires PyTorch Yes No
Memory model CPU↔GPU copies Unified (zero copy)
Kernel scheduling PyTorch dispatcher MLX lazy eval + Metal
L=140 latency (M3) ~68ms ~57ms
L=280 latency (M3) ~148ms ~104ms
L=426 latency (M3) ~233ms ~150ms
Logit MAE vs fp32 ~1e-7 ~1e-4

MPS is the correct fallback when mlx is not installed. MLX is the preferred path when it is.

Model coverage

Model Params MLX supported? Notes
ESMC-300M 300M ✅ ~600 MB bf16
ESMC-600M 600M ✅ ~1.2 GB bf16
ESMC-6B 6B ✅ ~12 GB bf16; recommend opt_level=2 on <24 GB RAM
ESM3 (open small) 1.4B — Different multimodal arch; out of scope for this PR
ESMFold2 — — Structure prediction; not applicable

Correctness

Closes / builds on

Test plan

  • load_esmc("biohub/ESMC-300M") on M-series Mac selects MLX backend
  • load_esmc("biohub/ESMC-300M", "mps") forces MPS path
  • load_esmc("biohub/ESMC-6B", "mlx") loads and runs on ≥16 GB Mac
  • cookbook/snippets/esmc_mlx.py runs end-to-end
  • Existing PyTorch test suite unaffected (mlx is an optional import)

🤖 Generated with Claude Code

Antigravity Assistant and others added 4 commits September 22, 2026 13:05
Antigravity Assistant
feat: add Apple Metal (MPS) support with automatic device resolution …
f183a7b 
…and fallbacks
@Tyronita @claude
Add native MLX backend for ESMC-300M / 600M (Apple Silicon)
3aed976 
Adds esm/models/esmc/mlx_model.py: a self-contained ESMC implementation
in Apple MLX — no PyTorch, no CUDA, unified memory, Metal GPU via
mx.fast.scaled_dot_product_attention.

Architecture matches EsmcForMaskedLM exactly:
  • pre-LN Transformer, RoPE (base=10000), QK-Norm, SwiGLU FFN
  • residue scaling: sqrt(n_layers/36) (ESM3 scheme)
  • weights loaded from the same HF safetensors checkpoint

Correctness: max logit MAE = 1.3e-4 vs PyTorch float32 on CPU.

Benchmark vs PyTorch MPS on Apple M3:
  L=140 : 57 ms vs 68 ms  → 1.19x
  L=280 : 104 ms vs 148 ms → 1.41x
  L=426 : 150 ms vs 233 ms → 1.56x
Speedup grows with sequence length (attention is O(L²)).

Also adds cookbook/snippets/esmc_mlx.py with a usage example.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@Tyronita @claude
mlx: update benchmark table + batch sweep in docstring
b92d645 
Full M3 benchmark table for ESMC-300M across opt levels 0-4 and sequence
lengths L=76-450. B=12 is sweet spot at 1.75x vs sequential (33.7 seq/s).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@Tyronita @claude
device: add is_mlx_available() + load_esmc() factory with MLX routing
e5c1486 
Extends the device utils from PR Biohub#409 with:

- is_mlx_available(): detects Apple MLX (mlx.core importable)
- load_esmc(model_name, device="auto"): top-level factory that selects
  the optimal ESMC backend for the current hardware:
    CUDA            → EsmcForMaskedLM on CUDA  (bfloat16)
    MLX on M-chip   → EsmcMLX          (native Metal, no PyTorch)
    MPS only        → EsmcForMaskedLM on MPS   (float32)
    CPU             → EsmcForMaskedLM on CPU

device="mlx" forces the MLX path explicitly. Both backends share the
same __call__(input_ids) -> logits interface so callers need no
device-specific branching.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@Tyronita Tyronita mentioned this pull request Sep 22, 2026
Open
Tyronita and others added 4 commits September 22, 2026 16:31
@Tyronita @claude
mlx: add ESMC-6B support (same arch, hidden=2560, 80 layers)
eb91cb7 
All three published ESMC checkpoints share the same pre-LN transformer
architecture — only depth and width differ. Adds ESMC-6B to CONFIGS:

  "biohub/ESMC-6B": hidden=2560, n_heads=40, n_layers=80, intermediate=6912

Memory: ~12 GB bf16, ~3 GB mxfp4.  Recommended opt_level=2 on machines
with <24 GB RAM (avoids JIT warmup spike); opt_level=3 on M-series Max/Ultra.

Also updates load_esmc() docstring and from_pretrained error message to
enumerate all three supported repos.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@Tyronita @claude
mlx: update benchmark table with correct isolated subprocess measurem…
d2fefb4 
…ents

Previous table measured MLX in isolation (bench_serial.py) against MPS in
isolation — not directly comparable. New table runs both backends in separate
subprocesses on the same ProteinGym WT sequences (4 warmup + 5 timed, M3 8GB):

  Mean speedup: 3.27× (geomean 3.0×) across L=70–536
  Range: 1.73× at L=161 → 7.93× at L=428

The earlier 14s outlier at L=211 in the same-process benchmark was Metal
compiler contention from MLX + MPS co-resident — not reproducible in isolation
(L=211 isolated = 260ms).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@Tyronita @claude
style: ruff lint + format pass on mlx_model.py and device.py
dcf1c0e 
Brings both files into line with the repo's ruff config (E4/E7/E9/F/I):
- sorted imports (I001)
- split semicolon statements onto separate lines (E702 × 3)
- removed unused B_curr variable (F841)
- reformatted to ruff style (line length, trailing commas, blank lines)

No logic changes.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@Tyronita @claude
fix(ty): resolve remaining type errors in mlx_model.py
31f7b85 
- Narrow list|None unions with `is not None` guards instead of
  `if return_hidden_states` so ty can infer the type
- Add mlx/mlx.** to ty replace-imports-with-any (optional dep absent
  from the CPU lint env, same pattern as flash_attn)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Tyronita