[Momoyeyu]

Closes #1355.

The streamed test builders resolved their device as "cuda" if cuda_available() else "cpu" (on main: tests/test_v07204.py:378, tests/test_v07200.py:2501, tests/test_v07202.py:1832, tests/test_issue975_backward_tail_embed_prefetch.py:53). On Apple Silicon the model was built on CPU while TrainingArguments picked mps:0, so the first train() step failed on a device mismatch, and 32 cases carried skipif(_mps_is_the_accelerator(), reason="MPS is untested for layer streaming ...").

What does this PR do?

• tests/conftest.py: mps_is_the_accelerator() (cached, False when CUDA is present) and accelerator_device(), next to cuda_available. accelerator_device() returns the device TrainingArguments picks by itself: cuda, else mps, else cpu.
• All four builder sites call accelerator_device().
• The three module-local _mps_is_the_accelerator copies (test_v07200.py, test_v07202.py, test_v07204.py) are deleted, and the three modules that re-imported one (test_issue975_..., test_issue1011_..., test_issue1049_...) import from conftest, so no test module probes the device itself (~51 CUDA-only tests have no marker, no documented way to run them and no results record; both known runs on real cards found defects (#776, and 9 failures from three defects on H100) #833).
• The MPS skip is removed from the 32 cases: test_v07204.py 11 (test_loss_matches_a_resident_run_of_the_same_loss ×4, test_layer_zero_adapter_receives_gradient ×4, test_kto_streams_at_batch_two, test_a_reference_using_loss_reads_more_layers_than_sft ×2), test_issue1011_streamed_save_guard.py 5, test_issue1049_untied_streamed_preference.py 10, test_issue975_backward_tail_embed_prefetch.py 5 (class-level), test_v07200.py 1 (test_one_training_step_actually_runs).
• Unchanged: both bitsandbytes MPS skips (test_v07202.py's skip_on_mps, now backed by the conftest helper, and test_v07204.py's "bitsandbytes has no 4-bit MPS kernels"), and the not mps_is_the_accelerator() mark on TestStreamedLargeLayerParityMps.

No case needed a strict xfail: all 32 pass on MPS.

No CPU pin for TrainingArguments. I tried pinning the CPU-pinned cases' trainer to CPU (ACCELERATE_USE_CPU) and took it out again: with the pin disabled, all 441 tests in the six files still pass on MPS. None of the device="cpu" bit-exact cases reaches a TrainingArguments step, so the pin would have been code no test exercises. If one ever does, it will fail loudly on macOS with the device-mismatch message rather than pass silently.

Verification

• Mac Studio M3 Ultra (MPS), Python 3.12.10, torch 2.14.0: the 32 cases pass; the six files give 441 passed, 26 skipped (the remaining skips are the CUDA gpu markers and the two bitsandbytes skips). Full suite: 29080 passed / 3 failed / 221 skipped on this branch against 29047 / 4 / 253 on main at the same commit. The skip count drops by exactly 32, and there are no new failures (the 3 failures also fail on main: soup data forge --judge-provider writes an empty-answer row for every failed judge call and still reports "synth complete" #1221 forge judge, Feature proposal: resumable offline AutoDistill pipeline with adaptive data generation #580 MLX venv launcher, test_lambda_submit_interrupt).
• RTX 4090 D (CUDA 13), Python 3.12.14, torch 2.14.0+cu130: the 32 cases pass; pytest -m gpu over the six files gives 20 passed both on this branch and on main, so the CUDA path is unchanged (accelerator_device() returns cuda there, as the old expression did). One failure in the six files, TestShardCacheIdentityBinding::test_changed_source_weights_invalidate_the_cache, is intermittent and not from this change: on the same machine it passes 3/3 in isolation on this branch and fails in a whole-file run on main.

Type of change

• Bug fix
• New feature
• Refactor / cleanup
• Documentation
• Tests

Checklist

• ruff check src/soup_cli/ scripts/ tests/ passes
• pytest tests/ -v passes, apart from the three failures above that also fail on main (Mac Studio full run)
• Updated relevant docs (README.md and the matching page under docs/) if needed: not needed, test-only
• Added a changelog fragment, or this change has no user-visible impact: no user-visible impact, test-only

[MakazhanAlpamys]

@Momoyeyu, thank you, and welcome to Soup. This is the PR #1355 needed, done carefully:

• You ran the 32 cases on real MPS hardware and on CUDA, and reported branch-vs-main counts for both. That is exactly the evidence CI cannot give us, since a macOS cell only tells us that something failed, not why.
• You pushed back on the issue's own proposed fix (pinning TrainingArguments to CPU for the CPU-pinned bit-exact cases) with a reason, instead of adding code no test exercises. The reason holds: I checked every CPU-pinned case (test_v07204.py:491, :509, :1182, :1195, test_issue1049_untied_streamed_preference.py:63, test_v07300.py:167, :203), none of them calls trainer.train(), and none was ever MPS-skipped. They already run and pass on GitHub's macOS runner today (all 11 cases, in #1258's CI log).
• One resolver next to cuda_available in tests/conftest.py, the three private copies gone, and the modules that re-imported a copy now import from conftest, as the #833 rule asks. Both bitsandbytes skips and the MPS-only parity test are untouched.

What I ran (Windows, Python 3.12, torch 2.14.0+cu130, this branch merged with current main; there is no MPS here, so the macOS side rests on your M3 Ultra run and on CI):

• The six files plus tests/test_issue833_gpu_marker.py with the GPU hidden: 445 passed, 22 skipped, identical to main test by test (467 cases, no outcome differs). The 32 cases take about 15 s on CPU.
• The same files with an RTX 5070 visible: 465 passed, 2 skipped on both trees, and accelerator_device() answers cuda there.
• Mutations of tests/conftest.py: accelerator_device() always returning "cpu" fails 40 cases on the GPU box (28 of them among your 32), all on device mismatch; ignoring CUDA fails the same 40. Forcing the MPS probe to True on this Windows box fails 51 cases with "PyTorch is not linked with support for mps devices" and turns on the six bitsandbytes skips, so every builder and every kept skip really does read the resolver.
• GitHub's macOS runner (image macos-26-arm64) passes TestStreamedLargeLayerParityMps on MPS in #1258's CI run today, which is a good sign for the streamed training steps this PR turns on there.

One thing to coordinate: #1258. It adds tests/test_issue1174_head_prefetch_after_lookup.py, which imports _mps_is_the_accelerator from tests.test_v07204 and uses it to skip three classes. The two PRs merge without a textual conflict, but on the combined tree collection stops at:

ImportError: cannot import name '_mps_is_the_accelerator' from 'tests.test_v07204'
!!!!!!!!!!!!!!!!!!! Interrupted: 1 error during collection !!!!!!!!!!!!!!!!!!!!

which would fail every test cell. So whichever of the two lands second needs one more push. If #1258 lands first, please rebase this branch and convert that file like the others: import mps_is_the_accelerator from tests.conftest for its skip or, if its three classes (TestBothTimingsAreBitIdentical, TestTheInstalledDefault, TestTheLoadsMoveButDoNotMultiply) pass on your M3 Ultra, drop _NO_MPS altogether, which is what #1355 had in mind for them.

Also on macOS: besides the 32, 18 setup-only tests that use the builders' default device (for example TestNoSecondModelInstance, TestStreamRuntimeIsReleased and TestNF4EndToEndSetup::test_setup_builds_a_real_trl_trainer_under_nf4) move from CPU to MPS on the macOS cells. They passed in your six-file run, so there is nothing to do. I mention it so that a red macOS cell is read correctly.

Optional, none of these block:

1. A test for the resolver itself. Today its MPS branch only runs on the three macOS cells: putting main's expression back into accelerator_device() leaves the six files green on any machine without MPS (411 passed, 22 skipped here). This test runs everywhere. On your branch it passes (5 passed), and it fails the always-"cpu" mutation (3 failures), the ignore-CUDA one (2) and the revert to main's expression (1). The #833 guard still passes with it in the tree, and a streamed train step run after it in the same pytest run still lands on cuda, so its cache_clear() leaves no stale answer behind. For example, as tests/test_issue1355_accelerator_device.py:

   """#1355 — the streamed builders build on the device ``TrainingArguments`` picks by itself.
   
   Pinned without a Mac or a card: otherwise the MPS branch of the resolver only ever
   runs on the macOS cells, and the CUDA branch only on a GPU box.
   """
   
   from __future__ import annotations
   
   from types import SimpleNamespace
   
   import pytest
   
   from tests import conftest
   
   
   @pytest.fixture()
   def fake_accelerators(monkeypatch):
       """Fake ``torch.cuda`` and ``torch.backends.mps``; both probes are cached, so clear them."""
       import torch
   
       def apply(*, cuda: bool, mps: bool | None) -> None:
           monkeypatch.setattr(torch.cuda, "is_available", lambda: cuda)
           backend = None if mps is None else SimpleNamespace(is_available=lambda: mps)
           monkeypatch.setattr(torch.backends, "mps", backend, raising=False)
           conftest.cuda_available.cache_clear()
           conftest.mps_is_the_accelerator.cache_clear()
   
       yield apply
       # Runs before monkeypatch restores torch, so the next caller re-probes the real one.
       conftest.cuda_available.cache_clear()
       conftest.mps_is_the_accelerator.cache_clear()
   
   
   @pytest.mark.parametrize(
       ("cuda", "mps", "expected"),
       [
           (True, True, "cuda"),
           (True, False, "cuda"),
           (False, True, "mps"),
           (False, False, "cpu"),
           (False, None, "cpu"),  # a torch build with no torch.backends.mps at all
       ],
   )
   def test_the_builders_device_is_cuda_then_mps_then_cpu(fake_accelerators, cuda, mps, expected):
       fake_accelerators(cuda=cuda, mps=mps)
       assert conftest.accelerator_device() == expected
       assert conftest.mps_is_the_accelerator() is (expected == "mps")

2. tests/test_v07202.py:34-37: the NF4 skip_on_mps that stays says "MPS is untested for NF4 streaming (measured on CUDA + CPU only)", which reads like the 32 skips this PR removes. Its real reason, that bitsandbytes has no 4-bit MPS kernels, lived in the docstring of the deleted helper. A one-line comment would stop a future cleanup from removing it.

3. The docstrings at tests/test_v07204.py:238 and :353 still name only two streaming dtypes, bf16 on CUDA and float32 on CPU. On MPS it is bf16 when the runtime accepts it, float32 otherwise.

4. tests/test_issue975_backward_tail_embed_prefetch.py:62-63: two blank lines are left where the class-level pytestmark was.

No changelog fragment is needed, since this only touches tests. The required checks were still queued when I looked, and the three macOS cells are the real check for the 32. Apart from the #1258 ordering above, this now waits only on those checks.

[MakazhanAlpamys]

@Momoyeyu, #1258 is merged (15a66781), so this is the case from my review: it is now the second of the two to land. Please rebase on main and convert tests/test_issue1174_head_prefetch_after_lookup.py like the other files. Import the resolver from tests/conftest.py instead of _mps_is_the_accelerator from tests.test_v07204, which this PR deletes. Without that change the combined tree fails at collection. Once you've pushed, I'll re-check that file on the merged tree, and this merges on green.

[Momoyeyu]

Thanks for checking all of this, including the mutations. #1258 landed first, so I rebased on main and converted tests/test_issue1174_head_prefetch_after_lookup.py. Its three classes pass on the M3 Ultra (14 passed), so _NO_MPS is dropped instead of re-pointed.

Also in the new commit:

• tests/test_issue1355_accelerator_device.py, your resolver test as written. It passes here (5 passed), and your three mutations fail it with 3, 2 and 1 failures, same as on your box.
• A one-line comment on skip_on_mps in test_v07202.py giving the bitsandbytes reason.
• Both dtype docstrings in test_v07204.py now cover MPS: bf16 when mps_supports_bf16() accepts it, float32 otherwise (layer_stream.py:249-251).
• The leftover blank line in test_issue975_backward_tail_embed_prefetch.py.

On the M3 Ultra, the six files, the #833 guard, #1174 and the new test give 481 passed, 27 skipped. ruff is clean.

[MakazhanAlpamys]

@Momoyeyu, thank you, that was quick and thorough. The file from #1258 is converted the better of the two ways: you ran its three classes on the M3 Ultra first and then dropped _NO_MPS rather than re-pointing it, which is what #1355 had in mind for them. The resolver test and all three smaller points are in, and the new dtype docstrings describe what layer_stream.py actually does on MPS.

Some good news from CI, too. The test (macos-latest, 3.11) cell of the run on your previous push finished green. That push has the same first commit, so its log is a real GitHub macOS runner result. All 32 cases passed there, along with the 18 setup tests that now build on the default device. The MPS parity test ran, which confirms MPS was the accelerator on that runner, and the six bitsandbytes cases skipped as they should.

What I ran (Windows, Python 3.12, torch 2.14.0+cu130 with the GPU hidden, on this branch alone and merged with current main):

• The six files, the #833 guard, test_issue1174_head_prefetch_after_lookup.py and the new test: 485 passed, 23 skipped on both trees, identical test by test. test_issue1174_head_prefetch_after_lookup.py itself gives 35 passed, 1 skipped, so the collection error from my earlier review is gone. Collecting the whole tests/ tree on the merged tree also finishes with no error. On main the same files give 480 passed, 23 skipped, and the only difference is the five new cases.

• tests/test_issue1355_accelerator_device.py fails on main (5 failures and 5 errors, module 'tests.conftest' has no attribute 'mps_is_the_accelerator'). Here it catches every resolver mutation I tried:

  • always returning "cpu": 3 failures
  • ignoring CUDA: 2
  • main's expression, i.e. the MPS branch dropped: 1
  • the MPS probe ignoring CUDA: 1

  Before, the third one could only be caught on a Mac.

• Forcing the MPS probe to True on this machine, which has no MPS, fails 67 cases. 65 of them fail with "PyTorch is not linked with support for mps devices", including all 14 of #1174's cases. The six bitsandbytes skips also turn on, so every builder, #1174's included, goes through the resolver.

• The fixture's cleanup: one pytest run of the resolver's cuda=True case followed by two real streamed train steps gives 3 passed. With the teardown's cache_clear() calls removed, both steps fail with "No CUDA GPUs are available". So the fixture leaves nothing stale behind for later tests.

Nothing else is needed from you. GitHub holds workflow runs from a first-time contributor until a maintainer approves them, and it asks again on every push, so the runs for this push are waiting on that. After that, this PR waits only on the required checks. The macOS cells will be the first GitHub run of #1174's three classes on MPS.

[Momoyeyu]

Thank you — a pleasure working with you on this one. Looking forward to the next. :)

[MakazhanAlpamys]

Merged as 26111cf0. Thank you, @Momoyeyu: you tested every round on real MPS hardware before asking for review, and you dropped the MPS skip from #1258's three classes only after they passed on your M3 Ultra, instead of just re-pointing the old helper. Here, the nine files gave 485 passed / 23 skipped on the merged tree, identical to main apart from the five new resolver cases. The resolver test fails on main and catches all four resolver mutations (3, 2, 1 and 1 failures). Forcing the MPS probe on a machine without MPS fails 67 cases, including all 32 and #1174's 14.