Skip to content

auto device probe runs MPS once before falling back to CPU - #257

Draft
jaredpalmer wants to merge 1 commit into
mainfrom
cursor/auto-mps-probe-once-aaa9
Draft

jaredpalmer wants to merge 1 commit into
mainfrom
cursor/auto-mps-probe-once-aaa9

Conversation

@jaredpalmer

Copy link
Copy Markdown
Owner

--device auto on a machine with no CUDA already asks its probe child to try MPS. When that attempt times out, or the backend rejects the dtype, select launches a second child with the same mps argument. CPU fallback then waits out two probe timeouts, 60 seconds with the default KEV_DEVICE_PROBE_TIMEOUT if the child hangs, and the warning names the same failure twice. kev.train, kev.evaluate, kev.benchmark, and kev.experiment all call select. On a torch build with no MPS backend the second child can only fail, and that unrelated error is added to the warning. Explicit --device mps and a healthy MPS auto probe were already one child. Reported by @luc-pimentel, with the timeout and float8 repro on the issue.

The child prints kev-try <device> and flushes it before the allocation. Auto prints kev-try mps when CUDA is absent, and kev-try cuda mps when CUDA is present and torch.backends.mps says an MPS backend is there. The parent starts a second child only for that second line. The warning names the device the child was about to test, so a failed MPS attempt shows up once, as mps. A child that dies during import torch, before it can print, is not relaunched. The next process would import the same torch.

This depends on the device being chosen before the operation that can hang or crash. torch.cuda.is_available() runs before the line. A hang inside that call, rather than inside the allocation, falls back to CPU after one timeout and does not try MPS. On Apple Silicon the CUDA check returns immediately, and the failures in the report are the allocation timeout and the dtype error, both after the line. A machine where CUDA init hangs and MPS would still have worked is the case this no longer covers. I have not seen that machine. What #231 fixed stays in place: no core file if torch crashes while loading, the probe timeout is still KEV_DEVICE_PROBE_TIMEOUT, and auto still tries the Apple GPU before CPU.

Not in this change: reading MPS availability in the parent, which is the process a driver crash must not kill, and the MLX serve path, which never runs this probe.

Fixes #255.

Test plan

  • tests/test_device.py: an auto child that reports kev-try mps and then times out, or raises the Float8 MPS error, is one probe and one warning clause. A CUDA child that reports kev-try cuda mps still gets one MPS probe, and both failures are named once if that probe fails too. A CUDA child that does not offer MPS is not followed by a second probe. A child that never names a device is not relaunched. Explicit cuda and mps stay one probe and fail closed. A healthy auto result is one child.
  • The same three outcomes with a real probe child whose torch is a stand-in, so the line the parent reads is the one _probe prints before torch.ones. No CUDA and an MPS dtype error. CUDA with no MPS backend. CUDA death with MPS available.
  • CI unit list locally, 606 passed and 26 skipped, plus one rerun. test_resume_is_bit_identical[1-gate0] failed once and passed on the next run. It trains with --device cpu, which returns before a probe. The mismatch was bf16 weights at about 1e-4 across the checkpoint, not a device string.
Open in Web Open in Cursor 

On a machine with no CUDA, the auto child already probes MPS. select then
launched the same probe again, so a hang waited out two timeouts and the
warning named the failure twice. The child now prints the device it is about
to test, and a second child runs only when that line says CUDA died and MPS
is available.

Fixes #255.

Co-authored-by: Jared Palmer <hello@jaredpalmer.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

--device auto on Apple Silicon runs the same MPS probe twice before falling back to CPU

2 participants