Performance of Small Language Model Pretraining on FABRIC: An Empirical Study

TL;DR

This study addresses democratizing small-language-model pretraining on academic testbeds by evaluating SLM pretraining on FABRIC using commodity GPUs. It systematically compares data, intra-operator, and inter-operator parallelism, including Alpa's shard and pipeshard plans, across homogeneous and heterogeneous GPU clusters with cross-site latency. The results show that Pipeshard generally yields the best training performance in two-site deployments, while single-site training often favors data-parallel approaches when memory permits; ZeRO2 serves as a memory-efficient fallback in heterogeneous setups. A practical technique-selection algorithm is proposed to choose the most efficient pretraining configuration, enabling faster generation of domain-specific embeddings for vector databases and semantic search on academic infrastructure.

Abstract

Large language models (LLMs) require enormous computing power to pretrain on massive datasets. When limited datasets are available, smaller-sized LLMs are better choice to pretrain (on user-specified datasets) by following the scaling laws of LLMs. Using pretrained models, vector embeddings can be generated for raw data and stored using vector databases to support modern AI applications and semantic search. In this work, we investigate the performance of pretraining techniques for smaller-sized LLMs on an experimental testbed (with commodity GPUs) available to academic users at no charge. We consider data parallelism, intra-operator parallelism, and inter-operator/pipeline parallelism, and their combinations for pretraining. We set up different GPU clusters with homogeneous and heterogeneous GPU hardware. Furthermore, we investigate the impact of network latency on pretraining performance especially when GPUs are geographically distributed. We used GPT-2 medium and large models and pretrained them using open-source packages, namely, Alpa and Ray. We observed that Alpa's execution plans that collectively optimized intra-operator and inter-operator/pipeline parallelism consistently performed the best when GPUs were geographically distributed. This was especially true when the network latencies were in 10's of milliseconds. Based on the insights gained from the experiments, we propose a systematic approach for selecting the appropriate pretraining technique to achieve high training performance/lower execution time as well as to reduce the number of GPUs used.

Figures (7)

• Figure 1: Examples of different parallelization techniques
• Figure 2: A FABRIC slice consisting of two VMs; each VM has 2 GPUs attached to it; the VMs are connected by L2STS
• Figure 3: SLM pretraining performance on TACC ($\mathbf{\times}$ denotes failed execution)
• Figure 4: SLM pretraining on UTAH-GPN ($\mathbf{\times}$ denotes failed execution)
• Figure 5: SLM pretraining on UTAH-MASS ($\mathbf{\times}$ denotes failed execution)