Speculative dynamic HMC

[wds15]

Summary:

The dynamic HMC algorithm can be parallelized to run the forward and backward sweeps in independence. See for a description and a decent evaluation the thread on discourse:

https://discourse.mc-stan.org/t/parallel-dynamic-hmc-merits/10895

By using 2 cores we can be speed up things by 35%.

Description:

A prototype implementation is here:

https://github.com/stan-dev/stan/blob/parallel-nuts-2/src/stan/mcmc/hmc/nuts/base_nuts.hpp#L159

That version is based on the 2.20.0 sampler.

The branch is not runnable without a special setup of math due to the requirement of the TBB. A variant of the branch https://github.com/stan-dev/math/tree/parallel-ad-tape-3 can make this one run, but no attempt is made at the moment to give clear instructions.

This issue is here to document that we can get 35% speedup if we do a speculative HMC

Reproducible Steps:

See discourse for all of this.

Current Output:

Expected Output:

Faster samples.

Additional Information:

Provide any additional information here.

Current Version:

v2.20.0

[yizhang-yiz]

Is there a way to evaluate memory footprint of presampling to max treedepth?

[wds15]

more memory is needed, for sure... don't know by heart.

[SteveBronder]

Already iterated this but It may be worth throwing together a branch of Stan that can use a bigger DGP of the cmdstan benchmarks + this so we can get an idea of how it works across a bunch of models

[SteveBronder]

Tangent to this but I liked Michaels suggestion of refactoring Stan a bit so that it can handle multiple chains and sharing data between them. That would also let us do weird things like adjust the tree depth and max delta of the model based on how other chains were doing

[yizhang-yiz]

Tangent to this but I liked Michaels suggestion of refactoring Stan a bit so that it can handle multiple chains and sharing data between them. That would also let us do weird things like adjust the tree depth and max delta of the model based on how other chains were doing

+1 to this. Could you point me to where he made the suggestion?

[wds15]

He made the suggestion in the parallel HMC thread... and I do agree.

I do want to look into that - imagine you fire up 6 chains and pool the warmup info. That would speedup sampling a lot. The crux is that stan is not tailored at all to handle multiple chains at once. The one-chain assumption is very hard coded into the services (which is why speculative HMC is easier from this perspective since it speeds up one chain).

My plan is to write a version which fires of many chains for warmup, pools that info and then discards the info to run just one chain. If you do that, then you can basically do one run for warmup and then use the restarting facilities which we have to get what you really want (multiple chains starting from a common warmup). At least this is a first quick way to test the merits of all of this.

[yizhang-yiz]

I still don't understand get why you said using 4-cores in this setting is not possible. Could you elaborate?

[SteveBronder]

I think the biggest wugh with multiple chains in one stan program is that our upstream services assume there is one chain.

It would look a bit odd but essentially we would just want an 'algorithm' class that takes in N instances of a model then every few iterations pulls some meta info from them and sends some metainfo back

[SteveBronder]

I was thinking about that when I was reading over the hints paper, they seem to do something like that but with the temperature for their annealing scheme (it wasn't exactly annealing but similar)

[SteveBronder]

@yizhang-yiz see bottom of here

https://discourse.mc-stan.org/t/parallel-dynamic-hmc-merits/10895/13?u=stevo15025

[wds15]

Why can’t we have a multi chain warmup which are terminated once done doing warmup. Then the CmdStan or whatever else interface fires up multiple chains with the so obtained adapted sampler. This way the services still handle only one chain at a time and it’s just the interface which makes this work...or some callback in Stan services itself.

We should find a way there.

[SteveBronder]

^I also like that!

[wds15]

@yizhang-yiz Our dynamic HMC implementation goes randomly forward and backward in time. Right now that is done serially. The speculative HMC allows to go simultaneous forward and backward in time. Since there is only backward and forward that's it - no more cores can be spend on that. Makes sense?

[bob-carpenter]

I started a discussion on Discourse about how to evaluate parallelization of log density or MCMC algorithm.

https://discourse.mc-stan.org/t/evaluating-parallelization-performance/10962

[bob-carpenter]

On Sep 12, 2019, at 7:48 PM, wds15 ***@***.***> wrote: He made the suggestion in the parallel HMC thread... and I do agree. I do want to look into that - imagine you fire up 6 chains and pool the warmup info. That would speedup sampling a lot. The crux is that stan is not tailored at all to handle multiple chains at once. The one-chain assumption is very hard coded into the services (which is why speculative HMC is easier from this perspective since it speeds up one chain).

The services code up specific algorithms, but it'd be trivial to write new ones or extend the existing ones for multiple chains.

My plan is to write a version which fires of many chains for warmup, pools that info and then discards the info to run just one chain. If you do that, then you can basically do one run for warmup and then use the restarting facilities which we have to get what you really want (multiple chains starting from a common warmup). At least this is a first quick way to test the merits of all of this.

This would be great.