MPS CI Support by huseyincavusbi · Pull Request #1278 · TransformerLensOrg/TransformerLens

huseyincavusbi · 2026-05-02T16:03:35Z

This PR implements MPS (Metal Performance Shaders) CI Runner Support as proposed in #1264.

The goal is to provide automated testing for the Apple Silicon research community while working within the limits of GitHub's Mac runners.

Key Changes:

New Test Suite: Added tests/mps/test_mps_basic.py with 11 smoke tests covering device detection, core tensor ops on Metal, and HookedTransformer forward passes/caching with small models (TinyStories-1M).
CI Automation: Introduced the mps-checks job in .github/workflows/checks.yml. It uses macos-latest and runs only on PRs/pushes to main.
Memory Management:
- Updated tests/conftest.py to proactively clear the MPS cache after every test using torch.mps.empty_cache().
- Configured the CI to ignore memory-intensive modules (e.g., model_bridge) to ensure stability.
Opt-in Mechanism: Respects TRANSFORMERLENS_ALLOW_MPS=1 to ensure safe defaults for Mac users.

Type of change

New feature (non-breaking change which adds functionality)

Checklist:

I have commented my code, particularly in hard-to-understand areas
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes
I have not rewritten tests relating to key interfaces which would affect backward compatibility

* Fix type of HookedTransformerConfig.device This is typed as `Optional[str]` but sometimes returns `torch.device`. Updated the code to just return the `str` instead of wrapping with a device. I'm not confident that every function which takes a device will always be passed a string, so I didn't change functions like warn_if_mps. Found while working on TransformerLensOrg#1219 * more cleanup * 3.0 CI Bugs (TransformerLensOrg#1261) * Fixing `utils` imports * skip gated notebooks on PR from forks * Updating notebooks * Ensure LLaMA only runs when HF_TOKEN is available --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>

TransformerLens 3.1.0

…imits

Copilot

Pull request overview

Adds Apple Silicon MPS coverage to CI by introducing an MPS-specific test suite and a macOS GitHub Actions job, alongside device-selection tweaks to make MPS opt-in by default.

Changes:

Added a new tests/mps smoke-test suite that validates basic tensor ops and a small HookedTransformer run on MPS.
Added an mps-checks GitHub Actions job on macos-latest to run unit/integration tests plus the new MPS smoke tests on PRs to main and pushes to main.
Updated device utilities and configs to better support MPS opt-in behavior, plus proactive torch.mps.empty_cache() cleanup in pytest fixtures.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
transformer_lens/utilities/devices.py	Adjusts device selection behavior/signatures to support MPS opt-in and updated warning typing.
transformer_lens/train.py	Updates training config typing and default device assignment.
transformer_lens/config/HookedTransformerConfig.py	Uses `get_device()` directly when defaulting `cfg.device`.
tests/unit/utilities/test_devices.py	Updates device utility unit tests for the new `get_device()` return type.
tests/mps/test_mps_basic.py	Adds MPS-only smoke tests covering device detection, core ops, and small-model forward/cache paths.
tests/mps/init.py	Declares the MPS test package.
tests/conftest.py	Adds MPS cache clearing after tests/classes/session to reduce CI OOM risk.
pyproject.toml	Registers a `no_mps` pytest marker.
.github/workflows/checks.yml	Adds the `mps-checks` CI job that runs on macOS and executes the MPS tests.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

jlarson4 · 2026-05-04T15:49:47Z

+def get_device() -> str:
    """Get the best available device, with MPS safety checks.

    MPS is only auto-selected when the environment variable
    ``TRANSFORMERLENS_ALLOW_MPS=1`` is set **and** the installed PyTorch
    version meets or exceeds ``_MPS_MIN_SAFE_TORCH_VERSION``.

    Returns:
-        torch.device: The best available device (cuda, mps, or cpu)
+        str: The best available device name (cuda, mps, or cpu)
    """
    if torch.cuda.is_available():
-        return torch.device("cuda")
+        return "cuda"


Roll back any changes related to get_device, there were some git history issues that caused this to crop up here for you.

I discussed this with another contributor and setup #1230 to standardize the type returned by get_device. There was a typing inconsistency, get_device() -> torch.device only returned a torch.device object in the default case, and returned a string in all other instances, creating a situation where the typing was silently wrong. We standardized everything to str, but for some reason that change never made it to dev, I have integrated it into this branch and it should no longer be a problem.

…ency

huseyincavusbi · 2026-05-02T23:40:44Z

Hi @jlarson4, I've updated the PR to address the automated feedback:

API Stability: Reverted get_device() to return torch.device objects.
Type Checks: Updated type hints across model classes to resolve mypy failures.
CI Trigger: Strictly restricted mps-checks to the main branch

jlarson4 · 2026-05-04T16:00:15Z

+          print(f'MPS built: {torch.backends.mps.is_built()}')
+          assert torch.backends.mps.is_available(), 'MPS not available on this runner!'
+          "
+      - name: MPS Unit Tests


Make sure you're passing $${{ secrets.HF_TOKEN }} to any tests that might load models so we can avoid HuggingFace rate limits where possible

jlarson4 · 2026-05-04T16:10:54Z

+      - name: MPS Unit Tests
+        run: >
+          uv run pytest tests/unit -v
+          --ignore=tests/unit/model_bridge/


We are ignoring a lot of tests here, some should be able to run within our memory limit. Try ignoring only these tests and let me know if you found something while setting this up that necessitates skipping more tests.

- name: MPS Unit Tests run: > uv run pytest tests/unit -v --ignore=tests/unit/model_bridge/test_optimizer_compatibility.py --ignore=tests/unit/model_bridge/test_gpt_oss_moe.py --ignore=tests/unit/model_bridge/test_component_inspection.py --ignore=tests/unit/model_bridge/test_key_analysis.py --ignore=tests/unit/model_bridge/test_benchmark_gated_hooks_fire.py --ignore=tests/unit/model_bridge/test_weight_processing_adapter_paths.py --ignore=tests/unit/model_bridge/test_bridge_generate_kv_cache.py --ignore=tests/unit/model_bridge/test_bridge_vs_hooked_transformer_patching.py --ignore=tests/unit/model_bridge/compatibility/

jlarson4 · 2026-05-04T16:16:54Z

+          --ignore=tests/unit/model_bridge/
+      - name: MPS Integration Tests
+        run: >
+          uv run pytest tests/integration -v


We should still be able to run grouped_query_attention and create_hooked_encoder that are ignored here.

jlarson4 · 2026-05-04T16:17:13Z

-        markers=["slow: marks tests as slow (deselect with '-m \"not slow\"')"]
+        markers=[
+            "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+            "no_mps: marks test as incompatible with MPS device (deselect with '-m \"not no_mps\"')",


If we aren't using this no_mps tag, we should drop it

brendanlong and others added 10 commits April 20, 2026 14:50

Merge pull request TransformerLensOrg#1277 from TransformerLensOrg/dev

6f56518

TransformerLens 3.1.0

feat: Add MPS CI runner support (TransformerLensOrg#1264)

5d412ca

ci: Enable runs on feature branch

7d9d6df

fix: Skip heavy model_bridge unit tests on MPS runner due to memory l…

6453913

…imits

fix: Ignore flaky grouped query attention tests on Mac runner

18b2f57

style: Run make format on test_mps_basic.py

b2787d2

style: Standardize MPS step naming convention in CI

5931009

ci: Revert MPS trigger to run only on main PRs and pushes

be01240

ci: Remove feature branch from global workflow triggers

ed30beb

Copilot AI review requested due to automatic review settings May 2, 2026 16:03

Copilot started reviewing on behalf of huseyincavusbi May 2, 2026 16:04 View session

Copilot AI reviewed May 2, 2026

View reviewed changes

huseyincavusbi added 6 commits May 2, 2026 23:14

fix: Restore torch.device return type in get_device for API stability

58ebdd3

docs: Align train and device config docstrings with implementation

9e816fe

fix: Update device type hints to Union[str, torch.device] for consist…

159bd8d

…ency

ci: Update mps-checks trigger to include dev branch

81d387c

test: Update device tests for torch.device compatibility and robustness

02ae6a0

ci: Restrict MPS trigger to main branch only

341ab71

jlarson4 mentioned this pull request May 4, 2026

Run GitHub CI on MacOS #598

Closed

7 tasks

Merge branch 'dev' into feat/mps-ci-support

4356758

jlarson4 reviewed May 4, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

MPS CI Support#1278

MPS CI Support#1278
huseyincavusbi wants to merge 17 commits intoTransformerLensOrg:devfrom
huseyincavusbi:feat/mps-ci-support

huseyincavusbi commented May 2, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

jlarson4 May 4, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

huseyincavusbi commented May 2, 2026

Uh oh!

jlarson4 May 4, 2026

Uh oh!

jlarson4 May 4, 2026

Uh oh!

jlarson4 May 4, 2026

Uh oh!

jlarson4 May 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

huseyincavusbi commented May 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Key Changes:

Type of change

Checklist:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

jlarson4 May 4, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

huseyincavusbi commented May 2, 2026

Uh oh!

jlarson4 May 4, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 May 4, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 May 4, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 May 4, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

huseyincavusbi commented May 2, 2026 •

edited

Loading