Skip to content

Fix identical default random keys across threads - #4857

Open
Chr0mium24 wants to merge 1 commit into
ml-explore:mainfrom
Chr0mium24:fix/thread-random-keys
Open

Chr0mium24 wants to merge 1 commit into
ml-explore:mainfrom
Chr0mium24:fix/thread-random-keys

Conversation

@Chr0mium24

Copy link
Copy Markdown
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: I used Claude Code to find the cause and write the patch and tests, all of which I reviewed and ran. I wrote the issue and this PR description myself.

Proposed changes

Fixes #4856

Since #3278, every thread creates its default random key from the same seed. So all threads that do not call mx.random.seed() or pass key= produce the same random numbers.

Changes

  • Add KeySequence::default_seed() in mlx/random.cpp. It returns a timestamp, computed once per process, plus an atomic counter. Each call returns a different seed.
  • KeySequence::default_() in mlx/random.h creates each thread's key from default_seed().
  • PyKeySequence::state() in python/src/random.cpp uses the same function. C++ and Python default keys share one counter, so they also get different seeds.
  • Move the <chrono> include from mlx/random.h and python/src/random.cpp to mlx/random.cpp.

The two goals of #3278 do not change. The seed is still created only when a key is first needed, so import does not touch the GPU. Each thread still has its own key, and the only shared state is the atomic counter.

Tests

  • Python: test_default_random_key_per_thread in python/tests/test_threads.py. Four threads and the main thread each draw 8 uniform values. All results must be different.
  • C++: test default key per thread in tests/random_tests.cpp. Four threads each draw one 32-bit value on the CPU. All values must be different.

Both new tests fail without this change and pass with it (the Python test on both CPU and GPU). The full Python suite (CPU and GPU) and the full C++ suite have no new failures.

Some tests fail on my machine (Apple M5, macOS 27.2) both with and without this change. I did not investigate them:

  • Python, CPU only: test_conv2d_winograd_disable, test_leaks (export), test_sdpa_split_large_values, test_fft_too_large, test_init_reduce_large_output
  • C++: 6 linalg tests (QR, LU, Cholesky, eigh, inversion, pseudo-inverse)

Not changed

mx.random.seed() still affects only the calling thread, and the docs still call it a global PRNG. I left this for you to decide (see #4856). After this change, multithreaded code gives reproducible results only if each thread calls seed() or passes key=.

This branch has not been deployed

No deployments
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.

[BUG] Threads without a seed get the same random numbers since 0.31.2

1 participant