Comprehensive Performance Modeling and System Design Insights for Foundation Models

TL;DR

This paper analyzes how different transformer models impact HPC system design, highlighting the importance of tailored parallelism strategies and system features for optimal performance in generative AI applications.

Contribution

It introduces a performance model that explores the complex interactions between transformer types, parallelization strategies, and HPC system features, providing insights for optimized system design.

Findings

Large Language Models perform well with 3D parallelism.

Long-sequence transformers require 4D parallelism.

System features need to be tailored to transformer type and scale.

Abstract

Generative AI, in particular large transformer models, are increasingly driving HPC system design in science and industry. We analyze performance characteristics of such transformer models and discuss their sensitivity to the transformer type, parallelization strategy, and HPC system features (accelerators and interconnects). We utilize a performance model that allows us to explore this complex design space and highlight its key components. We find that different transformer types demand different parallelism and system characteristics at different training regimes. Large Language Models are performant with 3D parallelism and amplify network needs only at pre-training scales with reduced dependence on accelerator capacity and bandwidth. On the other hand, long-sequence transformers, representative of scientific foundation models, place a more uniform dependence on network and capacity with necessary 4D parallelism. Our analysis emphasizes the need for closer performance modeling of different transformer types keeping system features in mind and demonstrates a path towards this. Our code is available as open-source.