|
| 1 | +- Feature Name: Reduce Sum Parallelization |
| 2 | +- Start Date: 2020-03-10 |
| 3 | +- RFC PR: 17 |
| 4 | +- Stan Issue: |
| 5 | + |
| 6 | +# Summary |
| 7 | +[summary]: #summary |
| 8 | + |
| 9 | +The reduce_sum function is an attempt to add a parallelization utility that is much easier to use than map-rect without requiring a substantial changes to the autodiff stack. |
| 10 | + |
| 11 | +# Motivation |
| 12 | +[motivation]: #motivation |
| 13 | + |
| 14 | +The goal of reduce_sum is to make it easier for users to parallelize their models by streamlining how arguments are handled, hiding work scheduling, and making it more difficult to program something that will accidentally be inefficient. |
| 15 | + |
| 16 | +# Guide-level explanation |
| 17 | +[guide-level-explanation]: #guide-level-explanation |
| 18 | + |
| 19 | +```reduce_sum``` is a tool for parallelizing operations that can be represented as a sum of functions, `g: U -> real`. |
| 20 | + |
| 21 | +For instance, for a sequence of ```x``` values of type ```U```, ```{ x1, x2, ... }```, we might compute the sum: |
| 22 | + |
| 23 | +```g(x1) + g(x2) + ...``` |
| 24 | + |
| 25 | +In terms of probabilistic models, comes up when there are N conditionally independent terms in a likelihood. In these cases each conditionally indepedent term can be computed in parallel. If dependencies exist between the terms, then this isn't possible. For instance, in evaluating the log density of a Gaussian process ```reduce_sum``` would not be very useful. |
| 26 | + |
| 27 | +```reduce_sum``` doesn't actually take ```g: U -> real``` as an input argument. Instead it takes ```f: U[] -> real```, where ```f``` computes the partial sum corresponding to the slice of the sequence ```x``` passed in. For instance: |
| 28 | + |
| 29 | +``` |
| 30 | +f({ x1, x2, x3 }) = g(x1) + g(x2) + g(x3) |
| 31 | +f({ x1 }) = g(x1) |
| 32 | +f({ x1, x2, x3 }) = f({ x1, x2 }) + f({ x3 }) |
| 33 | +``` |
| 34 | + |
| 35 | +If the user can write a function ```f: U[] -> real``` to compute the necessary partial sums in the calculation, then we can provide a function to automatically parallelize the calculations (and this is what ```reduce_sum``` is). |
| 36 | + |
| 37 | +If the set of work is represented as an array ```{ x1, x2, x3, ... }```, then mathematically it is possible to rewrite this sum with any combination of partial sums. |
| 38 | + |
| 39 | +For instance, the sum can be written: |
| 40 | + |
| 41 | +1. ```f({ x1, x2, x3, ... })```, summing over all arguments, using one partial sum |
| 42 | +2. ```f({ x1, ..., xM }) + reduce({ x(M + 1), x(M + 2), ...})```, computing the first M terms separately from the rest |
| 43 | +3. ```f({ x1 }) + f({ x2 }) + f({ x3 }) + ...```, computing each term individually and summing them |
| 44 | + |
| 45 | +The first form uses only one partial sum and no parallelization can be done, the second uses two partial sums and so can be parallelized over two workers, and the last can be parallelized over as many workers as there are elements in the array ```x```. Depending on how the list is sliced up, it is possible to parallelize this calculation over many workers. |
| 46 | + |
| 47 | +```reduce_sum``` is the tool that will allow us to automatically parallelize this calculation. |
| 48 | + |
| 49 | +For efficiency and convenience, ```reduce_sum``` allows for additional shared arguments to be passed to every term in the sum. So for the array ```{ x1, x2, ... }``` and the shared arguments ```s1, s2, ...``` the effective sum (with individual terms) looks like: |
| 50 | + |
| 51 | +``` |
| 52 | +g(x1, s1, s2, ...) + g(x2, s1, s2, ...) + g(x3, s1, s2, ...) + ... |
| 53 | +``` |
| 54 | + |
| 55 | +which can be written equivalently with partial sums to look like: |
| 56 | + |
| 57 | +``` |
| 58 | +f({ x1, x2 }, s1, s2, ...) + f({ x3 }, s1, s2, ...) |
| 59 | +``` |
| 60 | + |
| 61 | +where the particular slicing of the ```x``` array can change. |
| 62 | + |
| 63 | +Given this, the signature for reduce_sum is: |
| 64 | + |
| 65 | +``` |
| 66 | +real reduce_sum(F func, T[] x, int grainsize, T1 s1, T2 s2, ...) |
| 67 | +``` |
| 68 | + |
| 69 | +1. ```func``` - User defined function that computes partial sums |
| 70 | +2. ```x``` - Argument to slice, each element corresponds to a term in the summation |
| 71 | +3. ```grainsize``` - Target for size of slices |
| 72 | +4-. ```s1, s2, ...``` - Arguments shared in every term |
| 73 | + |
| 74 | +The user-defined partial sum functions have the signature: |
| 75 | + |
| 76 | +``` |
| 77 | +real func(T[] x_subset, int start, int end, T1 arg1, T2 arg2, ...) |
| 78 | +``` |
| 79 | + |
| 80 | +and take the arguments: |
| 81 | +1. ```x_subset``` - The subset of ```x``` (from ```reduce_sum```) for which this partial sum is responsible (```x[start:end]```) |
| 82 | +2. ```start``` - An integer specifying the first term in the partial sum |
| 83 | +3. ```end``` - An integer specifying the last term in the partial sum (inclusive) |
| 84 | +4-. ```arg1, arg2, ...``` Arguments shared in every term (passed on without modification from the reduce_sum call) |
| 85 | + |
| 86 | +The user-provided function ```func``` is expect to compute the ```start``` through ```end``` terms of the overall sum, accumulate them, and return that value. The user function is passed the subset ```x[start:end]``` as ```x_subset```. ```start``` and ```end``` are passed so that ```func``` can index any of the tailing ```sM``` arguments as necessary. The trailing ```sM``` arguments are passed without modification to every call of ```func```. |
| 87 | + |
| 88 | +The ```reduce_sum``` call: |
| 89 | + |
| 90 | +``` |
| 91 | +real sum = reduce_sum(func, x, grainsize, s1, s2, ...) |
| 92 | +``` |
| 93 | + |
| 94 | +can be replaced by either: |
| 95 | + |
| 96 | +``` |
| 97 | +real sum = func(x, 1, size(x), s1, s2, ...) |
| 98 | +``` |
| 99 | + |
| 100 | +or the code: |
| 101 | + |
| 102 | +``` |
| 103 | +real sum = 0.0; |
| 104 | +for(i in 1:size(x)) { |
| 105 | + sum = sum + func({ x[i] }, i, i, s1, s2, ...); |
| 106 | +} |
| 107 | +``` |
| 108 | + |
| 109 | +As an example, in the regression: |
| 110 | +``` |
| 111 | +data { |
| 112 | + int N; |
| 113 | + int y[N]; |
| 114 | + vector[N] x; |
| 115 | +} |
| 116 | +
|
| 117 | +parameters { |
| 118 | + vector[2] beta; |
| 119 | +} |
| 120 | +
|
| 121 | +model { |
| 122 | + beta ~ std_normal(); |
| 123 | + y ~ bernoulli_logit(beta[1] + beta[2] * x); |
| 124 | +} |
| 125 | +``` |
| 126 | + |
| 127 | +the likelihood term: |
| 128 | + |
| 129 | +``` |
| 130 | +y ~ bernoulli_logit(beta[1] + beta[2] * x); |
| 131 | +``` |
| 132 | + |
| 133 | +can be written equivalently (up to a constant of proportionality) as a sum over terms: |
| 134 | + |
| 135 | +``` |
| 136 | +for(n in 1:N) { |
| 137 | + target += bernoulli_logit_pmf(y | beta[1] + beta[2] * x); |
| 138 | +} |
| 139 | +``` |
| 140 | + |
| 141 | +Because this sum can be broken up into partial sums, ```reduce_sum``` can be used |
| 142 | +to parallelize this model. Writing the function for the partial sums and |
| 143 | +updating the model block to use ```reduce_sum``` gives: |
| 144 | + |
| 145 | +``` |
| 146 | +functions { |
| 147 | + real partial_sum(int[] y_subset, |
| 148 | + int start, int end, |
| 149 | + vector x, |
| 150 | + vector beta) { |
| 151 | + return bernoulli_logit_lpmf(y_subset | beta[1] + beta[2] * x[start:end]); |
| 152 | + } |
| 153 | +} |
| 154 | +
|
| 155 | +data { |
| 156 | + int N; |
| 157 | + int y[N]; |
| 158 | + vector[N] x; |
| 159 | +} |
| 160 | +
|
| 161 | +parameters { |
| 162 | + vector[2] beta; |
| 163 | +} |
| 164 | +
|
| 165 | +model { |
| 166 | + int grainsize = 100; |
| 167 | + beta ~ std_normal(); |
| 168 | + target += reduce_sum(partial_sum, y, |
| 169 | + grainsize, |
| 170 | + x, beta); |
| 171 | +} |
| 172 | +``` |
| 173 | + |
| 174 | +# Reference-level explanation |
| 175 | +[reference-level-explanation]: #reference-level-explanation |
| 176 | + |
| 177 | +The argument packing/unpacking is avoided by using lots of C++ parameter packs, the parallelization is managed by the existing TBB autodiff extensions (each TBB worker thread gets its own separate nested autodiff stack), and by requiring that the output of all the functions to be a scalar, this calculation can be done efficiently with nested forward mode autodiff. In this way the autodiff work is parallelized without having to add any parallelization mechanism to the reverse mode sweep. |
| 178 | + |
| 179 | +There are hidden costs associated with the nested autodiff, however. If there are N items, TBB, depending on the grainsize, schedules chunks of work to be done by different threads. For each of these chunks of work, we figure out how many autodiff variables there are on the input and make copies of them on the specific thread-local stacks. We make sure that any autodiff done on these copied variables does not affect the autodiff variables on the main stack. Without this copying, we would have to deal with race conditions. |
| 180 | + |
| 181 | +As such, there are up to N + M * P autodiff variable copies performed for each reduce_sum where: |
| 182 | + |
| 183 | +1. N is the number of individual operations over which we are reducing |
| 184 | +2. M is the number of blocks of work that TBB sections work up into |
| 185 | + |
| 186 | +and |
| 187 | + |
| 188 | +3. P is the number of autodiff variables in ```(arg1, arg2, ...)``` (... meaning all the trailing arguments) |
| 189 | + |
| 190 | +reduce_sum was designed specifically for the case that ```N >> P```. |
| 191 | + |
| 192 | +Regarding the first term (N), if ```x``` does not contain autodiff types, no autodiff copy is performed. |
| 193 | + |
| 194 | +M (the number of blocks of works that TBB chooses) will be decided by the number of work individual operations (N) and the grainsize. It should roughly be N / grainsize. If grainsize is small compared to the number of autodiff variables in the input arguments, then the overhead of copying disconnecting the nested autodiff stack from the main one will probably become evident. Experimentally, a large grainsize can be detrimental to performance overall, though, so grainsize will probably have to be implemented manually via user experimentation. |
| 195 | + |
| 196 | +Regarding P, consider the variables defined by the C++ code: |
| 197 | +``` |
| 198 | +var a = 5.0; // one autodiff variable |
| 199 | +double b = 5.0; // zero autodiff variables |
| 200 | +int c = 5; // zero autodiff variables |
| 201 | +std::vector<var> d = {1.0, 2.0} // two autodiff variables |
| 202 | +std::vector<double> e = {1.0, 2.0} // zero autodiff variables |
| 203 | +std::vector<int> f = {1, 2} // zero autodiff variables |
| 204 | +Eigen::Matrix<var, Eigen::Dynamic, 1> g(3); // three autodiff variables |
| 205 | +Eigen::Matrix<double, Eigen::Dynamic, 1> h(3); // zero autodiff variables |
| 206 | +Eigen::Matrix<var, 1, Eigen::Dynamic> i(4); // four autodiff variables |
| 207 | +Eigen::Matrix<double, 1, Eigen::Dynamic> j(4); // zero autodiff variables |
| 208 | +Eigen::Matrix<var, Eigen::Dynamic, Eigen::Dynamic> k(5, 5); // twenty-five autodiff variables |
| 209 | +Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> l(5, 5); // zero autodiff variables |
| 210 | +``` |
| 211 | + |
| 212 | +Now, together as a collection of arguments: |
| 213 | +``` |
| 214 | +(a, b, c) // P = 1 |
| 215 | +(e, g, i, l) // P = 7 |
| 216 | +(a, d, g) // P = 6 |
| 217 | +``` |
| 218 | + |
| 219 | +Repeat autodiff arguments are counted multiple times: |
| 220 | +``` |
| 221 | +(a, a, a) // P = 3 |
| 222 | +``` |
| 223 | + |
| 224 | +# Drawbacks |
| 225 | +[drawbacks]: #drawbacks |
| 226 | + |
| 227 | +The main drawback of this new function is the danger that this is replaced by some more advanced and easier to use parallelization mechanism. If that happens, it will be very easy to implement a failsafe backwards compatibility version of reduce_sum that leans on the fact that: |
| 228 | + |
| 229 | +``` |
| 230 | +real sum = reduce_sum(func, x, s1, s2, ...) |
| 231 | +``` |
| 232 | + |
| 233 | +can be replaced by: |
| 234 | + |
| 235 | +``` |
| 236 | +real sum = func(x, 1, size(x), s1, s2, ...) |
| 237 | +``` |
| 238 | + |
| 239 | +where ```func``` was always provided by the user. |
| 240 | + |
| 241 | +In isolated testing, we can get linear speedup with the number of cores. In practical models, the speedups seem much more limited. See discussions in https://discourse.mc-stan.org/t/parallel-autodiff-v4/13111. For memory bound problems, it is difficult to get more than a 2x or 3x speedup. For larger problems, more is possible. In any case, ```reduce_sum``` is just as efficient as ```map_rect```. |
| 242 | + |
| 243 | +Choosing grainsize (the ideal number of terms the tbb scheduler will try to compute in each reduce) is a tricky thing. It is possible for the grainsize to be too small (in which case autodiff overhead slows down the sum) or get the grainsize too large (in which case things slow down too -- probably because of memory usage). We will have to add to the docs instructions on how to pick this, and perhaps an example of performance vs. grainsize for a given model and number of threads. |
| 244 | + |
| 245 | +# Rationale and alternatives |
| 246 | +[rationale-and-alternatives]: #rationale-and-alternatives |
| 247 | + |
| 248 | +The reduce_sum parallelization mechanism is desirable largely because it doesn't require any large reworking of the autodiff system and can be implemented with threads in the TBB pretty painlessly. |
| 249 | + |
| 250 | +In the future it is possible that the scalar output assumption is lifted, but that will only realistically be reasonable once it is possible to do parallel computions during the reverse mode sweep. This currently isn't possible with Stan's autodiff. |
| 251 | + |
| 252 | +There was some discussion of backends for this functionality (MPI vs. threading). There is no reason this could not be ported to work with MPI, though the performance characteristics would change. |
| 253 | + |
| 254 | +# Prior art |
| 255 | +[prior-art]: #prior-art |
| 256 | + |
| 257 | +The previous example of parallelization in Stan is map-rect. Argument packing made using map-rect difficult, and the goal here is to overcome that difficulty. |
| 258 | + |
| 259 | +# Unresolved questions |
| 260 | +[unresolved-questions]: #unresolved-questions |
| 261 | + |
0 commit comments