Computational Bottlenecks of Training Small-scale Large Language Models

TL;DR

This paper investigates the training bottlenecks of Small-scale Language Models (up to 2B parameters) on cloud hardware, using a cost-aware framework that blends tokens-per-second with loss-per-token to derive Loss/Dollar. Through a broad grid search across GPU types, batch sizes, communication schemes, and attention mechanisms, the authors quantify practical trade-offs for cost-efficient training. Key findings show that FlashAttention substantially improves cost-efficiency for SLMs, that expensive GPUs are not always worthwhile, and that Distributed Data Parallel remains effective for smaller models while Fully Sharded Data Parallel excels for larger models with bigger batch sizes. The work provides actionable, hardware-aware guidance to researchers and organizations with limited resources, aiming to democratize access to training capable SLMs.

Abstract

While large language models (LLMs) dominate the AI landscape, Small-scale large Language Models (SLMs) are gaining attention due to cost and efficiency demands from consumers. However, there is limited research on the training behavior and computational requirements of SLMs. In this study, we explore the computational bottlenecks of training SLMs (up to 2B parameters) by examining the effects of various hyperparameters and configurations, including GPU type, batch size, model size, communication protocol, attention type, and the number of GPUs. We assess these factors on popular cloud services using metrics such as loss per dollar and tokens per second. Our findings aim to support the broader adoption and optimization of language model training for low-resource AI research institutes.

Figures (5)

• Figure 1: FlashAttention is more cost-efficient for smaller models. Maximum $\text{Token}/\text{Dollar}$ across GPU-types, GPU-number, and communication type when we use FlashAttention. FlashAttention shows a significant $\text{Token}/\text{Dollar}$ improvement over vanilla attention in smaller models and batch sizes. OOM runs are shown with 0. $\text{Token}/\text{Dollar}=$1k means training for 1k tokens costs $1.
• Figure 2: H100 GPUs are not cost-efficient for training SLMs. Maximum $\text{Token}/\text{Dollar}$ for different GPU-type across batch-size and communication types. We use FlashAttention.
• Figure 3: DDP is the best scheme for training SLMs. Maximum $\text{Token}/\text{Dollar}$ for different GPU-nodes on A100-80GB across different batch-sizes. We use FlashAttention in our models. for a single node, we use 2, 4, and 8 GPUs while for 2 and 4 nodes we use 16 and 32 GPUs respectively.
• Figure 4: For SLMs increasing global batch size saturates cost-efficiency before GPU memory is fully utilized. Maximum $\text{Token}/\text{Dollar}$ for different global batch-sizes across across different GPU-types, GPU numbers, and various per-device batch-sizes. We use FlashAttention in our models.
• Figure 5: Curve fitting to extract our SLM model sizes using LLaMa-2/3 models. Left: Hidden dimension extraction. Right: Number of decoder blocks.