Comprehensive Performance Modeling and System Design Insights for Foundation Models

TL;DR

This analysis emphasizes the need for closer performance modeling of different transformer types keeping system features in mind and demonstrates a path towards this.

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.

Figures (11)

• Figure 1: On $16384$B200 with NVS domain size $n_{\text{NVS}}=8$, microbatch size $1$ for GPT3-1T with 1D TP: (top) Configurations (Config.) chosen with PP fixed to $n_p = 64$ and the others varying. Example, Config. D is $(m,n_t,n_d,n_p)=(128,8,32,64)$ and so on. (bottom) Time for each configuration broken down by the time spent on compute, memory accesses, various communications (DP, TP, PP) and PP bubbles. In config. D (TP $n_t=8$), we observe a local minimum in time per iteration.
• Figure 1: Time for $\mathcal{AG}$ as a function of communication volume on 32 A100 GPUs. We compare empirical numbers on Perlmutter through NCCL tests nccltests to our theoretical formulae for different fast domain (NVLink) sizes. For NVL 2, 2 GPUs per node are used and for NVL 4, 4 GPUs per node are used. We see that more GPUs per node effectively increases the SlingShot bandwidth, leading to smaller times.
• Figure 2: Following Fig, \ref{['fig:rationale']}, we change the following: we fix TP $n_t=8$ and vary PP and DP on two different NVS domain sizes. (top) We observe a local minimum at PP $n_p=64$. (bottom) For larger NVS, we observe that the local minimum shifts to low PP $n_p=2$, with NVS used to hide DP costs.
• Figure 2: (a) For GPT3-1T with 2D TP, we fix TP $n_t=32, n_p=1$ and vary $n_1, n_2$ to get the the first five configurations and then switch to $n_p = 128$ and repeat the same with $n_t=8$ for the last four configurations. We use a B200 GPU with NVS domain size $n_{\text{NVS}}=64$, global batch size $4096$ on $16384$ GPUs and also show time for each configuration broken down by the components. We see similar behavior as 2D TP SUMMA but the memory consumed is very high and hence the large PP configurations are chosen. (b) For ViT with 2D TP, we first fix $n_t=16$ and vary $n_1,n_2$, then switch to $n_p=16$ and repeat the same. The memory used is sensitive to the parallelization regime.
• Figure 3: For GPT3-1T with 2D TP SUMMA, we fix TP $(n_t,n_p)=(32,1)$ and vary $n_1, n_2$ to get the the first five configurations and then switch to $(n_t,n_p)=(8,128)$ and repeat the same for the rest. We show two NVS domain sizes. With larger NVS, large DP (low PP) is preferred.