Automated optimization of distributed tensor programs for large DNN training and inference

Abstract

Modern large-scale machine learning relies heavily on complex computations expressed as tensor programs, typically distributed across diverse hardware clusters. Optimizing these programs for performance presents significant challenges due to model scale, hardware heterogeneity, intricate parallelism strategies, and dynamic execution patterns. This thesis focuses on automated optimization of distributed tensor programs, developing systems and methodologies to automatically enhance the performance of both training and inference for large models.

First, we introduce TAG for optimizing DNN training, viewing the process as deploying a tensor program onto heterogeneous devices. TAG employs a graph neural network (GNN) processing both the computation graph (tensor operations) and device topology, coupled with Monte-Carlo tree search, to automatically derive optimized distributed execution plans. It further automates communication optimization via sufficient factor broadcasting. TAG demonstrates up to 4.56x speed-up, showcasing effective automated optimization for heterogeneous training programs.

Second, addressing performance bottlenecks in SPMD parallelism for trillion-parameter models, we present HiDup, that optimizes distributed tensor programs with computation-communication overlapping. We propose interleaving the execution of two microbatches, where computation of a microbatch overlaps communication of the other. A dynamic programming algorithm automates the search for optimal tensor sharding strategies within this restructured program, maximizing overlap and achieving up to 61% training speed-up.

Third, we develop HAP to further accelerate SPMD tensor program on heterogeneous clusters. We formulate the partitioning of the tensor program as a program synthesis problem, automatically generating an optimized distributed program from a single-device program using A*-based search. HAP automatically co-optimizes tensor sharding strategy, optimal device-specific sharding ratios (via linear programming), and efficient tensor communication methods. HAP achieves up to 2.41x speed-up, demonstrating effective automated optimization for SPMD programs on heterogeneous hardware.

Finally, we tackle the automated optimization of dynamic workflows composed of multiple tensor programs. Conditional execution in these workflows hinders static optimization. We introduce DyOrc, an workflow serving system deploying components as independently scalable services. DyOrc features: (i) speculative scheduling for effective batching of tensor computations; (ii) multi-tier message passing for efficient inter-program tensor communication; and (iii) proactive loading to hide startup latency. DyOrc improves serving latency by 4-198% on diverse dynamic workloads.