Stanli: Stan language interpreter, embedded in Python, or it can run in the browser and is often faster than CmdStan!

General
1

Check it out: stanli: Stan Language Interpreter, browser edition

It’s using stanc3, the Stan Math library (though it’s using the autodiff in weird ways for perf reasons), and the Stan algorithms verbatim with no upstream patches. Numerics are pretty close to CmdStan and tested to remain that way.

The Python package (or binary executable) is the thing that’s actually faster than CmdStan at median; the browser wasm is often about the same speed minus C++ compilation time. At this point mostly just ODE models are slower, about 0.6x CmdStan with some models being up to 6x faster in stanli.

I made this with an AI this past weekend after talking with @ariddell and getting lots of good ideas from him. I think at the least it occupies another interesting point in the design space for a system like Stan and I think shows how much perf we are leaving on the table by using a var class that allocates and pointer chases for each scalar, which IIRC was mostly to support conditionals on parameters (and perhaps ODEs back in the day, or at least ODE RHS functions that are called a variable number of times). But we can partition the graph into static and dynamic parts and make the static parts (eg the entire model of every model in posteriordb) faster.

LMK what you think! Maybe don’t read the code yet; I’m refactoring it to be cleaner.

2

Awesome, static has been on my wishlist for a long time!

3

Cool!

One part I’m trying to understand better, how much of the adjoint calculations in Stan-math are used and how much is not? I see in stanli/runtime/src/adjoint.cpp at fb707d131541dc9db424a57a17bb159f0db45a9b · seantalts/stanli · GitHub that there are custom adjoints for some functions but this is a small subset. Is part of the ODE issue nested autodiff?

4

Most of the time, most of the functions use the Stan Math adjoints. adjoint.cpp is just some basic scalar op core ad code, and it exists so we can generate fast reverse ad programs for ODE RHS functions, sections of the graph that are conditional on parameters, and sometimes (depending on a cost model) sections of the graph that are not vectorizable / re-rollable into a nice vectorized form. There are about 30 scalar op rules in there copied from Stan Math that we couldn’t re-use directly unfortunately.

There are 3 other categories of AD in stanli:

  1. Densities prim code are run through the rvar recorder where we compute values and partials (via partials_propagator) and then convince the propagator to depost them into our pre-allocated tape with the build() method
  2. Real stan-math rev, nested-tape replay (legacy.hpp / legacy_fns.cpp / matrix_fns.cpp / ode.cpp). We couldn’t get these to do the nice thing above so we run the full rev mode tape replay.
  3. Regions we couldn’t generate an adjoint program for also get the full tape replay.

And yeah, ODEs are slow because of the nested autodiff, it’s basically conditional on parameters as well and can’t be preallocated or run as efficiently.

5

Very cool! @stevebronder and I have discussed doing this a few different times, but one of the things we were lacking was real evidence that it would actually be worth the engineering effort.

6

Yeah, this is an amalgamation of a bunch of old ideas, the most impactful ones from many discussions with Bob and Steve and everyone when I was working on Stan before Covid :)

I’m mostly so pleased AI made this possible to create so quickly. It would have taken me 6 months probably?

7

That was more or less our estimate, and a lot of that was because we were imagining pulling every backward pass lambda in the math library into a separate function so that the static part would have something to call (the equivalent of all the _bwd functions in your kernel/ folder). This is exactly the kind of thing that moves across a ‘worth it’ time boundary when it doesn’t need to be a human doing the same work over and over…

What I’d never considered is how doing the static/dynamic split also makes the idea of interpreting the model rather than compiling it much more palatable, since you dramatically collapse the size of the stan::math template surface you need to instantiate in the precompiled library. Very cool how those things interact!

Not to get too far ahead of things, but are you interested in upstreaming (some of…?) this work after the dust settles? Either in the form of ‘makes something like stanli easier’ (e.g., a more stable, not-so-text-based format for outputting the MIR from stanc?) or a bigger project with the eventual goal of making stanli itself obsolete. I’d be very curious where you think would be the best place to start with such a thing.

8

This is brilliant! Having a toolchain-less Stan would be a huge usability improvement for users that want fast iteration (or just quick results) and don’t need maximum performance.

I’m working on an alternative R interface for Stan and I’ve just done an initial integration for this, pretty painless and allows for the same experience (also works under wasm/webr):

pak::pak("andrjohns/stanr")

library(stanr)

bernoulli_model <- "
  data {
    int<lower=0> N;
    array[N] int<lower=0, upper=1> y;
  }
  parameters {
    real<lower=0, upper=1> theta;
  }
  model {
    theta ~ beta(1, 1);
    y ~ bernoulli(theta);
  }
  generated quantities {
    array[N] real log_lik;
    for (n in 1:N) {
      log_lik[n] = bernoulli_lpmf(y[n] | theta);
    }
  }
"
data <- list(N = 10, y = c(1, 0, 1, 1, 0, 1, 0, 0, 1, 0))

mod <- stan_model(code = bernoulli_model, backend = "stanli")
fit <- mod$sample(data = data)
# Chain [1] Iteration: 1800 / 2000 [ 90%]  (Sampling)
# 
#  Elapsed Time: 0.01 seconds (Warm-up)
#                0.009 seconds (Sampling)
#                0.019 seconds (Total)
# 
# Chain [2] Iteration: 1900 / 2000 [ 95%]  (Sampling)
# Chain [3] Iteration: 2000 / 2000 [100%]  (Sampling)
# 
#  Elapsed Time: 0.01 seconds (Warm-up)
#                0.01 seconds (Sampling)
#                0.02 seconds (Total)
# 
# Chain [1] Iteration: 1900 / 2000 [ 95%]  (Sampling)
# Chain [2] Iteration: 2000 / 2000 [100%]  (Sampling)
# 
#  Elapsed Time: 0.011 seconds (Warm-up)
#                0.009 seconds (Sampling)
#                0.02 seconds (Total)
# 
# Chain [1] Iteration: 2000 / 2000 [100%]  (Sampling)
# 
#  Elapsed Time: 0.011 seconds (Warm-up)
#                0.01 seconds (Sampling)
#                0.021 seconds (Total)

fit$log_prob(0.78)
# [1] -9.208129

fit$loo()
...
#         Estimate  SE
# elpd_loo     -1.6 1.0
# p_loo         0.2 0.1
# looic         3.1 2.1
# ------
# MCSE of elpd_loo is NA.
# MCSE and ESS estimates assume independent draws (r_eff=1).
# 
# Pareto k diagnostic values:
#                          Count Pct.    Min. ESS
# (-Inf, 0.7]   (good)     2     20.0%   3547    
#    (0.7, 1]   (bad)      0      0.0%   <NA>    
#    (1, Inf)   (very bad) 8     80.0%   <NA>    
# See help('pareto-k-diagnostic') for details.

Having a no-compilation backend for brms would very well received by users imo.

P.S. Feel free to break/change the stanli API/implementation as you like, I’m happy to adjust as needed

9

I think upstreaming some of the graph optimizations from stanli into stanc3 could be good, sizeable cool projects for PL grad students or other contributors who want to do cool compilers stuff. Loop reroll (WAY better ways to code that in OCaml) is the main obvious one, maybe also constant folding and data hoisting (though this effectiveness is limited in the compiled Stan model since you don’t provide data until later).

Another project that falls out of this: compile-time reverse mode emitted from stanc3. For whatever portion of the model is parameter-branch-free, we could generate a double-only forward and backward pair, falling back to the normal path for any blocks of code that can’t work this way. This gives “static” autodiff and, with re-rolling, the vast majority of stanli’s perf wins.

I’m not sure a fast-startup interpreter is ever going to be obsolete! I think it’s an amazing interactive development or pedagogical tool and I expect it to stick around, and there’s an entire class of optimizations that become possible once you have access to the data in the optimizer. But I hope that it proves out the “return” portion of ROI for a few of the engineering projects we’ve had around for a while as well and it’s up to whoever to go ahead and start on them upstream where applicable. I feel pretty happy with my stanli side project for now and have a day job that is currently fairly busy. I am happy for anyone to pick up any projects inspired by this.

10

That’s what I was thining! And I’ve been hearing this complaint for many years now… Hope it takes off in classrooms and model iteration workflows.

Super cool! Thanks for integrating that already, this is awesome. Let me know if you have any issues.

I’m not too familiar with how that works, but I think it could already be used as a brms backend, just needs to be added to the brms project. Maybe I can put up a PR for that, hah.

11

That’s what I was thining! And I’ve been hearing this complaint for many years now… Hope it takes off in classrooms and model iteration workflows.

I think it absolutely will, not to mention corporate environments where getting a toolchain installed can be a nightmare in its own right!

I’m not too familiar with how that works, but I think it could already be used as a brms backend, just needs to be added to the brms project. Maybe I can put up a PR for that, hah.

I’ll be putting up a PR to add my package as a backend, so stanli will be added for free!

12

Both of those would be interesting compiler projects for sure.

I remember a few discussions with @rybern over the years where we mentioned how useful it would be to give the Stan compiler the data to specialize the code gen for. This was never seriously pursued in part because re-compiling after changing the data just seemed like too big a hurdle

Oh sure! I didn’t mean that the project idea would be obsolete, but rather there’s a hypothetical future where stanli as a package itself would be, because upstream Stan ends up supporting the same usecases

You might find Stan Playground interesting. It’s been reasonably successful in the classroom environment, but it does still rely on a central server somewhere to actually compile the Stan models to WASM. A version based on these ideas would be able to avoid that (and generally be a quicker edit/run loop)

13

Yep, that was also the issue when we were talking about it back in the day. It makes the API suck… unless you have an interpreter :)

I’m curious what you have in mind. In my head, I don’t think we would ever get rid of compiled Stan because it will always have the best “production” performance (assuming someone runs with some of these ideas upstream). So I am imagining there will always be a world in which there is a separate interpreter package. But what did you have in mind?

This is super cool! I think Stanli would be a great backend here for the obvious reasons that it runs faster and ‘compiles’ faster but also it would mean hosting Stan Playground is literally just static files and requires no server or server compute, maintenance, downtime, overload, etc. There is an npm package: https://www.npmjs.com/package/@seantalts/stanli

14

Nothing concrete in general, just that if you were interested and the project continued that said ‘separate interpreter package’ could be not-so-separate and part of Stan proper somewhere.

For the static autodiff specifically, I think my ideal upstream version would look something like splitting each function in rev into a fwd function, a backward function, and a function-specific struct that gets passed as context to each (to provide a bit more type safety than the untyped pointer math that the kernels in stanli are doing on the context data).

Example based on the fft code currently in stan-dev/math

So instead of the current rev definition of an fft:

template <typename V, require_eigen_vector_vt<is_complex, V>* = nullptr,
          require_var_t<base_type_t<value_type_t<V>>>* = nullptr>
inline plain_type_t<V> fft(const V& x) {
  if (unlikely(x.size() <= 1)) {
    return plain_type_t<V>(x);
  }

  arena_t<V> arena_v = x;
  arena_t<V> res = fft(to_complex(arena_v.real().val(), arena_v.imag().val()));

  reverse_pass_callback([arena_v, res]() mutable {
    auto adj_inv_fft = inv_fft(to_complex(res.real().adj(), res.imag().adj()));
    adj_inv_fft *= res.size();
    arena_v.real().adj() += adj_inv_fft.real();
    arena_v.imag().adj() += adj_inv_fft.imag();
  });

  return plain_type_t<V>(res);
}

You’d have something like

stuct fft_ctx {
  arena_t<Eigen::Vector<std::complex<double>>> arena_v;
  arena_t<Eigen::Vector<std::complex<double>>> res;
};


template <typename V>
inline plain_type_t<V> fft(const V& x, fft_ctx& ctx) {
  if (unlikely(x.size() <= 1)) {
    return plain_type_t<V>(x);
  }

  arena_t<V> arena_v = x;
  arena_t<V> res = fft(to_complex(arena_v.real().val(), arena_v.imag().val()));
  ctx.arena_v = arena_v;
  ctx.res = arena_v;


  return plain_type_t<V>(res);
}

inline fft_bkwd(fft_ctx& ctx) {
    auto adj_inv_fft = inv_fft(to_complex(ctx.res.real().adj(), ctx.res.imag().adj()));
    adj_inv_fft *= ctx.res.size();
    ctx.arena_v.real().adj() += adj_inv_fft.real();
    ctx.arena_v.imag().adj() += adj_inv_fft.imag();
  });
}

You can then keep backwards compatibility by defining fft as something like

template <typename V>
inline plain_type_t<V> fft(const V& x) {
  if (unlikely(x.size() <= 1)) {
    return plain_type_t<V>(x);
  }
  
  fft_ctx ctx;
  auto res = fft_fwd(V, ctx);
  reverse_pass_callback([ctx]() mutable {   fft_bkwd(ctx)  });
  return res;
}

With that all in place in the math library, the compiler could chose between generating fft calls and using the tape as before, or between generating fft_fwd/fft_bkwd pairs in a static situation. This would let us keep it as an optimization pass similar to SoA, at least while it was being developed in stanc3

15

This is extremely cool!

I (and claude) ran the benchmarks locally, but added the --O1 flag to the compiler and updated the benchmarks to use google benchmark. Code and table are below. Oddly, setting --O1 shrank a lot of the speed differences in some of the benchmarked models. Also would it be hard to have stanli take in the optimized mir instead of the transformed mir? That would also give stanli access to a lot of these optims.

I tried my darnd’est to get stan up to the speed of stanli. Mixed success! Just setting some compiler flags and --O1 was enough to get cmdstan up to or faster than stanli for a few models, but there are a handful of tricks going on in stanli that are very cool and near impossible for cmdstan to match without backporting them to the stan compiler.

The below graph’s ns/grad number is the average nanoseconds per gradient evaluation. Speedup is just stanli/cmdstan

model unconstrained params stanli ns/grad CmdStan ns/grad speedup
`radon_pooled` 3 103,174 688,497 6.67x
`arK` 7 3,863 20,352 5.27x
`radon_ic` 391 239,551 952,364 3.98x
`radon_county_intercept` 388 157,058 751,744 4.79x
`nes` 10 28,123 89,649 3.19x
`eight_schools_nc` 10 386 446 1.15x
`election88_full` 90 756,709 1,428,294 1.89x
`bym2_offset_only` 3845 54,836 94,929 1.73x
`dogs` 3 52,029 79,771 1.53x
`kidscore_momiq` 3 2,853 4,977 1.74x
`lsat_model` 1006 117,619 152,073 1.29x
`state_space_slss` 389 36,299 43,700 1.20x
`hmm_example` 4 48,137 49,457 1.03x
`garch11` 4 10,557 12,338 1.17x
`hmm_drive_0` 6 326,306 272,144 0.83x
`normal_mixture` 3 159,234 162,175 1.02x
`low_dim_gauss_mix` 5 187,607 197,962 1.06x
`wells_dist100ars_model` 3 19,675 18,368 0.93x
`iohmm_reg` 29 1,088,012 736,116 0.68x
`radon_county` 389 118,430 122,811 1.04x
`arma11` 4 11,886 6,430 0.54x
`diamonds` 26 72,496 28,106 0.39x
`ldaK2` 7 522,558 270,922 0.52x

radon_pooled is the best example of where stan would benefit from even a simple for loop collapser. The model spends 99% of its time in one loop

  for (n in 1 : N) {
    target += normal_lpdf(log_radon[n] | mu[n], sigma_y);
  }

So if the compiler could rewrite this to be single line statement that would turn on a ton of vectorization

target += normal_lpdf(log_radon | mu, sigma_y);

arK is another for loop collapse that is really interesting. It is impressive that it can take that double loop scheme and work out the vectorized version

Before:

  for (t in (K + 1) : T) {
    real mu;
    mu = alpha;
    
    for (k in 1 : K) {
      mu = mu + beta[k] * y[t - k];
    }
    
    y[t] ~ normal(mu, sigma);
  }

After:

  vector[T - K] mu = rep_vector(alpha, T - K);

  for (k in 1:K) {
    mu += beta[k]
          * to_vector(y[(K + 1 - k):(T - k)]);
  }

  to_vector(y[(K + 1):T]) ~ normal(mu, sigma);

I think there is a way to reduce the second loop as well, but that requires some data logic that I think would be very wonky.

kidscore_momiq comes in as 1.74x faster than cmdstan and that one has nothing to do with a for loop collapser. Stanli ran a gradient for this model in 2,853ns while stan took 4,977ns.

The model is pretty much just the below two lines

model {
  sigma ~ cauchy(0, 2.5);
  kid_score ~ normal(beta[1] + beta[2] * mom_iq, sigma);
}

I need to look more into why this model is still faster, but I think this one could be misses from stanc or just the fact that stanli sets up the autodiff graph once and then fills in new values on the forward and reverse pass. Stan reuses the memory, but not the nodes themselves. Stanc could phrase the inner part of the normal as an fma, or it could make it a normal_id_glm.

I still need to look into more of the models to pick apart why they are faster, but overall this is very cool and exciting! Having a fast interpreted version of stan is awesome. Would you be open to making a pull request in the stanc compiler for the for loop collapser / any other optims you think would be good to have? The for loop collapser would be awesome in the compiler because then that would also turn on the SoA matrix types for a lot of stan models.

16

Hey Steve! Awesome, really appreciate you digging in here. I hope this project gives some nice hard numbers for some of these perf improvement directions that have been around. Sorry about not running -O1, I forgot that existed / maybe it’s new since I last looked at Stan in 2020 or so :P

Not hard, but it doesn’t add any perf really. Maybe 1% on some models. Though that’s partly because I have the goal of matching vanilla CmdStan numerics, which means stanli has to desugar eg fma / lmultiply after stanc3 optimizes them. I think I’ll still make this change and maybe see how bad the numerics differences get with those changes…

This one I can do, though I was on vacation the past week or so and had a lot more time to spend on this than I do now (big new project at work). I had Fable queue up a few PRs for the loop vectorizer that gives up to 3x speedup for some models in stanc3; I have to review them before I send them over to you guys. Feel free to crib wholesale, it’s all BSD license. I think the two biggest perf improvements are both primarily stanc3 changes (“vectorization” / rerolling, and static autodiff) though I think Brian’s API for it in stan math makes sense and would need an adapter layer.

The dream would be that all the complicated code for that stuff leaves Stanli and ends up in stanc3 and Stan math and Stanli becomes very simple. In that world I’d expect it to probably be slower than CmdStan in most cases - the only class of improvement would be the data constant folding. But I think Stanli would still be useful as something that’s easier to install.

17

Can R users use your stanr package to use WALNUTS?

18

Yep! Calling $sample(engine = "walnuts") is available for both normal compiled Stan models and the stanli backend

19

Did you see the issue I opened on stanr Github?

20

This is really cool! Do you have any plan to add nutpieR too? GitHub - andytimm/nutpieR: R interface for the fast nuts-rs MCMC sampler · GitHub