ProTrain: Efficient LLM Training via Memory-Aware Techniques

• URL: http://arxiv.org/abs/2406.08334v1
• Date: Wed, 12 Jun 2024 15:40:06 GMT
• Title: ProTrain: Efficient LLM Training via Memory-Aware Techniques
• Authors: Hanmei Yang, Jin Zhou, Yao Fu, Xiaoqun Wang, Ramine Roane, Hui Guan, Tongping Liu,
• Abstract summary: This paper proposes ProTrain, a novel training system that balances memory usage and performance by coordinating memory, computation, and IO. ProTrain improves training throughput by 1.43$times$ to 2.71$times compared to the SOTA training systems.
• Score: 18.30799115938978
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: It is extremely memory-hungry to train Large Language Models (LLM). To solve this problem, existing work exploits the combination of CPU and GPU for the training process, such as ZeRO-Offload. Such a technique largely democratizes billion-scale model training, making it possible to train with few consumer graphics cards. However, based on our observation, existing frameworks often provide coarse-grained memory management and require experienced experts in configuration tuning, leading to suboptimal hardware utilization and performance. This paper proposes ProTrain, a novel training system that intelligently balances memory usage and performance by coordinating memory, computation, and IO. ProTrain achieves adaptive memory management through Chunk-Based Model State Management and Block-Wise Activation Management, guided by a Memory-Aware Runtime Profiler without user intervention. ProTrain does not change the training algorithm and thus does not compromise accuracy. Experiments show that ProTrain improves training throughput by 1.43$\times$ to 2.71$\times$ compared to the SOTA training systems.

Related papers

• APOLLO: SGD-like Memory, AdamW-level Performance [61.53444035835778]
  Large language models (LLMs) are notoriously memory-intensive during training. Various memory-efficient Scals have been proposed to reduce memory usage. They face critical challenges: (i) costly SVD operations; (ii) significant performance trade-offs compared to AdamW; and (iii) still substantial memory overhead to maintain competitive performance.
  arXiv  Detail & Related papers  (2024-12-06T18:55:34Z)
• SSDTrain: An Activation Offloading Framework to SSDs for Faster Large Language Model Training [13.283682311968752]
  SSDTrain is an adaptive activation framework offloading to high-capacity GPU memory. It is compatible with popular deep learning frameworks like PyTorch, Megatron, and DeepSpeed. Results demonstrate that SSDTrain reduces 47% of the activation peak memory usage.
  arXiv  Detail & Related papers  (2024-08-19T14:09:48Z)
• MEMO: Fine-grained Tensor Management For Ultra-long Context LLM Training [24.066283519769968]
  Large Language Models (LLMs) have been trained using extended context lengths to foster more creative applications. We propose MEMO, a novel framework for fine-grained activation memory management. MeMO achieves an average of 1.97x and 1.80x MFU compared to Megatron-LM and DeepSpeed.
  arXiv  Detail & Related papers  (2024-07-16T18:59:49Z)
• Block Selective Reprogramming for On-device Training of Vision Transformers [12.118303034660531]
  We present block selective reprogramming (BSR) in which we fine-tune only a fraction of total blocks of a pre-trained model. Compared to the existing alternatives, our approach simultaneously reduces training memory by up to 1.4x and compute cost by up to 2x.
  arXiv  Detail & Related papers  (2024-03-25T08:41:01Z)
• GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection [133.45193150403537]
  Training Large Language Models (LLMs) presents significant memory challenges due to the growing size of weights and GPU states. In this work, we propose Gradient Low-Rank Projection (GaLore) as a memory-efficient training strategy. Our 8-bit GaLore further reduces memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline.
  arXiv  Detail & Related papers  (2024-03-06T07:29:57Z)
• Time-, Memory- and Parameter-Efficient Visual Adaptation [75.28557015773217]
  We propose an adaptation method which does not backpropagate gradients through the backbone. We achieve this by designing a lightweight network in parallel that operates on features from the frozen, pretrained backbone.
  arXiv  Detail & Related papers  (2024-02-05T10:55:47Z)
• Fast Machine Unlearning Without Retraining Through Selective Synaptic Dampening [51.34904967046097]
  Selective Synaptic Dampening (SSD) is a fast, performant, and does not require long-term storage of the training data. We present a novel two-step, post hoc, retrain-free approach to machine unlearning which is fast, performant, and does not require long-term storage of the training data.
  arXiv  Detail & Related papers  (2023-08-15T11:30:45Z)
• Task-oriented Memory-efficient Pruning-Adapter [3.0751447761822903]
  We propose a task-oriented Pruning-Adapter method that achieve a high memory efficiency of training and memory. No significant decrease in accuracy in GLUE tasks, achieving training and inference efficiency at the same time.
  arXiv  Detail & Related papers  (2023-03-26T12:18:00Z)
• Failure Tolerant Training with Persistent Memory Disaggregation over CXL [7.700500756012469]
  This paper proposes TRAININGCXL that can efficiently process large-scale recommendation datasets in the pool of disaggregated memory. To this end, we integrate persistent memory (PMEM) and GPU into a cache-coherent domain as Type-2. The evaluation shows that TRAININGCXL achieves 5.2x training performance improvement and 76% energy savings, compared to the modern PMEM-based recommendation systems.
  arXiv  Detail & Related papers  (2023-01-14T05:59:07Z)
• On-Device Training Under 256KB Memory [62.95579393237751]
  We propose an algorithm-system co-design framework to make on-device training possible with only 256KB of memory. Our framework is the first solution to enable tiny on-device training of convolutional neural networks under 256KB and 1MB Flash.
  arXiv  Detail & Related papers  (2022-06-30T17:59:08Z)
• Online Convolutional Re-parameterization [51.97831675242173]
  We present online convolutional re- parameterization (OREPA), a two-stage pipeline, aiming to reduce the huge training overhead by squeezing the complex training-time block into a single convolution. Compared with the state-of-the-art re-param models, OREPA is able to save the training-time memory cost by about 70% and accelerate the training speed by around 2x. We also conduct experiments on object detection and semantic segmentation and show consistent improvements on the downstream tasks.
  arXiv  Detail & Related papers  (2022-04-02T09:50:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.