This is Bob.
JAX on my mind
I’ve been thinking a lot about JAX lately. JAX is appealing to a computer scientist like me, due to its beautifully compositional architecture for coding autodiff on GPUs. Had JAX existed when we started coding Stan in 2011, we would’ve used that rather than rolling our own autodiff system. Apparently Theano did exist at that time, but we didn’t hear about it until long after releasing the first version of Stan.
Why JAX?
Originally, my interest was sparked by Matt Hoffman’s work on Cheese and Meads samplers (too lazy to look up their capitalization pattern, so going Andrew style) that use massively parallel HMC steps on GPU (Matt was driven to abandon NUTS because its recursive structure is anathema to GPU acceleration). It continues to be kindled by statements from people like Elizaveta Semenova, who announced during the recent StanCon at Oxford that she had to give up Stan and moved to NumPyro because she couldn’t code neural nets easily in Stan and Stan doesn’t scale on the GPU (we have some operations that can be sent to GPU, but we can’t keep the whole eval in kernel). The full blown fire is due to Justin Domke and Abhinav Agrawal’s work on normalizing flows—they have a repo, vistan, that implements their methods, including a greatly improved version of autodiff variational inference (ADVI) and a version of real non-volume preserving (realNVP) normalizing flows. Gilad Turok, who’s a research analyst here and is applying to grad school for next year, is almost done with a better engineered version we can submit to Blackjax.
We haven’t gotten far, but so far we haven’t found a density yet that the realNVP flows, followed by importance resampling, didn’t fit nearly perfectly. For example, it could fit a 1000-parameter hierarchical IRT 2PL without any identifying strategy other than priors, which is a model that Stan cannot fit well at all due to the product of funnels structure of the posterior geometry. Justin and Abhinav’s ADVI and realNVP variants rely on JAX’s ability to massively parallelize the log density and gradient calculations in order to stabilize the stochastic gradients used to calculate KL divergence up to a constant (i.e., the ELBO).
Comparing Stan and JAX performance
At the recent StanCon in Oxford, Simon Maskell presented a result I’ve been meaning to generate for a year or two, which is a comparison of JAX and Stan on CPU. Here’s Simon (or at least his arm!) presenting the result:
Working out the algebra problem Simon’s slide title presents (JG = 20 * JC, JG= 5 * SC, where JG is JAX on GPU, JC is JAX on CPU, and SC is Stan on CPU), we see that JAX on GPU is 5 times faster than Stan on CPU, which is in turn 4 times faster than JAX on CPU. For the evaluation, Simon was using a high-end consumer-grade (RTX) NVIDIA GPU—we will re-run Simon’s results on our state-of-the-art GPUs over the next couple of months and report back. Tensorflow did a similar evaluation with similar results a few years ago, but I wasn’t able to find it.
CPU to GPU over next decade?
This is actually kind of a bummer in that it means we couldn’t really move to JAX as a back end for most of our users because it’s not very performant on CPUs. My guess is that the median Stan user in terms of computation has a 5-year old Window’s notebook with no GPU. Part of my interest in things like JAX is my feeling that in 10 years, this is no longer going to be the case, and almost everyone trying to fit a statistical model will be working on a cluster.
Are there any fundamental architecture issues that make it so JAX/XLA performance on CPUs couldn’t be improved? I imagine that TPUs and GPUs (in that order) have been the focus of the XLA team and the main priority of JAX and Tensorflow’s users. It’s possible that with some dedicated engineering effort they could be brought up to parity with Stan.
Bob:
You’re my open-science hero because you wrote Stan and you can still write dispassionately about its weaknesses as well as its strengths.
I will say, as a user of Stan, that I absolutely love how the Stan modeling language allows me to write a model, define an objective function, and perform predictive simulations without having to think much about the computation. There can never be perfect compartmentalization of modeling and computing, but for many problems Stan takes us closer to that ideal, more than any other language I’ve seen.
Code being easy to read was one of my goals for the Stan language. When you’re a programmer in industry or on a large-scale open-source project like Stan, you spend much more time reading code than writing code.
For expressiveness, I feel like there was a fundamental decision in the early design that I wish I could do over. I think the blockless version of Stan that Maria Gorinova developed in her thesis would have had a lot of advantages. On the other hand, the way it’s currently designed makes it very clear what the data, parameters, etc. are—it’s much more closely aligned to how the computation happens.
The other PPLs that people use are more graphical modeling oriented than they are define-a-log-density oriented. That also has some advantages in terms of being able to do things like automatic SBC or PPCs. And it also has some drawbacks if you want to do things like define a custom density or constraint.
Bob, have you ever tried Turing.jl in Julia?
It is very much a “define a log density function” approach. It uses Julia itself as the language, adding a few specialized syntaxes through its use of macros. So
a ~ Normal(0,1)
For example means make `a` a parameter with prior distribution Normal(0,1) (and this Normal is a struct that’s a subtype of UnivariateDistribution, there are many others.)
Having the whole of the Julia programming language as your modeling language is just amazingly powerful and you immediately benefit from whatever other packages there are in Julia. For example ApproxFun is one I really like, defining functions in terms of Chebyshev or other orthogonal polynomials. There are also various spline libraries, there’s special functions, the whole Julia differential equations solver system, the optimization subsystems, geospatial coordinate transforms, etc etc.
At the moment I’m working on a project involving migration of people within countries. we create 2D polynomial representations of desirability over spatial coordinates, calculate distances between locations using a haversine distance from Distances.jl and have a 2D surface over transitions between density (more or fewer people per km^2) all in like 10-12 lines of Turing code using multiple extant packages.
We are fitting by first approximate optimization and then a short sample run to get a single sample in near equilibrium at the end, because computing with the model is slow (500,000 data points) but it’s still a vast improvement over the state of the art in migration models.
I declared Turing.jl the “best” PPL for supporting Bayesian workflow in my ProbProg talk. But that’s because I took it to be a graphical modeling language, not a simple differentiable density definition language like Stan. When you go toward Stan’s notion of just defining a log density, you can no longer extract graphical models to do, say, forward simulation for calibration or prior predictive checks.
Which Julia autodiff system are you using? I realy like how Zygote.jl is architected in order to make writing new custom functions easy. I find the code there easier to read there than in the math papers about matrix derivatives.
My understanding from talking to some of the Turing.jl devs about a year ago is that all the Julia autodiff libraries have partial coverage of the Julia libraries. Also, Edward Roualdes developed the initial code for BridgeStan around a year and a half ago so he could use Stan models within Julia, because he found Julia’s autodiff to be too slow. I’d love to hear an update of how things stand now.
Turing has the macro @addlogprob! which allows you to in theory do any density you want.
At the moment, as I say, we’re just using optimization with LBFGS followed by a short run of diffusive metropolis hastings to further equilibriate, and then taking the final point as our single “optimum” sample. That takes something like 5-20 mins to fit. (on the order of 500k data points, on the order of 100 parameters for the US data, we were doing 6 different age groups in German data so around 6x that many parameters there. These days we are thinking of fitting each age group separately. Also I might throw Pathfinder at it.
I tried running NUTS with both ForwardDiff.jl, and ReverseDiff.jl with and without compiled tape, and the thing wanted a couple days to do the sampling, and it just wasn’t worth the hassle for this particular project. I’ve heard very good things about Enzyme.jl which auto-diffs the LLVM code, and I even tried it once in the early days when the model was significantly more broken, maybe I should go back and try that again now that the model works well with optimization!
Yes, none of the autodiff systems are perfect, but Enzyme is likely to open up a lot more options I think. People are telling me it’s stable enough for everyday use for the most part.
Finally, I intend to try writing my own “White Box” sampler that actually works with Turing models and try that out. (we discussed this algorithm in 2020 here: https://discourse.mc-stan.org/t/the-typical-set-and-its-relevance-to-bayesian-computation/17174/27 )
@Dashiell there is no fundamental architectural issue per se, JAX from a framework designer is a very thing frontend over an MLIR backend with JAX’s own representation, JAXPR, sitting inbetween. There has not been that much incentive to optimize the MLIR compilation pipelines on the CPU for JAX, but it can be done. If one would actually desire it (and can live with a hacky solution), one could quite probably pull some of the MLIR compilation pipeline on the CPU from the modern Flang compiler which is already MLIR-based. As far as I know Xanadu also does custom compilation coming from JAX to be more performant. But I am not sure whether that is on the CPU.
If I may ask, why would one want to move to a JAX backend itself? Would it not be more advantageous for an established framework like Stan to potentially rather use an MLIR compilation backend, potentially like IREE, or OpenXLA to tap into the same ecosystem without being dependent on the JAX frontend?
Thanks, @Ludger. I didn’t even know that XLA (low level linear algebra accelerator) had a linked middle representation layer and code generator (MLIR)! That does sound like a better idea. Similarly, we could’ve targeted something lower-level than C++ from Stan like LLVM, but I had no idea how to do that or how successful LLVM was going to be back when we started in early 2011.
The suggestion to target a lower-level intermediate representation layer is in line with what Matthijs Vákár suggested recently when visiting—that we should build an XLA-based autodiff. Matthijs was key in moving from the ad hoc C++ parser I wrote using Boost’s tools to our current functional OCaml version using Menhir that has a proper middle representation and code generation layer so we could optimize our C++ code generation (that’s an ongoing operation).
I had just assumed that CPU wasn’t a priority for JAX or even for XLA. It can be challenging to take the same client API and optimze for two different compilation targets simultaneously. But I wouldn’t be surprised at all if serious engineering effort could make it faster than Stan on CPU, even without disturbing anything that makes it fast on GPU. But I really have no direct experience, so that’s a wild guess.
Oh, I think you did a lot of things right there. LLVM was not stable at this point of time, and every once in a while it introduces a number of breaking changes, which one then needs to go and fix. So with the benefit of hindsight, your design choices look a lot more stable to me.
(Being slightly biased here, having worked on Enzyme in the past) I don’t think building XLA-level Autodiff is a good pitch, you would want another compiler plugin to ideally do that for you s.t. you don’t need to maintain that part of the infrastructure, and keep it up to date with the rapidly changing main branch of MLIR. Something like Enzyme (enzyme.mit.edu), once stable on MLIR, would be a reasonable approach in my opinion. That would still retain the ability to write custom MLIR transformation passes to leverage Stan-specific knowledge.
I think it can be made faster on CPU, but seeing the competitive pressure to run well on GPU for large LLM training, I am not sure I would expect Google, or others, to invest significant manpower into this any time soon sadly..
“…Part of my interest in things like JAX is my feeling that in 10 years, this is no longer going to be the case, and almost everyone trying to fit a statistical model will be working on a cluster.”
I have my doubts that adding computer resources is going to make their models better.
From my industry, hybrid statistics-ML perspective, there are several related things going on:
– Computation is moving to cloud, and one can run IDEs like RStudio or VSCode to a remote compute pretty much just like to local, conditional on a good enough network connection. This opens up easy multi-CPU (with no fan noise), GPUs, and clusters too. Usability will depend on the model, for my statistical models I haven’t yet find GPUs useful, but that depends.
– Good GPU-like chips, maybe with JAX support, will be commonplace in laptops because they are needed for local AI. Recent Apple devices are examples of that (with poor JAX support).
– There will be plenty of open generic models for feature extraction and even for more statistical tasks like time-series prediction. These will integrate to some statistical work, not to other. (Say, if you have open survey questions, you may want to have LLM summaries, classifications of embeddings from those.) Some of these models are available through APIs, but some may need to be run by yourself. That may sometimes mean cloud is more useful than one’s laptop.
– (Probably more, these just currently in my mind.)
Of course it’s possible that sampling algorithms themselves go parallel in a way that makes GPU usable for all models. (I have no insight to that.)
@James B. Shearer: The goal is to make inference faster. One might argue that this will help develop better models by accelerating the development/testing process (aka workflow).
@Janne Sinkkonen: I’ve been seeing the same thing. I keep forgetting that people want to do consumer-level AI. The new ARM archiecture on Mac is really great for running statistical models where it’s hard to pipeline data (e.g., large-scale spatio-temporal models) or you want to run multiple chains in parallel—we’ve seen cases where it’s around 4 times as fast as a similarly priced Intel/AMD machine for running just four chains of Stan models.
GPU is just way faster to run in parallel if you use a GPU-friendly sampling algorithm and you can get everything stable in single-precision arithmetic. There is a lot of work going on now to exploit parallelism in algorithms for both MCMC and variational methods like normalizing flows.
Hi Bob,
Can you share any tips for getting the most out of the new M1 chips? Using the default cmdstanr installation procedure and sampling with `chains = 4, parallel_chains = 4` has my machine using ~0% GPU load and 50% CPU load, which seems like it’s leaving performance on the table.
I think the minimal M1 has 8 cores, so it’s not surprising that using parallel_chains=4 only uses 50% CPU load. My Xeon iMac Pro has 8 cores, but you won’t see it getting to 50% CPU load with 4 chains because memory is a bottleneck. So you also want to measure memory pressure to see how well the M1 is doing; see https://support.apple.com/guide/activity-monitor/check-if-your-mac-needs-more-ram-actmntr34865/mac for instructions on how to do that. If you’re not redlined on memory, add more parallel chains. With 8 cores, you can cut the sampling phase in half compared to 4 cores in terms of iterations and get the same expected effective sample size—but warmup is still serial, so you still need to have warmup converge for all 8 chains (which often happens in fewer than our default 1000 iterations, but we haven’t instrumented good ways to monitor that).
Thanks Bob, I’ll give this a shot
Bob: One of the nice things about Apple’s M-series is the Unified Memory, which means CPU, GPU, and NPU can access the same memory — no transfer to/from the GPU is necessary. Apple’s MLX framework (similar to PyTorch, with its similarity to numpy) specifies a “channel” (cpu, gpu) to process, not to store. (The NPU is specialized and wouldn’t be useful for Stan anyhow.)
Old post, but re-reading this reminded me of Bob making the same comment about hardware requirements at StanCon. I have low-cost access to a high-performance compute cluster with thousands of GPUs through my university. What I do not have is funding for a mathematically and computationally talented graduate student who can spend endless hours reparameterizing models with difficult geometry so that they can work in Stan. If leveraging GPU can get a model sampling faster, or even just enable faster iteration over model design, that would probably be a huge win for anyone in the US working in a university setting. I could not care less whether or not Stan is fast on my personal laptop.