Training Ultra Long Context Language Model with Fully Pipelined Distributed Transformer

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Yao, Jinghan, Jacobs, Sam Ade, Tanaka, Masahiro, Ruwase, Olatunji, Subramoni, Hari, Panda, Dhabaleswar K.
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915284756987904
author Yao, Jinghan
Jacobs, Sam Ade
Tanaka, Masahiro
Ruwase, Olatunji
Subramoni, Hari
Panda, Dhabaleswar K.
author_facet Yao, Jinghan
Jacobs, Sam Ade
Tanaka, Masahiro
Ruwase, Olatunji
Subramoni, Hari
Panda, Dhabaleswar K.
contents Large Language Models (LLMs) with long context capabilities are integral to complex tasks in natural language processing and computational biology, such as text generation and protein sequence analysis. However, training LLMs directly on extremely long contexts demands considerable GPU resources and increased memory, leading to higher costs and greater complexity. Alternative approaches that introduce long context capabilities via downstream finetuning or adaptations impose significant design limitations. In this paper, we propose Fully Pipelined Distributed Transformer (FPDT) for efficiently training long-context LLMs with extreme hardware efficiency. For GPT and Llama models, we achieve a 16x increase in sequence length that can be trained on the same hardware compared to current state-of-the-art solutions. With our dedicated sequence chunk pipeline design, we can now train 8B LLM with 2 million sequence length on only 4 GPUs, while also maintaining over 55% of MFU. Our proposed FPDT is agnostic to existing training techniques and is proven to work efficiently across different LLM models.
format Preprint
id arxiv_https___arxiv_org_abs_2408_16978
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Training Ultra Long Context Language Model with Fully Pipelined Distributed Transformer
Yao, Jinghan
Jacobs, Sam Ade
Tanaka, Masahiro
Ruwase, Olatunji
Subramoni, Hari
Panda, Dhabaleswar K.
Distributed, Parallel, and Cluster Computing
Artificial Intelligence
Machine Learning
Large Language Models (LLMs) with long context capabilities are integral to complex tasks in natural language processing and computational biology, such as text generation and protein sequence analysis. However, training LLMs directly on extremely long contexts demands considerable GPU resources and increased memory, leading to higher costs and greater complexity. Alternative approaches that introduce long context capabilities via downstream finetuning or adaptations impose significant design limitations. In this paper, we propose Fully Pipelined Distributed Transformer (FPDT) for efficiently training long-context LLMs with extreme hardware efficiency. For GPT and Llama models, we achieve a 16x increase in sequence length that can be trained on the same hardware compared to current state-of-the-art solutions. With our dedicated sequence chunk pipeline design, we can now train 8B LLM with 2 million sequence length on only 4 GPUs, while also maintaining over 55% of MFU. Our proposed FPDT is agnostic to existing training techniques and is proven to work efficiently across different LLM models.
title Training Ultra Long Context Language Model with Fully Pipelined Distributed Transformer
topic Distributed, Parallel, and Cluster Computing
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2408.16978