Apple Silicon native run manager

[clockworksquirrel]

Summary

This PR adds an experimental native compute-provider path for macOS Apple Silicon, plus documentation and guardrails for using a native client when the standard Linux/NVIDIA CUDA Docker path cannot expose the local Apple Silicon MPS accelerator.

The main experimental target in this PR is Apple Silicon on macOS using PyTorch MPS. It also adds an experimental Windows ARM64 helper for CPU-only smoke/development runs. Windows DirectML, NPU, and GPU acceleration are not wired in here.

What changed

• Added native run-manager support for launching a local psyche-solana-client binary through --native-client.
• Added a macOS Apple Silicon build helper:
  • scripts/build-native-silicon.sh
  • builds run-manager and psyche-solana-client
  • configures PyTorch/libtorch environment variables
  • patches macOS rpaths so binaries can find PyTorch dynamic libraries
• Added an experimental Windows ARM64 CPU-only build helper:
  • scripts/build-native-windows-arm64.ps1
  • targets aarch64-pc-windows-msvc
  • verifies native Windows ARM64, Rust target installation, Python architecture, and torch import
  • emits a reusable PowerShell runtime env script for PyTorch DLL lookup
• Hardened non-CUDA model setup:
  • avoids CUDA-only assumptions on MPS/CPU
  • makes Liger lazy/CUDA-only
  • avoids CUDA torch.compile on non-CUDA devices
  • downgrades flash_attention_2 to sdpa on non-CUDA devices
• Added MPS BF16 probing and overrides:
  • PSYCHE_MPS_BF16=1 forces BF16
  • PSYCHE_MPS_BF16=0 forces float16
  • default behavior probes BF16 support and falls back when needed
• Optimized HfAuto repo-file checkpoint loading:
  • indexes safetensors metadata without loading all tensors into a CPU state_dict
  • streams tensor payloads shard-by-shard
  • avoids mutating caller-owned in-memory state_dict inputs
  • adds clearer shape/copy errors and unexpected-key checks
  • handles tied-weight aliases for common HF embedding/head ties
  • keeps DTensor behavior honest: DTensor paths still materialize each source tensor before redistribution
• Added native-mode compatibility checks:
  • rejects unsupported Torchtitan/non-CUDA paths
  • rejects ephemeral checkpoint runs for native mode
  • enforces single local rank for non-CUDA devices
  • preserves coordinator client-version checks, with an explicit override only for run-admin-approved cases
• Gated tikv-jemallocator off on Windows so the Solana client can build on Windows/MSVC targets.
• Added native silicon docs and linked them from end-user docs.
• Updated README to identify this as an Apple Silicon native fork and to summarize Apple Silicon/Windows ARM64 scope.

Platform support in this PR

Platform	Status	Notes
macOS Apple Silicon	Experimental native path	Uses native run-manager/client binaries with PyTorch MPS for local acceleration. Single local rank only. Sustained runs require burn-in on the exact model, sequence length, batch shape, dtype, and coordinator configuration.
Windows ARM64	Experimental CPU-only	Builds native Windows ARM64 binaries for --device cpu. DirectML, NPUs, and Windows GPU acceleration are not implemented.
Linux NVIDIA CUDA	Existing upstream path	Docker + CUDA remains the production compute-provider path.
Other accelerators	Not supported here	ROCm, Vulkan, DirectML, NPUs, and non-macOS MPS are outside this PR.

Apple Silicon sizing note

Single-rank means one local worker process, not an automatic short-duration cap. On high-memory Apple Silicon systems, some 12B-class Hugging Face checkpoints may load and may complete limited local train/verify steps when the exact run uses modest sequence lengths, small micro-batches, gradient checkpointing, and BF16 or float16 weights. This is not a general 12B support guarantee.

Loading larger 20B-30B checkpoints on the same machine is not evidence that live training or verification will fit. Sustained runs also use memory for activations, gradients, optimizer state, logits, allocator overhead, and any unsupported MPS operations that fall back to CPU.

Documentation

New primary doc:

• psyche-book/src/enduser/native-silicon.md

Also touched:

• README native fork summary
• compute-provider quickstart link
• join-run wallet guidance
• run-manager README

The docs include:

• Apple Silicon build/run instructions
• Windows ARM64 CPU-only build/run instructions
• Solana wallet/keypair and funding requirements
• native-mode caveats and run-admin/version compatibility warnings

Validation

Local checks run on macOS Apple Silicon:

git diff --check
black --check python/python/psyche/models/hf_transformers.py
prettier --check README.md psyche-book/src/enduser/native-silicon.md
python3 -m py_compile python/python/psyche/models/hf_transformers.py
PYTHON_SYS_EXECUTABLE=/opt/homebrew/bin/python3 \
LIBTORCH_USE_PYTORCH=1 \
LIBTORCH_BYPASS_VERSION_CHECK=1 \
cargo check -p psyche-solana-client --no-default-features
PYTHON_SYS_EXECUTABLE=/opt/homebrew/bin/python3 \
LIBTORCH_USE_PYTORCH=1 \
LIBTORCH_BYPASS_VERSION_CHECK=1 \
cargo check -p run-manager
PYTHON_SYS_EXECUTABLE=/opt/homebrew/bin/python3 \
scripts/build-native-silicon.sh --python

A focused Python helper smoke test also covered safetensors indexing, tied-alias helper behavior, and shape-mismatch error reporting.

The Windows ARM64 helper has been statically reviewed and mechanically checked where possible, but still needs target-host validation on a real Windows ARM64 machine before treating that path as proven.

Caveats

• This does not make Apple Silicon equivalent to a CUDA provider. MPS support is narrower and some PyTorch operations may fall back to CPU.
• BF16 availability on MPS depends on the local macOS, PyTorch, and hardware stack; float16 fallback remains necessary for some environments.
• Native mode should only join coordinator runs where the run administrator accepts this client/fork/version.
• --native-client-version should only be used after the run administrator confirms the accepted version; it should not be used just to bypass the compatibility gate.
• Large models, long contexts, or large vocabularies may exceed unified memory or be too slow to be useful on Apple Silicon.