Speaker: Zhihao Jia

Affiliation: Stanford University

Date and location: Friday, May 17, 12:30-1:30pm; Wozniak Lounge (430 Soda Hall)

Abstract: Existing deep learning frameworks deploy DNN architectures on modern hardware devices by applying a sequence of heuristic optimizations. For example, current frameworks use data and model parallelism to parallelize DNN training across distributed clusters, and use rule-based operator fusions to optimize DNN computation graphs. These heuristic approaches achieve improved runtime performance in general but miss subtle optimizations for particular DNN architectures.

In this talk, I will present two projects that use search-based approaches to optimize deep learning computation. First, I will introduce SOAP, a comprehensive search space of parallelization strategies for DNN training that includes strategies to parallelize a DNN in the Sample, Operator, Attribute, and Parameter dimensions. I will present FlexFlow, a deep learning engine that uses guided randomized search of the SOAP space to automatically discover fast parallelization strategies, outperforming existing data/model parallelism by up to 3.3x. Second, I will present XFlow, a DNN computation graph optimizer with automatically generated graph substitutions. For a DNN architecture, XFlow uses these auto-generated substitutions to constitute a search space of candidate computation graphs, and uses cost-based backtracking search to discover optimized graphs, which outperform graphs generated by existing rule-based optimizers by up to 2.9x. Finally, I will conclude the talk by discussing future work and research opportunities enabled by emerging DNN architectures, such as graph neural networks.

Bio: Zhihao Jia is a PhD candidate in the Computer Science Department at Stanford University, working with Prof. Matei Zaharia and Prof. Alex Aiken. His research interests lie at the intersection of systems and deep learning. In particular, he is interested in designing automated systems for optimizing deep learning computation.