[feat][JAX] Add Apple Silicon MPS support

[bvolpato]

Summary

Add opt-in Apple Silicon support for SkyRL's JAX backend using jax-mps.

Changes

• Add an mps dependency extra for Apple Silicon macOS.
• Restore dependency resolution for the current JAX and Transformers stack.
• Route causal attention through XLA when the jax-mps SDPA patch is active.
• Add MPS-compatible fallbacks for attention-mask and KV-cache updates.
• Preserve dynamic_update_slice boundary behavior in the KV-cache fallback.
• Keep ephemeral JAX sampler syncs in memory instead of writing temporary archives.
• Invalidate the in-memory sampler checkpoint ID after optimizer steps and training checkpoint restores so later snapshot loads cannot reuse stale weights.
• Document setup, supported workflows, and current limitations.
• Add focused backend, generator, attention, cache-boundary, and sampler-checkpoint coverage.

Validation

The checkpoint-restore follow-up has two CPU regressions covering both optimizer-restore settings. Both fail before the fix and pass afterward using actual training and sampler checkpoint round trips. The focused optimizer and ephemeral-checkpoint controls also pass. Required pre-commit hooks pass.

Earlier Linux validation after rebasing onto main at 38ef74d3:

• JAX backend, generator, and attention suites (30 passed)
• DeepSeek CPU suite (9 passed with eight forced XLA CPU devices)
• Sampler checkpoint and optimizer regressions (2 passed)
• Ruff, Black, and Gitleaks
• uv lock --check and git diff --check

The same branch code also passed the documentation type check and production build (66 pages), plus a focused type check with no branch-added diagnostics compared with main, before the final dependency-only rebase.

Earlier branch validation on an M4 Max covered the MPS JAX backend, DeepSeek-V3 against its Hugging Face reference, and Qwen3-0.6B forward, backward, optimizer update, and sampling. Those hardware results predate the latest rebase and fallback fixes; current fixes have CPU regression coverage but have not been rerun on Apple hardware.

Limitations

• jax-mps remains experimental and supports one Apple GPU.
• Some JAX operations remain unsupported by the plugin.
• Causal attention uses the XLA implementation when the jax-mps patch is active.
• KV-cache updates use a dense einsum fallback until batched scatter support lands upstream.
• Buffer-donation warnings may still appear.

---

Note

Medium Risk
Changes shared JAX attention/generation paths and training–sampling weight coherence; incorrect invalidation or MPS fallbacks could affect correctness on all platforms when jax-mps is installed, though regressions are partly covered by new tests.

Overview
Adds opt-in Apple Silicon support for the JAX backend via a new mps extra (jax-mps on macOS arm64), plus docs for JAX_PLATFORMS=mps LoRA training/sampling and an end-to-end Tinker RL smoke path.

Runtime workarounds for jax-mps limits: causal attention (with padding masks) forces XLA SDPA when the jax-mps patch is active; decoding uses MPS-specific KV-cache updates and attention-mask updates instead of vmap(dynamic_update_slice) / .at[...].set.

Sampler weight lifecycle: ephemeral save_weights_for_sampler syncs can skip disk (persist=False) while marking in-memory weights as loaded; optimizer steps and training checkpoint restore clear loaded_checkpoint_id so sampling cannot reuse stale LoRA weights. Dependency overrides widen transformers / ml_dtypes pins beyond Linux-only for the JAX stack on macOS.

Tests cover MPS attention routing, KV-cache fallback, ephemeral sampler checkpoints, and post-restore sampler reload.

^{Reviewed by Cursor Bugbot for commit 8ef4c5a. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces experimental support for Apple Silicon (macOS 14+) in the JAX backend using the jax-mps PJRT plugin. It includes documentation updates, dependency adjustments in pyproject.toml, and custom workarounds in the attention layers and KV cache update logic to bypass current jax-mps limitations. Additionally, it implements ephemeral in-memory sampler checkpoints to avoid redundant disk writes. Feedback on the changes highlights a behavioral discrepancy in the custom MPS KV cache update logic, where negative start indices are incorrectly wrapped instead of clamped to zero as standard JAX dynamic_update_slice does. Suggestions have been provided to simplify this logic and update the corresponding unit tests.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 56f7fd8. Configure here.}