TBA: Faster Large Language Model Training Using SSD-Based Activation Offloading
Authors:
Kun Wu,
Jeongmin Brian Park,
Xiaofan Zhang,
Mert Hidayetoğlu,
Vikram Sharma Mailthody,
Sitao Huang,
Steven Sam Lumetta,
Wen-mei Hwu
Abstract:
The growth rate of the GPU memory capacity has not been able to keep up with that of the size of large language models (LLMs), hindering the model training process. In particular, activations -- the intermediate tensors produced during forward propagation and reused in backward propagation -- dominate the GPU memory use. To address this challenge, we propose TBA to efficiently offload activations…
▽ More
The growth rate of the GPU memory capacity has not been able to keep up with that of the size of large language models (LLMs), hindering the model training process. In particular, activations -- the intermediate tensors produced during forward propagation and reused in backward propagation -- dominate the GPU memory use. To address this challenge, we propose TBA to efficiently offload activations to high-capacity NVMe SSDs. This approach reduces GPU memory usage without impacting performance by adaptively overlapping data transfers with computation. TBA is compatible with popular deep learning frameworks like PyTorch, Megatron, and DeepSpeed, and it employs techniques such as tensor deduplication, forwarding, and adaptive offloading to further enhance efficiency. We conduct extensive experiments on GPT, BERT, and T5. Results demonstrate that TBA effectively reduces 47% of the activation peak memory usage. At the same time, TBA perfectly overlaps the I/O with the computation and incurs negligible performance overhead. We introduce the recompute-offload-keep (ROK) curve to compare the TBA offloading with other two tensor placement strategies, keeping activations in memory and layerwise full recomputation. We find that TBA achieves better memory savings than layerwise full recomputation while retaining the performance of keeping the activations in memory.
△ Less
Submitted 19 August, 2024;
originally announced August 2024.
ASAP: Accelerated Short-Read Alignment on Programmable Hardware
Authors:
Subho S. Banerjee,
Mohamed El-Hadedy,
Jong Bin Lim,
Zbigniew T. Kalbarczyk,
Deming Chen,
Steve Lumetta,
Ravishankar K. Iyer
Abstract:
The proliferation of high-throughput sequencing machines ensures rapid generation of up to billions of short nucleotide fragments in a short period of time. This massive amount of sequence data can quickly overwhelm today's storage and compute infrastructure. This paper explores the use of hardware acceleration to significantly improve the runtime of short-read alignment, a crucial step in preproc…
▽ More
The proliferation of high-throughput sequencing machines ensures rapid generation of up to billions of short nucleotide fragments in a short period of time. This massive amount of sequence data can quickly overwhelm today's storage and compute infrastructure. This paper explores the use of hardware acceleration to significantly improve the runtime of short-read alignment, a crucial step in preprocessing sequenced genomes. We focus on the Levenshtein distance (edit-distance) computation kernel and propose the ASAP accelerator, which utilizes the intrinsic delay of circuits for edit-distance computation elements as a proxy for computation. Our design is implemented on an Xilinx Virtex 7 FPGA in an IBM POWER8 system that uses the CAPI interface for cache coherence across the CPU and FPGA. Our design is $200\times$ faster than the equivalent C implementation of the kernel running on the host processor and $2.2\times$ faster for an end-to-end alignment tool for 120-150 base-pair short-read sequences. Further the design represents a $3760\times$ improvement over the CPU in performance/Watt terms.
△ Less
Submitted 23 May, 2018; v1 submitted 6 March, 2018;
originally announced March 2018.
