Search papers, labs, and topics across Lattice.
This paper introduces Einsummable, a system that automates the distribution of AI computations across multi-GPU servers by modeling operations as relational joins followed by aggregations. By optimizing the decomposition of computations to minimize communication costs, Einsummable can outperform manually tuned implementations, achieving a geometric-mean runtime of 8.97 ms on LLaMA transformer blocks compared to 13.80 ms for hand-tuned PyTorch. The approach eliminates the need for manual device assignments and communication operations, streamlining the parallelization process for AI workloads.
Einsummable achieves a remarkable 35% speedup over hand-tuned PyTorch implementations by automatically optimizing multi-GPU computation distribution without manual intervention.
Distributing an AI computation across the GPUs of a multi-GPU server is one of the central problems in systems-for-AI. We present Einsummable, a prototype system that accepts a PyTorch-like description of an AI computation and automatically distributes it across a multi-GPU server, with no device assignments, sharding annotations, or communication operations written by the programmer. Einsummable models every operation as a relational join followed by an aggregation over tensor relations, in which the tuples contain sub-tensors. Each operation exposes its possible decompositions through what we call"join-agg specs". An optimizer then selects decompositions across the whole computation to minimize a communication-cost proxy. Because it searches decompositions rather than a menu of named strategies, Einsummable discovers plans that mesh-based auto-parallelizers cannot. Each decomposed operation is implemented by synthesizing an exchange program, which is a topology-aware generalization of Volcano's exchange operator. Einsummable invokes no canned collectives: all communication and aggregation is special-purpose, derived at compile time. Despite being fully automatic, Einsummable can outperform custom-designed implementations. For example, on LLaMA transformer blocks on an eight-GPU A100 server, Einsummable achieves a geometric-mean runtime of 8.97 ms, versus 13.80 ms for hand-tuned PyTorch and 15.90 ms for vLLM.
