Part of the TPC Seminar Series


Speaker: Sajal Dash, Research Scientist at Oak Ridge National Laboratory
Date: Wednesday, May 8, 2024
Time: 10:00  A.M. to 11:15 A.M.  (Central Time)
Location: Virtual

Abstract:

Large language models (LLMs) have demonstrated remarkable success as foundational models, benefiting various downstream applications through fine-tuning. Recent studies on loss scaling have demonstrated the superior performance of larger LLMs compared to their smaller counterparts. Nevertheless, training LLMs with billions of parameters poses significant challenges and requires considerable computational resources. For example, training a one trillion parameter GPT-style model on 20 trillion tokens requires a staggering 120 million exaflops of computation. This research explores efficient distributed training strategies to extract this computation from Frontier, the world’s first exascale supercomputer dedicated to open science. We enable and investigate various model and data parallel training techniques, such as tensor parallelism, pipeline parallelism, and sharded data parallelism, to facilitate training a trillion-parameter model on Frontier. We empirically assess these techniques and their associated parameters to determine their impact on memory footprint, communication latency, and GPU’s computational efficiency. We analyze the complex interplay among these techniques and find a strategy to combine them to achieve high throughput through hyperparameter tuning. We have identified efficient strategies for training large LLMs of varying sizes through empirical analysis and hyperparameter tuning. For 22 Billion, 175 Billion, and 1 Trillion parameters, we achieved GPU throughputs of 38.38%, 36.14%, and 31.96%, respectively. For the training of the 175 Billion parameter model and the 1 Trillion parameter model, we achieved 100% weak scaling efficiency on 1024 and 3072 MI250X GPUs, respectively. We also achieved strong scaling efficiencies of 89% and 87% for these two models.

Biography:

Sajal Dash is a research scientist at Oak Ridge National Laboratory specializing in HPC and AI. He is exploring distributed training methodologies for training large-scale LLMs and their applications in various scientific domains. He recently led an effort by Oak Ridge National Laboratory to develop an optimized training strategy for training LLMs with up to a trillion parameters. Sajal completed his Ph.D. in Computer Science at Virginia Tech. His Ph.D. dissertation titled “Exploring the Landscape of Big Data Analytics Through Domain-Aware Algorithm Design” focused on solving large-scale domain problems by leveraging domain knowledge with properties of big data and solving a big data problem in cancer biology by efficiently distributing the combinatorial workload across nodes with regularizing memory access patterns. His research also includes catastrophic forgetting in continual learning and a mixture of experts in LLMs.