Skip to content

Avoid probing MPS twice after automatic device selection fails - #256

Open
gezerui888 wants to merge 1 commit into
jaredpalmer:mainfrom
gezerui888:fix/avoid-duplicate-mps-probe
Open

gezerui888 wants to merge 1 commit into
jaredpalmer:mainfrom
gezerui888:fix/avoid-duplicate-mps-probe

Conversation

@gezerui888

Copy link
Copy Markdown

On Apple Silicon, select("auto") already probes MPS. If that child fails or times out, the fallback probes MPS again before selecting CPU. This repeats the same failure and can double the wait.

I have the child announce its selected device before allocating a tensor. When a failed automatic probe reports MPS, the parent goes directly to CPU. If the child fails before reporting a device, the existing MPS fallback remains. The announcement uses a separate marker from the success result, so it cannot make a failed probe look successful.

Fixes #255.

Test plan

  • Python CI lane from .github/workflows/ci.yml: 598 passed, 26 skipped.
  • Final pytest tests/test_device.py tests/test_conventions.py -q: 91 passed, including regression cases for missing output, CUDA fallback, MPS errors, timeouts and signals.
  • Apple Silicon reproduction with torch.float8_e4m3fn: one failed MPS probe before CPU fallback, instead of two.

@gezerui888

Copy link
Copy Markdown
Author

Hi @jaredpalmer, could you review my fix for the repeated MPS probe when you have a moment? I would also appreciate approval for the pending CI workflow.

@gezerui888

Copy link
Copy Markdown
Author

Hi @jaredpalmer, I would appreciate a review when you have time. Could you also approve the pending CI workflow?

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

1 participant