Repository navigation
Fix identical default random keys across threads - #4857
Open
Chr0mium24 wants to merge 1 commit into
Open
Chr0mium24 wants to merge 1 commit into
Chr0mium24 wants to merge 1 commit into
Conversation
This branch has not been deployed
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.
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 passkey=produce the same random numbers.Changes
KeySequence::default_seed()inmlx/random.cpp. It returns a timestamp, computed once per process, plus an atomic counter. Each call returns a different seed.KeySequence::default_()inmlx/random.hcreates each thread's key fromdefault_seed().PyKeySequence::state()inpython/src/random.cppuses the same function. C++ and Python default keys share one counter, so they also get different seeds.<chrono>include frommlx/random.handpython/src/random.cpptomlx/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
test_default_random_key_per_threadinpython/tests/test_threads.py. Four threads and the main thread each draw 8 uniform values. All results must be different.test default key per threadintests/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:
test_conv2d_winograd_disable,test_leaks(export),test_sdpa_split_large_values,test_fft_too_large,test_init_reduce_large_outputNot 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 callsseed()or passeskey=.