2019/10/23 by Alexey Radul, Brian Patton, Radul, Alexey +7 · 1 citation
Computer Science · #Advanced Data Storage Technologies #Distributed #Distributed and Parallel Computing Systems #FOS: Computer and information sciences #Machine Learning (cs.LG) #Parallel #Parallel Computing and Optimization Techniques #Programming Languages (cs.PL) #and Cluster Computing (cs.DC)
paper · pdf · doi:10.48550/arxiv.1910.11141
openalex publication_date 2019/10/23 · openalex created_date 2019/11/01 · openalex updated_date 2026/07/28
We present a general approach to batching arbitrary computations for accelerators such as GPUs. We show orders-of-magnitude speedups using our method on the No U-Turn Sampler (NUTS), a workhorse algorithm in Bayesian statistics. The central challenge of batching NUTS and other Markov chain Monte Carlo algorithms is data-dependent control flow and recursion. We overcome this by mechanically transforming a single-example implementation into a form that explicitly tracks the current program point for each batch member, and only steps forward those in the same place. We present two different batching algorithms: a simpler, previously published one that inherits recursion from the host Python, and a more complex, novel one that implemenents recursion directly and can batch across it. We implement these batching methods as a general program transformation on Python source. Both the batching system and the NUTS implementation presented here are available as part of the popular TensorFlow Probability software package.