Repository navigation
auto device probe runs MPS once before falling back to CPU - #257
Draft
jaredpalmer wants to merge 1 commit into
Draft
jaredpalmer wants to merge 1 commit into
jaredpalmer wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
--device autoon a machine with no CUDA already asks its probe child to try MPS. When that attempt times out, or the backend rejects the dtype,selectlaunches a second child with the samempsargument. CPU fallback then waits out two probe timeouts, 60 seconds with the defaultKEV_DEVICE_PROBE_TIMEOUTif the child hangs, and the warning names the same failure twice.kev.train,kev.evaluate,kev.benchmark, andkev.experimentall callselect. 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 mpsand 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 printskev-try mpswhen CUDA is absent, andkev-try cuda mpswhen CUDA is present andtorch.backends.mpssays 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, asmps. A child that dies duringimport 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 stillKEV_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 reportskev-try mpsand then times out, or raises the Float8 MPS error, is one probe and one warning clause. A CUDA child that reportskev-try cuda mpsstill 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._probeprints beforetorch.ones. No CUDA and an MPS dtype error. CUDA with no MPS backend. CUDA death with MPS available.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.