Repository navigation
Support the free threaded build #2807
Description
Activity
I'm able to build mlx and run the tests on my M3 Macbook Pro. I'm not sure why, but I had to force cmake to find the free-threaded python install I built using pyenv:
diff --git a/CMakeLists.txt b/CMakeLists.txt index 3487c22c..933fc536 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -272,6 +272,7 @@ target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:fmt::fmt-header-only>) if(MLX_BUILD_PYTHON_BINDINGS) message(STATUS "Building Python bindings.") + set(Python_ROOT_DIR /Users/goldbaum/.pyenv/versions/3.14.1t) find_package( Python 3.8 COMPONENTS Interpreter Development.ModuleAfter that, all but one of the tests pass:
(.venv) ± PYTHON_GIL=0 python -m unittest discover python/tests .....................s......................................................F.......................s................................................................................s........sssssssssssssssss..............................................................................s..............ss..................................................................................................................................................................................................................................................................s.......................................................................s.......................transform(x)=State(a=array(11, dtype=int32), b=array(20, dtype=int32)) . ====================================================================== FAIL: test_grad_ids_pre_post (test_autograd.TestAutograd.test_grad_ids_pre_post) ---------------------------------------------------------------------- Traceback (most recent call last): File "/Users/goldbaum/Documents/mlx/python/tests/test_autograd.py", line 785, in test_grad_ids_pre_post self.assertEqual(init_id, id(arrs[0])) ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^ AssertionError: 4265975377344 != 4265975363904 ---------------------------------------------------------------------- Ran 657 tests in 67.071sIt looks like there is a multithreaded backend for
mlx.launch, I don't know offhand if the tests execute it. That said, pretty encouraging that all the C++ bindings seem to work correctly.@awni @angeloskath hope you don't mind the ping - I think the discussion in this thread on a PR earlier this year that you two participated in is relevant to the test failure I've found.
wjakob/nanobind@d793091 seems like it might be related, maybe? But this isn't a multithreaded test, so maybe not. All the uses of
nb::castseems to use the defaultautomaticrv_policyas well.I'm a little mystified why there's any behavior difference. Any chance either of you have a little extra context?
I guess I can also add some print debugging inside
py_value_and_gradand try running with both builds to see where the divergence in behavior starts...I'm now thinking that the test succeeding on the GIL-enabled build is a red herring. Checking id's doesn't really verify anything, since the id can be the same for two distinct objects if the interpreter happens to re-use an object address.
My guess is that because the free-threaded build / mimalloc has different object re-use semantics, this test has different behavaior.
If I modify the test slightly:
import mlx.core as mx def fun(arrs): return arrs[0] arrs = [mx.array(1.0)] arr = arrs[0] mx.grad(fun)(arrs) assert arr is arrs[0]That fails on both the free-threaded and GIL-enabled build.
So I think it's probably safe to ignore this failure and either skip the test on the free-threaded build or simply delete the test because it's not checking something useful.
It's pretty weird to me that the modified test you posted fails. The whole point is that the objects should be the same (not distinct).
I think that's not actually happening. I think you made a mistake back here by writing the test using
idinstead of retaining a reference and usingisand the object was never preserved, it just hapenned to work out that you always get back an object with the same id because CPython re-uses objects.Yup sure seems like we're just getting lucky with that test. We should fix it properly
Reacted by somedatapacketWe should fix it properly
Let me file a new issue for this. I tried to understand the object re-use semantics but am not very familiar with nanobind and didn't get much further than some vague printf debugging before I noticed the test doesn't make any sense in the first place.
Moving on to look at thread-safety, if I run the unit tests under unittest-ft to run them under a thread pool, I see:
goldbaum at Nathans-MBP in ~/Documents/mlx on nanobind ± PYTHON_GIL=0 unittest-ft python/tests -[AGXG15XFamilyCommandBuffer tryCoalescingPreviousComputeCommandEncoderWithConfig:nextEncoderClass:]:1091: failed assertion `A command encoder is already encoding to this command buffer' [1] 6496 abort PYTHON_GIL=0 unittest-ft python/testsNote that this crash happens on the GIL-enabled build and the free-threaded build, so there's some kind of use of global state that probably needs Python-level locking and this is likely a pre-existing issue on the GIL-enabled build.
I looked at adding free-threaded CI today. There's a crash in the MPI tests on Linux that I need to chase down: https://github.com/ngoldbaum/mlx/actions/runs/20174769146/job/57919840988. I only get that far on the aarch64 runner, the Linux x86_64 runner crashes at the end of the installation step in such a way that it doesn't even produce logs in the github UI: https://github.com/ngoldbaum/mlx/actions/runs/20174769146/job/57919840967. Not sure what's happening there.
I can't do any debugging on the Mac CI because they use self-hosted runners.
After disabling the MPI test, I did get the CPU aarch64 Linux builds to pass with Python 3.14t: https://github.com/ngoldbaum/mlx/actions/runs/20175170976/job/57922579676.
I think I'll try to figure out what's going wrong on the Linux tests and come back to Mac tests later. Maybe you could give my GitHub account sufficient permissions so the tests run automatically for my PRs?
I'd also appreciate feedback on whether you'd prefer I only add one 3.14t CI job or if you'd be OK with adding one on Linux CPU, one Linux GPU, and one on MacOS. Currently all but the nightly cron tests uses Python 3.10.
Currently, mlx and any package that depends on mlx fails to install on the free-threaded build:
And indeed, there aren't any cp314t wheels on PyPI.
It'd be really nice to be able to use mlx on the free-threaded build. Both PyTorch and jaxlib ship free-threaded wheels.
Also see https://py-free-threading.github.io for more about using the free-threaded build and https://py-free-threading.github.io/porting-extensions/ for more about updating extensions.
Myself and my team at Quansight have been actively working on adding support for the free-threaded build in community packages for the past year and a half. I can also loop in CPython core devs and other experts who can offer advice.
Currently our short-term goal is to get vLLM running on the free-threaded build. This is in support of that, since mlx is an indirect dependency on apple silicon macs, via xgrammar. See vllm-project/vllm#28762.