Dynamic Tensor Rematerialization

Kirisame, Marisa; Lyubomirsky, Steven; Haan, Altan; Brennan, Jennifer; He, Meian; Roesch, Jared; Chen, Tianqi; Tatlock, Zachary

doi:10.48550/arxiv.2006.09616

Cited by 11 publications

(29 citation statements)

References 0 publications

Supporting

Mentioning

Contrasting

Order By: Relevance

“…Techniques for training large-scale models. In addition to parallelization, there are other complementary techniques for training large-scale models, such as memory optimization [8,10,17,19,24,42], communication compression [4,48], and low-precision training [31]. Alpa can incorporate many of these techniques.…”

Section: Related Workmentioning

confidence: 99%

Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning

Zheng¹,

Li²,

Zhang³

et al. 2022

Preprint

View full text Add to dashboard Cite

Alpa automates model-parallel training of large deep learning (DL) models by generating execution plans that unify data, operator, and pipeline parallelism. Existing model-parallel training systems either require users to manually create a parallelization plan or automatically generate one from a limited space of model parallelism configurations, which does not suffice to scale out complex DL models on distributed compute devices. Alpa distributes the training of large DL models by viewing parallelisms as two hierarchical levels: inter-operator and intra-operator parallelisms. Based on it, Alpa constructs a new hierarchical space for massive model-parallel execution plans. Alpa designs a number of compilation passes to automatically derive the optimal parallel execution plan in each independent parallelism level and implements an efficient runtime to orchestrate the two-level parallel execution on distributed compute devices. Our evaluation shows Alpa generates parallelization plans that match or outperform handtuned model-parallel training systems even on models they are designed for. Unlike specialized systems, Alpa also generalizes to models with heterogeneous architectures and models without manually-designed plans. * Lianmin, Zhuohan, and Hao contributed equally. Part of the work was done when Lianmin interned at Amazon and Zhuohan interned at Google.

show abstract

Section: Related Workmentioning

confidence: 99%

Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning

Zheng¹,

Li²,

Zhang³

et al. 2022

Preprint

View full text Add to dashboard Cite

show abstract

“…For example, [Jain et al, 2020b] proposed an Integer Linear Program (ILP) to find the optimal rematerialization strategy suitable for an arbitrary Directed Acyclic Graph (DAG) structure. [Kirisame et al, 2020] presented a cheap dynamic heuristic DTR that relies on scores encouraging to discard (i) heavy tensors (ii) with a long lifetime and (iii) that can be easily recomputed.…”

Section: Rematerialization Of Activationsmentioning

confidence: 99%

Survey on Large Scale Neural Network Training

Gusak¹,

Cherniuk²,

Shilova³

et al. 2022

Preprint

View full text Add to dashboard Cite

Modern Deep Neural Networks (DNNs) require significant memory to store weight, activations, and other intermediate tensors during training. Hence, many models don't fit one GPU device or can be trained using only a small per-GPU batch size. This survey provides a systematic overview of the approaches that enable more efficient DNNs training. We analyze techniques that save memory and make good use of computation and communication resources on architectures with a single or several GPUs. We summarize the main categories of strategies and compare strategies within and across categories. Along with approaches proposed in the literature, we discuss available implementations.

show abstract

“…Gradient checkpointing [18,19,20,21] trades computation for memory by dropping some of the activations in the forward pass and recomputing them in the backward pass. Swapping [22,23,24,25,26] utilizes the huge amount of available CPU memory by swapping tensors between CPU and GPU.…”

Section: Memory-efficient Training Systemsmentioning

confidence: 99%

“…We implement highly-optimized Table 1: Usability Comparison of Memory Saving Systems. The systems include Checkmate [19], DTR [21],…”

Section: Activation Compressed Layersmentioning

confidence: 99%

See 1 more Smart Citation

ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training

Chen,

Zheng,

Yao

et al. 2021

Preprint

View full text Add to dashboard Cite

The increasing size of neural network models has been critical for improvements in their accuracy, but device memory is not growing at the same rate. This creates fundamental challenges for training neural networks within limited memory environments. In this work, we propose ActNN, a memory-efficient training framework that stores randomly quantized activations for back propagation. We prove the convergence of ActNN for general network architectures, and we characterize the impact of quantization on the convergence via an exact expression for the gradient variance. Using our theory, we propose novel mixed-precision quantization strategies that exploit the activation's heterogeneity across feature dimensions, samples, and layers. These techniques can be readily applied to existing dynamic graph frameworks, such as PyTorch, simply by substituting the layers. We evaluate ActNN on mainstream computer vision models for classification, detection, and segmentation tasks. On all these tasks, ActNN compresses the activation to 2 bits on average, with negligible accuracy loss. ActNN reduces the memory footprint of the activation by 12×, and it enables training with a 6.6× to 14× larger batch size. We implement ActNN as a PyTorch library at https://github.com/ucbrise/actnn.

show abstract

Dynamic Tensor Rematerialization

Cited by 11 publications

References 0 publications

Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning

Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning

Survey on Large Scale Neural Network Training

ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training

Contact Info

Product

Resources

About