Repository navigation
Add MLX backend support for Nutpie compilation - #254
cetagostini wants to merge 40 commits into
Conversation
|
Thanks, that looks great! |
Good point, that simple addition brings between 5% to 20% more performance! @aseyboldt |
|
@aseyboldt solve the test issue to work only on macs with intel chips. |
|
@aseyboldt can you give me a hand? The test failing its strange. My local pass everythig. |
aseyboldt
left a comment
aseyboldt
left a comment
There was a problem hiding this comment.
That failure is annoying. For some reason the results seem to differ between different machines? I think we really should figure out what's going on here. Maybe it helps if we print the first couple of values in warmup_posterior to see if the initial values are already different, or if small differences accumulate?
| updated.update(**updates) | ||
|
|
||
| # Convert to MLX arrays if using MLX backend (indicated by force_single_core) | ||
| if self._force_single_core: |
There was a problem hiding this comment.
We should not use that argument to detect mlx.
How about we add an attribute _convert_data_item or so to the dataclass, that contains a function that transforms data arrays? We could then also use that for jax.
Address review comments and CI failures on PR pymc-devs#254: * MLX is not thread-safe (Metal command-buffer race, ml-explore/mlx#2133), so always set ``force_single_core=True`` for ``backend="mlx"`` regardless of ``gradient_backend``. This unblocks the default config (``gradient_backend="pytensor"``) with ``chains>=2``, which previously aborted with "A command encoder is already encoding to this command buffer". * Decouple shared-data type conversion from ``_force_single_core``: ``PyFuncModel`` now carries an optional ``_shared_data_converter`` callable (set to ``mx.array`` for MLX) used by ``with_data``. The old code was abusing ``_force_single_core`` as a "is MLX backend" proxy. * Drop the stale ``raw_logp_fn`` plumbing in the MLX backend. The transform adapter is flowjax-based (JAX-only), so MLX could never expose a usable raw logp. * Pin ``mlx<0.31`` in ``pyproject.toml`` extras and the CI install. mlx 0.31.x crashes inside ``Compiled::eval_gpu`` when the sampler worker thread evaluates auto-fused element-wise kernels (ml-explore/mlx#3329), which is the root cause of the macOS aarch64 pymc test segfaults. mlx 0.29.x and 0.30.x are unaffected. * Add a runtime guard in ``_compile_pymc_model_mlx`` that raises a helpful ``RuntimeError`` if mlx>=0.31 is installed anyway, instead of segfaulting partway through sampling. * Make ``test_deterministic_sampling_mlx`` a smoke test (drop ``array_compare`` + reference file). MLX sampling is not bit-identical across machines/MLX versions, mirroring the situation already documented for ``test_normalizing_flow``. * Mark ``test_dims_model[mlx-*]`` as ``xfail`` while pymc-devs/pytensor#1350 is open (PyTensor's MLX linker has no ``XTensorFromTensor`` op yet). * Mirror logp's ``mx.array(_x)`` conversion in the expand closure for consistency. * Ignore local AI tooling folders (``.ai/``, ``.cursor/``, ``.claude/``). Made-with: Cursor
Introduces MLX as a backend option in compile_pymc_model, allowing gradient computation via MLX or Pytensor. Updates dependency groups to include MLX, extends internal functions to handle MLX mode, and adds corresponding tests for deterministic sampling with MLX.
Bump MLX version requirement to >=0.29.0 in pyproject.toml for all relevant extras. In compile_pymc.py, JIT compile the logp function using mx.compile for improved performance, aligning with JAX backend behavior.
Address review comments and CI failures on PR pymc-devs#254: * MLX is not thread-safe (Metal command-buffer race, ml-explore/mlx#2133), so always set ``force_single_core=True`` for ``backend="mlx"`` regardless of ``gradient_backend``. This unblocks the default config (``gradient_backend="pytensor"``) with ``chains>=2``, which previously aborted with "A command encoder is already encoding to this command buffer". * Decouple shared-data type conversion from ``_force_single_core``: ``PyFuncModel`` now carries an optional ``_shared_data_converter`` callable (set to ``mx.array`` for MLX) used by ``with_data``. The old code was abusing ``_force_single_core`` as a "is MLX backend" proxy. * Drop the stale ``raw_logp_fn`` plumbing in the MLX backend. The transform adapter is flowjax-based (JAX-only), so MLX could never expose a usable raw logp. * Pin ``mlx<0.31`` in ``pyproject.toml`` extras and the CI install. mlx 0.31.x crashes inside ``Compiled::eval_gpu`` when the sampler worker thread evaluates auto-fused element-wise kernels (ml-explore/mlx#3329), which is the root cause of the macOS aarch64 pymc test segfaults. mlx 0.29.x and 0.30.x are unaffected. * Add a runtime guard in ``_compile_pymc_model_mlx`` that raises a helpful ``RuntimeError`` if mlx>=0.31 is installed anyway, instead of segfaulting partway through sampling. * Make ``test_deterministic_sampling_mlx`` a smoke test (drop ``array_compare`` + reference file). MLX sampling is not bit-identical across machines/MLX versions, mirroring the situation already documented for ``test_normalizing_flow``. * Mark ``test_dims_model[mlx-*]`` as ``xfail`` while pymc-devs/pytensor#1350 is open (PyTensor's MLX linker has no ``XTensorFromTensor`` op yet). * Mirror logp's ``mx.array(_x)`` conversion in the expand closure for consistency. * Ignore local AI tooling folders (``.ai/``, ``.cursor/``, ``.claude/``). Made-with: Cursor
Keep only the upstream-issue references (mlx#2133, mlx#3329) since those are non-obvious context; drop the explanatory comments whose content is already conveyed by the surrounding code. Made-with: Cursor
Deep-dive into the cross-machine non-determinism flagged in review: the MLX backend evaluates the logp/gradient on the Metal GPU, where float64 is unsupported, so ``mx.array(_x)`` silently downcasts nutpie's float64 positions to float32. Measured relative error is ~1e-7 (float32 eps), and because NUTS is a chaotic integrator that error accumulates across the leapfrog steps and varies with GPU/Metal/kernel fusion -- which is why the array_compare reference test could never match between machines. * Cast inputs explicitly to ``mx.float32`` in the logp and expand closures so the precision contract is intentional and cannot silently change with the MLX version or a global default-dtype override, and document the root cause at the conversion site. * Replace the bare smoke test with a machine-independent statistical correctness check: sample HalfNormal and assert the posterior recovers the analytic moments within ~4x the Monte Carlo error. * xfail ``test_pymc_model_store_extra[mlx-pytensor]``: PyTensor's MLX linker Split dispatch passes a numpy array to ``mx.split`` (read as an equal-split count) and breaks on the uneven splits in the ZeroSumNormal/Dirichlet gradient graphs. mlx/mlx is unaffected. Made-with: Cursor
The pymc test suite (and pymc-dev) started crashing with SIGSEGV on Linux/Windows. Reproduced in a clean uv venv matching CI exactly: the fault is in native pandas (e.g. pd.date_range in the unmodified test_zarr_store), not in nutpie or the MLX backend. Root cause: pandas 3.0 wheels are compiled against the numpy 2.5 C ABI but their metadata only requires numpy>=1.26, while numba caps numpy<2.5. uv therefore resolves pandas 3.0.4 together with numpy 2.4.6, and the ABI mismatch segfaults in pandas' compiled extensions. Pinning pandas<3.0 (which is built against the numpy 2.x ABI) resolves it; the full pymc suite then passes. Mirrors the existing arviz<1.0.0 pin. Drop once numba supports numpy>=2.5. Made-with: Cursor
The pymc (dev) and docs jobs installed the locally-built wheel first (--no-deps --no-index), then ran a second `uv pip install nutpie[extra]` without --find-links dist. That second resolution pulled nutpie from PyPI and downgraded it to a release predating the newly-added extras: nutpie 0.13.4 has no 'dev' extra (pytest never installed -> "pytest: command not found") and 0.16.8 has no 'docs' extra (jupyter/nbformat never installed -> `quarto render` fails with ModuleNotFoundError). Add --find-links dist to both extras installs so the local wheel stays the preferred (highest-version) candidate and the new extras resolve against it. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
584f115 to
5a1cd3d
Compare
|
@cetagostini I rebased this branch onto My commits start at To keep the rebase but drop my changes: To go back to exactly what you had before the force-push: Either one followed by a force-push puts the PR branch back. |
|
Benchmarks for the MLX backend on an M4 Max (64 GB, macOS 26.4, mlx 0.32, jax 0.11, numba 0.67). Every backend computes in float32: models are built under At 4 chains, MLX is faster than both CPU backends in one case in the sweep: linear regression at 10^6 observations (1.48x). Below 10^5 observations it is 3-12x slower than the faster CPU backend. Each gradient call pays a GPU round trip of about 150-250 us, and NUTS cannot overlap it with other work, because each leapfrog step needs the gradient on the CPU before the next step starts. Both runs below predate two removals from this branch, the float64 logp accumulator and an experimental chain-batching path. The accumulator's CPU hop cost 40-90 us per call: up to 21% of MLX's cost on small models and 0-4% at N >= 10^5. So the MLX numbers below overstate the current per-call cost by up to 21% on small models and by 4% or less at N >= 10^5. Scaling sweepWall-clock us per gradient evaluation across all chains, warmup included, at 4 chains with tune=200 and draws=200. The speedup is the faster CPU backend divided by MLX.
At 4 chains the GPU is not saturated below about 10^5 observations, so the round trip dominates the cost. More chains put more work on the GPU per unit time. At 32 chains, MLX is 1.44x faster than the faster CPU backend on linear regression at 10^5 and 2.49x at 10^6, and 1.15x and 1.54x on the hierarchical logistic. At 10^3 and 10^4 observations it stays at 0.08-0.61x. The 32-chain runs used the chain-batching path, which evaluated all chains in one GPU call and is not part of this PR. How to make use of the many-short-chains regime in NUTS, where each chain runs its own warmup, is an open question and out of scope here. pymc-model-catalogue34 of 37 catalogue models ran, at 4 chains with tune=1000 and draws=1000. Two ( At 4 chains the catalogue models do not saturate the GPU. numba computes nearly every gradient in 3-50 us, while MLX's cost per gradient per chain stays at 0.2-0.7 ms from 2 to 1,603 parameters, which is the round trip and not compute. The smallest gap is Posterior accuracy was scored as Wasserstein-1 distance to a float64 numba reference, relative to the split-half Monte Carlo floor. MLX agrees with float64 wherever float32 numba and JAX do. Every model that disagrees with the reference on MLX also disagrees on float32 numba and JAX, so none of the mismatches is MLX-specific. Catalogue sampling times
float32 logp without an accumulatorMLX with and without the float64 accumulator, at 4 chains with tune=1000 and draws=1000, gives posteriors and NUTS diagnostics that agree within Monte Carlo error, including at |logp| = 1.4e7, where a float32 logp resolves about 1 nat. The acceptance rate and energy-error spread are the diagnostics logp rounding would disturb first. This result is why I removed the accumulator.
The 10^5 row is measured against a float64 numba reference. The 10^6 rows compare the two MLX variants directly, with the floor taken from the accumulator run's split halves. |
Isn't there room to improve here? What about the so fabled "unified memory"? |
This branch has not been deployed
4 participants
Introduces MLX as a backend option in compile_pymc_model, allowing gradient computation via MLX or Pytensor. Updates dependency groups to include MLX, extends internal functions to handle MLX mode, and adds corresponding tests for deterministic sampling with MLX.