Add blas_cores argument to pm.sample #7318
Merged
Conversation
This file contains 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
aseyboldt
force-pushed
the
blas-cores
branch
2 times, most recently
from
51daa0c
to
15cbbf4
Compare
ricardoV94
reviewed
May 16, 2024
pymc/sampling/mcmc.py
Outdated
| @@ -499,6 +504,13 @@ def sample( | |||
| Which NUTS implementation to run. One of ["pymc", "nutpie", "blackjax", "numpyro"]. | |||
| This requires the chosen sampler to be installed. | |||
| All samplers, except "pymc", require the full model to be continuous. | |||
| blas_cores: int or "auto" or None, default = "auto" | |||
| The total number of threads blas and openmp functions should use during sampling. If set to None, | |||
There was a problem hiding this comment.
+1
Do you think we should already default to "auto", or first release something where the default is None so that this can be tested a bit more?
There was a problem hiding this comment.
I think the new default makes much more sense. This often shows up in MvNormal models and it's very tricky for beginners to debug
pymc/sampling/mcmc.py
Outdated
Comment on lines
664
to
671
| if cores < 1: | ||
| raise ValueError("`cores` must be larger or equal to one") | ||
|
|
||
| if chains < 1: | ||
| raise ValueError("`chains` must be larger or equal to one") | ||
|
|
||
| if blas_cores is not None and blas_cores < 1: | ||
| raise ValueError("`blas_cores` must be larger or equal to one") |
There was a problem hiding this comment.
Remove it for the sake of less code? I don't believe anybody was ever hurt by this and couldn't figure out the problem?
ricardoV94
approved these changes
May 16, 2024
|
pre commit failed |
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.
Description
We currently do not configure blas in any way. This can lead to very bad behavior if we sample in several threads:
Many blas implementations default to using one worker thread per hardware thread in the machine. But if we sample in parallel with multiprocessing, each chain will use an independent thread pool, so we end up starting
chains*hardware_chainsworker threads. Combined with some spinnlocking that some blas implementations seem to do, this can lead to terrible performance.This PR adds a
blas_coresargument topm.sample(), and then usesthreadpoolctlto control how many worker threads we start.If it is set to
None, we don't do anything, and keep the current behavior of just using whatever the blas implementation uses as default. If set toauto(the default) use the cores argument to guess a decent number of blas worker threads. If it is set to an integer, we use that number of total blas worker.See for instance here for a model that shows bad behavior without this PR.
Type of change
📚 Documentation preview 📚: https://pymc--7318.org.readthedocs.build/en/7318/