AutoChunk: Automated Activation Chunk for Memory-Efficient Deep Learning Inference
Authors: Xuanlei Zhao
Shenggan Cheng
Guangyang LU
Haotian Zhou
Bin Jia
Yang You
Venue: ICLR 2024 poster (2024)
Keywords: Long Sequence Inference
Inference
Compiler
Activation Memory
Machine Learning Infrastructure
Low Resource Inference
Primary area: infrastructure, software libraries, hardware, etc.

Abstract: Large deep learning models have achieved impressive performance across a range of applications. However, their large memory requirements, including parameter memory and activation memory, have become a significant challenge for their practical serving. While existing methods mainly address parameter memory, the importance of activation memory has been overlooked. Especially for long input sequences, activation memory is expected to experience a significant exponential growth as the length of sequences increases. In this approach, we propose AutoChunk, an automatic and adaptive compiler system that efficiently reduces activation memory for long sequence inference by chunk strategies. The proposed system generates chunk plans by optimizing through multiple stages. In each stage, the chunk search pass explores all possible chunk candidates and the chunk selection pass identifies the optimal one. At runtime, AutoChunk employs code generation to automatically apply chunk strategies. The experiments demonstrate that AutoChunk can reduce over 80% of activation memory while maintaining speed loss within 10%, extend max sequence length by 3.2x to 11.7x, and outperform state-of-the-art methods by a large margin.

Published as a conference paper at ICLR 2024
AUTOCHUNK: AUTOMATED ACTIVATION CHUNK FOR MEMORY-EFFICIENT LONG SEQUENCE INFERENCE
Xuanlei Zhao" Shenggan Cheng', Guangyang Lu², Haotian Zhou², Bin Jia", Yang You National University of Singapore
AHPC-AI Technology Inc.
ABSTRACT
Large deep learning models have achieved impressive performance across range of applications. However, their large memory requirements, including parameter memory and activation memory, have become a significant challenge for their practical serving. While existing methods mainly address parameter memory, the importance of activation memory has been overlooked. Especially for long input sequences, activation memory is expected to experience significant exponen- tial growth as the length of sequences increases. this approach, we propose AutoChunk, automatic and adaptive compiler system that efficiently reduces activation memory for long sequence inference by chunk strategies. The proposed system generates chunk plans by optimizing through multiple stages. In each stage, the chunk search pass explores all possible chunk candidates and the chunk selection pass identifies the optimal one. At runtime, AutoChunk employs code gen- eration to automatically apply chunk strategies. The experiments demonstrate that AutoChunk can reduce over 80% of activation memory while maintaining speed loss within 10%, extend max sequence length by 3.2x to 11.7x, and outperform state-of-te of-the-art methods by large margin.
INTRODUCTION
In recent times, significant progress has been made large deep learning models, with their remarkable capabilities demonstrated across range of domains, including natural language processing (e.g., GPT-3 (Brown et al.] 20201), computer vision (e.g.. ViT Dosovitskiy et al. (021)). multimodal applications (e.g., DALL-E (Ramesh et al. (022)) and protein prediction (e.g., AlphaFold (umpeer et al. 2021)). ). As
GFT AdtuCClaam GIT se Aunt-Chance
70
with
60
Alstoffold AstuCrunnk
Activation Memory (GB)
50-
the scale of models increases, the substantial demand for memory resources emerges as major bottleneck for their application.
30
20
Model memory can be distinguished into parame- ter memory and activation memory. Researchers have made efforts to reduce parameter memory cost through techniques like data movement (Ren (et al. 2202 1 Aminabadi et al. (202)) and parallelism (Shoeybi al., | Rajbhandari et al. 2020). How- ever, relatively little attention has been given to the activation memory inference, which refers the storage of intermediate tensors during the model's computation. However, with the increasing size and complexity of models, activation memory becoming a critical consideration for long input sequences, such as documents, images and expected to experience significant exponential growth as the length of sequences increases, as shown in Figure which makes their inference challenging and costly.
10.
10a
10a
10-
Sequence Longth
Figure 1:
Comparison of activation memory usage with and without AutoChunk.
information. Activation memory is
"This work done during Xuanlei's intemship HPC-AT Technology Inc.

---

Published as a conference paper at ICLR 2024
Merorr Consumplap
Memory Concumption
ACTIVATION
ACTIVATION
OUTPUT INPUT
OUTPUT
INPUT
INPUT
INPUT
INPUT
INPUT
Time
(a) Origin activation memory.
(b) Activation memory with chunk.
Figure Demonstration activation chunk.
One approach to mitigate activation memory is quantization (Han et al.; 2016 Krishnamoorth 2018), but its compression capability is limited and and may introduce errors to the results. Another method is fused kernels like FlashAttention (Dao et al. (2022), yet they are primarily tailored for specific modules such Attention (Vaswani al.) (017). Recent studies (Jumper et al. [022]] Liu et al. 2022 Kitaev et al. 2020) have proposed chunk as a solution for activation memory in attention and feed-forward. As depicted in Figure 2 chunk decomposes intermediate activation into multiple segments and computes them sequentially to reduce the peak activation memory usage. represents an effective and promising solution for handling activation memory in more general scenarios. However, there are also several limitations that need to be addressed: 1) Chunk aims to reduce activation memory but may sacrifice computational efficiency, posing speed challenges for real-time serving tasks. 2) Many chunk settings are very speed-sensitive and require careful optimization. 3) Manual chunk design for different models and various input sequences within model is often suboptimal, falling significantly short of the desired level of efficiency, and is so labor-intensive that it becomes nearly impractical even for experts.
To address the challenges outlined above, propose AutoChunk, automatic and adaptive compiler system that efficiently reduces activation memory for long sequence inference. In compiler passes, AutoChunk generates chunk execution plans by optimizing the plan with dynamic programming through multiple stacks of chunk search and chunk selection. It first identifies all potential chunks for the given model based its analysis on computation graph and memory status, providing flexibility for various chunk settings. Then it employs the chunk selection method to determine the best chunking strategy for the given scenario, striking a balance between speed and memory usage. In runtime, we utilize code generation to apply and recompile the chunk plans to the code automatically. Our experiments shows promising results that AutoChunk reduce 80% of activation memory while maintaining speed loss within 10%, and extend max sequence length by 3.2x to 11.7x. And significantly surpasses sfate-of-the-ard expert-designed chunk strategy and fused kernels.
We summarize our contributions follows:
We introduce an innovative method for enabling automated chunk for long sequence in- ference to reduce activation memory effectively. Our approach allows for searching and generating chunks of any dimension and settings, making it the first of its kind.
We propose a novel metric for evaluating the speed loss of different chunks based on our observation of the uneven distribution of memory cost. Through this metric, we can select the best chunk that reduces memory usage significantly while minimizing speed loss.
We have demonstrated AutoChunk can reduce 80% of activation memory while maintaining speed loss within 10%, and extend max sequence length by 3.2x to 11.7x. AutoChunk surpasses both expert-designed chunk strategies and fused kernels in performance.
PRELIMINARY AND RELATED WORK
2.1 ACTIVATION MEMORY
Activation memory refers to the intermediate tensor memory used during the model's computation in inference. For module represented as Y F(X), there are three parts of activation, which
2

---

Published as a conference paper at ICLR 2024
imputs X, outputs Y and intermediate activation A. As illustrated in Figure 2a the cumulative activation memory can be expressed as:
M = mem (X) mem(Y) + mem(A)
However, in contemporary neural networks, the activation memory for long sequence tend to increase rapidly due to three primary factors: 1) The introduction of more complex modules to enhance performance. 2) The adoption of larger models, which results in increased dimension for every tensor. 3) The necessity to process even longer sequences for addressing complex tasks. As depicted in Figure [ the activation memory demand for models handling long sequences undergoes substantial exponential growth as the sequence length increases, potentially exceeding the parameter memory by several orders of magnitude.
2.2 CHUNK
To mitigate the issue of activation memory in attention and feed-forward during inference, the chunk method (Jumper et al. [202 Liu et al. 2022 Kitaev et al. (2020) has been proposed. To apply chunk F(X). ) the input X is initially partitioned into segments denoted as (2,, .......... Fn,) where n is called chunk size. Subsequently, for each segment 2;, we calculate 22 its corresponding output using the module, resulting y; F(z:). Finally, we concatenate all segments' outputs obtain the final output Y = [VI; : Val. As shown in Figure 2b the activation memory with chunk can be computed as:
module represented
M = mxm(X) + mYm(Y) + mem(A)/n.
(2)
This process effectively reduces the intermediate activation memory in the computation by factor of n since the size of intermediate tensors is directly related the input size. Considering that the intermediate activation constitutes most of the activation memory, the chunk method leads to significant reduction in activation memory requirements.
However, although chunk is simple and effective, its application is still limited for the following reasons: 1) Chunk inherently reduces activation at the cost of computational efficiency. Inadequately designed chunk can result in significant speed degradation, rendering unsuitable for most real tasks. 2) Several critical settings chunk, such as chunk regions, chunk dimensions, and chunk sizes, directly impact its speed performance. Optimal settings for all these parameters are necessary for good speed. 3) The manual crafting of chunk strategies for individual models and the varied input sequences is usually suboptimal, notably failing to achieve the desired efficiency levels. Furthermore, this process very labor-intensive, often surpassing the capabilities of even human experts. These challenges constrain the practical feasibility of applying chunk methods.
2.3 DEEP LEARNING COMPILER
For machine learning compilers such as Tensorflow XLA (Sabne 2020), TorchInductor and TVM (Chen et al. 2018), optimization techniques like operator fusion and loop tiling have been employed to enhance computational speed. However, these methods tend overlook activation memory considerations and lack the capability to effectively optimize memory utilization over long-range operators from a global perspective. And Jain et al. (2020) aims to reduce activation memory in training automatically by checkpointing (Chen et al. 2016), but is not applicable to inference.
3 AUTOCHUNK
This section outlines the system design of AutoChunk. We begin with the definition of chunk and the system's objectives in Section 3.1 In Section 3.2, provide overview of the AutoChunk system. Section our chunk search technique for identifying all potential chunks. In Section 3.4 we present chunk selection approach for selecting the most effective chunks.
3

---

Published as a conference paper at ICLR 2024
Nx
Score 1,50
Compller
Activat on
Computational Graph
Memorr Budge
Noda
Estimation Pass
Search Pass
Selection Pass
Runtime
Function
Function
Chunk Execution
Figure Overview of AutoChunk's compiler passes and runtime architecture.
3.1 PROBLEM FORMULATION
We first need to formulate the definition of chunk. Without loss of generalization, considering neural network Y = F(X). its chunk formulation can be represented as:
y = F X""),
(3)
where inputs X = X2,. outputs Y [Y1 .... Yzl. and is represent the number of inputs and outputs respectively. Chunkable inputs, denoted by X [X]=] "X,", X ..... X are inputs that can be divided into chunks for sequential computation; non-chunkable inputs, denoted by X [Xn+1, X.+2,..., X. j. are inputs that cannot be chunked such leaf nodes and residual; and chunkable outputs, denoted by y [Y," Y ..... Yon), are outputs that are produced in chunks. Here, = indicates that input X,, is split into chunks along dimension cz, and computed sequentially, and the number of chunks is called chunk size. In this approach, the goal of AutoChunk find best chunk regions and their corresponding settings X X y
3.2 OVERVIEW
The aim this approach is to identify the optimal chunk configuration for any given model, with the goal of controlling activation memory within limit while minimizing speed loss. To this end, we design AutoChunk, which is a compiler system that generates chunk execution plans by optimizing the plan through multiple stacks of chunk search and chunk selection. In chunk search, AutoChunk identifies all possible chunk possibilities in the model, with respect to the given memory budget. In chunk selection, AutoChunk tries minimize the speed loss for chunk strategy by choosing the optimal chunk from all candidates. Through this optimization process, AutoChunk generates chunk plans for the entire model that meets our requirements for speed and memory.
To achieve this, AutoChunk implements novel compilation passes as Figure illustrates. To be spe- cific, given memory budget and model's computation graph in the form of intermediate representation (IR), AutoChunk generates chunks, leveraging three distinct passes. The estimation pass estimates the activation memory cost and identifies the peak activation memory node for a given computation graph. Then, in the chunk search pass, AutoChunk employs bottom-up breath-first search algorithm to explore every possible chunk region in the IR, gathering potential chunk strategies. Finally, the chunk selection pass dynamic programming (DP) to select the optimal chunk configuration that reduces memory usage and minimizes speed loss. By iteratively repeating the above steps, AutoChunk generates chunks until the memory budget is met.
Given the generated chunk plan, AutoChunk employs code generation based on Py Torch FX (Paskke et al. 2019) and recompile the computation graph with chunk plans. During the runtime pass, the chunk execution is invoked based on the modified computation graph. User can simply call our wrapped function model autochunk(madd)) memory budget) apply AutoChunk.
4

---

Published as a conference paper at ICLR 2024
3.3 CHUNK SEARCH
In chunk search, AutoChunk utilizes a novel bottom-up breadth-first algorithm to explore the chunk space and identify all possible chunk solutions. Our algorithm can cover strategies of any dimensions and constitute individual solutions into chunk flows for any models, which allows us to efficiently cover the entire chunk space. In the following section, we describe the space of chunk strategies and the design of our algorithm.
Chunk Flow. In order to formulate the definition of legal chunk region, we introduce the concept of chunk flow. Chunk flow refers to the path of chunk dimension where it flows across multiple consecutive nodes in the computational graph, and the dimension the chunk flow passes node is the node's 's chunk dimension. A legal chunk flow's inputs and outputs should be able to be divided into several parts and compute sequentially. And importantly, the outputs value should remain the same after chunk, which means that chunk flow may be broken by certain computation or reshape. Following Equation 3 considering functions denoted as Y = F(X) and = G(Y), a legal chunk flow can be denoted as:
F(IX).... X.) Y2] G(YY Y21) = |Z Z Z.
(4)
where X, Y and Z are divided and computed into e parts from dimension i, j and respectively, each chunk outputs are equal to original outputs, and chunk flow passes through dimension of X, dimension j of Y and dimension of Z.
Chunk Space. AutoChunk defines the chunk space by proposing four key rules based on chunk flow. Given a computational graph, there are numerous possible ways to execute the chunk plan. For example, where to chunk, how long should chunk region be, which dimension to chunk for every node, what the chunk size is, and how many chunks there should be. Therefore, AutoChunk will first identify all legal chunks in the chunk space.
To be specific, given module, we propose the following rules determine whether can be a legal chunk: 1) Basic Chunk Rule: The basic requirement is that the function can be transformed into the form our definition for chunk in Equation where part of inputs and all outputs should be chunkable. 2) Output Alignment Rule: The outputs of the function must remain same after chunk. 3) Flow Traceability Rule: For each output, there should be at least one chunk flow which can trace back up the inputs without any interruption due reshape computation, which makes sure that the chunk can be applied to the whole function. 4) Unique Setting Rule: Every node in the function should be passed by chunk flows of same chunk settings. These four rules can be formulated as:
Rule 1&2: : F(X,, x==)===)=> =Y=
(5)
Rule 3: VY Y Xx2* t e s.t. ChunkFlow(Cyy Caz);
(6)
Rule 4: Vn N, EI ChunkSetting(c.).
(7)
where N refers to nodes set of the graph, ChunkFlout Flow(z, y) refers to chunk flow that passes r and y. and Chunk-etting((z zz) refers to the chunk dimension and size for node x.
Algorithm Design. By adhering to the aforementioned criteria, we design the chunk search algorithm, as presented in Algorithm [T] Chunk search algorithm takes computational graph G, memory budget M and peak activation node np inputs, and returns all possible chunks. It then searches through node pairs representing the chunk region, which refers to chunk start and end nodes and contains the peak node within the region. Once chunk region is selected, the algorithm gathers the outputs of the chunk region and iterates through each dimension of the outputs to determine whether it constitutes a legal chunk.
To validate the chunk dimension, the algorithm employs a bottom-up breadth- first search algorithm. Starting from the chunk dimension of the output, the algorithm traces the chunk flow in the directed acyclic computation graph towards the input of the chunk region. When the search algorithm passes node, it determines the corresponding chunk dimensions for this node based on the chunk flow.
The algorithm then applies the previously mentioned criteria 7 to check whether this new chunk dimension violates any rules. If this node and its chunk dimension passes the check, it will be added
5

---

Published as a conference paper at ICLR 2024
to the chunk flow, and search continues. If the algorithm successfully reaches the inputs of the chunk region without violating any criteria, it constitutes a legal chunk.
Complexity Optimization. As shown
Algorithm 1: AutoChunk's chunk search algorithm
in Algorithm [ the proposed chunk search algorithm possesses a compu- tational complexity of O( nooee It is evident that number of nodes is
Input: Computational graph G, memory budget M and peak activation node "p
Output: Possible Chunks C'
paramount to the algorithm's 's complex- ity. From our conclusion, there is need to traverse the entire graph but its neighboring nodes sufficient. Therefore, we limit the search to a care- fully designed local window with
C
for chunk.region in GeNNodePairs(G np) do // do bottom up search oulpul_nodes f- GecOuuNoods(chhnnseeeeiin |/ check every chunk dim for chunk dim in dim(outpus nodes) do nodes t- outpul nodes
size of k < Nnode. which reduces the
complexity O (Node) Further-
cur _cu_cchuk-- + (
more, It is time-consuming if we opt for a complete graph search for every
while nodes BottomUpBFSSnoodes) do
new new_ _ttu_ccuukk--
chunk region. So we propose a two- stage search method. the first stage,
GetChunk(nodes, cur-chunk)
if new .chunk satisfy Equ. 5 then
we first examine its inputs and out- 12
cur whwnkkupdate(nnne..m. (nw._chuwwk)c chunk)
puts by checking whether any chunk
else
flow exists between them. If true, pro- ceed to the second stage and search all
break
end if
nodes within the chunk region. By this
if reach inputs then
method, our algorithm's complexity now becomes OÇk Nnode), where { refers to filter passing rate.
C.update(cur ( cchunk)
end if end while end for 21 end for
19
Graph Optimization. As stated in Sec- tion 3.3 we define legal chunk regions by chunk flow. However, this approach
20
may not be optimal under certain circumstances, especially when multiple chunks affect each other. For instance, multiple chunk flows may exist within a chunk region, and only some of them require chunk. However, our algorithm is incapable of distinguishing between them. To address this, when detecting multiple dimension flows within a chunk region, we additionally examine whether some irrelevant flows should be moved out.
3.4
CHUNK SELECTION
Chunk selection is aimed to identify the best chunk that meets the memory constraints while minimiz- ing the impact on speed. By using our critical observations as guide, we employ primary rule for chunk selection, and define loss function that accurately measures the speed reduction of any chunk method without having to run actual code and utilize dynamic. We uses dynamic programming (DP) to search for optimal chunk solution globally.
Memory Estimation. AutoChunk prioritizes reducing the activation memory cost of the model during inference. Therefore, for the identified chunks, the first step is verify if they comply with our memory budget. To achieve this, we combine this chunk's 's config with all chunks generated before, and use our chunk memory analysis algorithm built on PyTorch estimate the activation memory usage under such circumstances. To make our estimation accurate, we not only consider the intermediate activation and residual tensors, but also calculate the memory cost due to continuous operation, which may have large effect on the result.
Selection Algorithm. In light of the potential inter-dependencies among chunks, we propose a dynamic programming (DP) algorithm search the optimal chunk strategy. Specifically, given model, we iteratively conduct passes until memory limit is met. Each pass generates a candidate chunk 8t. The goal of the algorithms is to find the best chunk strategy S [$1, ..., st] that can minimize speed loss while satisfying memory cap. To that end, we propose two loss functions from macro and micro perspectives to assess the performance of each chunk.

---

Published as a conference paper at ICLR 2024
OP
Lh:4
Activation Distrbutton
0.8
0.2
Figure 4: Examples of activation memory distribution. X-axis refers to operators the model.
Based on our observation, more than 70% of nodes have an activation memory consumption that is less than 30% of the maximum memory usage, as shown in Figure This allows us to achieve an 70% reduction in activation cost by optimizing just 30% of nodes. Therefore, there no need to chunk an entire module but several consecutive nodes is enough, as the activation memory distribution inside a module likely to be uneven. To achieve this, we can formulate the macro cost function as:
Limcero NNNooe + 8N
where Naode is the number of nodes, Nops denotes the total floating-poin ppinng operations. With this conclusion guide the macro way chunk, there comes chunk in micro perspective. The first insight may be we should choose nodes with least ops and node numbers because the less computation is, the less speed will be affected. But our experiment shows that nodes with higher computation density is less likely to be affected by chunk because they still have relatively more parallelism even the computation decomposed. And the strides of different chunk dimensions significantly affect the chunk efficiency because of different I/O cost for decomposing and combining tensors, so we should give more weight to dimensions with larger strides. We design the micro cost function as:
question that how should choose
Liniero ? N density +
XN sttride
where Ndensity is computation density calculated FLOPs per node, and N is the stride of the selected dimension. In summary, we need to minimize the cost function as follows:
L Timacro + Limiero
(10)
Then we can use this cost function to estimate the performance of every chunk and search the global optimal chunk strategy S with dynamic programming conjunction with beam search:
min
L(si); s.t. peak memory < memory budget.
(11)
4 EVALUATION
This section presents the evaluation of AutoChunk's performance in inference. All experiments carried out on the NVIDIA Tesla A 100 80GB platform with Pytorch. We select GPT (prefill stage), VIT, AlphaFold and UNet (Ronneberger al. 2015)
experimental models. UNet refers the variant used in Stable Diffusion (Rombach et al. (2022), which consists of multiple stacks of ResNet (He al. 2015) and Transformer (Vaswani al. (017) blocks. The hyper parameters of the cost functions in Equations 8 and are automatically tuned.
4.1 END-TO-END PERFORMANCE
In this section, we aim to answer three key questions: 1) To what extent does the reduction in activation memory affect speed for AutoChunk? 2) How beneficial is AutoChunk when the fused
7

---

Published as a conference paper at ICLR 2024
CPT 12-1---13---0000 --2-1)))
GPT 10-1.5 5-4041
GPT 18-445-04661
GF 18-05-20485
0.8
Normalized Throughput (tokens/S)
0.7
1.0
0.9
AlsioPsll 18=13-312
ActiP- 15-1-11001
Aphubtt 1=4,5-322
Alshadd 15655
1.0
0.9
0.8
0.7
VIT-O-9,5-0
1.0
0.9
100%
50%
40%
30%
20%
100%
50%
40%
30%
20%
100%
50%
40%
30%
20%
100%
50%
40%
30%
20%
Activation Memory Ratio
Figure 5: Throughput of AutoChunk for baseline models under various activation memory constraints. The x-axis indicates the ratio of activation memory usage compared the baseline. The y-axis represents throughput normalized with respect to the baseline.
AphaFold
0.4
Strategies
Speed
All strategies
100%
Memery
No computation density
84.5%
No dimension strides
75.2%
No number of nodes
89.2%
0.1
No flops
91.9%
graph optimization
67.3%
192
Sequence Length
The impact of different strategies on speed.
Table 1:
Figure 6: The activation memory of AutoChunk and fused kernels, normalized to the memory of fused kernels only.
attention kernel is already used? 3) Does our automated method outperform the expert-designed chunk strategy?
Performance Against Baseline. We evaluate throughput of AutoChunk for baseline models in Figure 5) When utilizing 40% 50% of the original activation memory, AutoChunk effectively manages to limit throughput loss to within 3%, signifying negligible impact on speed while effectively halving the activation memory cost for all model types. Furthermore, when operating with only 20% of the original activation memory, AutoChunk's 's speed reduction experiences a slight increase but remains below 10% which acceptable for inference. Notably, for larger sequences, are able maintain the original speed or even accelerate the inference process. This validates the effectiveness of AutoChunk in reducing activation memory while maintaining speed.
Performance Against Fused Attention Kernel. In contemporary models, the attention mechanism prevalent, and several fused kernels alleviate the ((n²) memory cost of and reduce the peak activation memory. To validate the effectiveness of AutoChunk when fused kernel is already used, we apply memory-efficient attention (Rabe & Staats, 2022) and evaluate AutoChunk's ability to reduce activation memory further. And we control the speed loss of AutoChunk 5%. As shown in Figure
6 when using fused attention kernels, AutoChunk able to reduce over 70% activation memory further minor loss in speed. This demonstrates that even when the activation memory of attention is reduced, other parts of the model still maintain high activation memory for long sequence, validating the broad applicability of approach.
8

---

Published as a conference paper at ICLR 2024
A pAnTeis WES chunc pfruut of chuR AuttOurt
1.20
Alshaled
30
Chank:
2.30
Troughtst
2.05
20
0.58
250
2536
2326
2043
3072
Figure 7: Minimum memory comparison of Expert-Designed chunk and AutoChunk for Al- phaFold.
Figure 8: Throughput comparison of expert-
designed chunk and AutoChunk for AlphaFold. Throughput is normalized by the former.
Performance Against Expert-Designed Chunk. We compare the throughput of memory cost of origin AlphaFold, AlphaFold with expert-designed chunk from OpenFold (Ahdritz et al.) 2022). and AlphaFold with AutoChunk. We compare the minimum memory usage and throughput under same memory limit for various sequence lengths. When the sequence length exceeds 1024, AlphaFold runs out of memory. In Figure 7 AutoChunk can reduce the minimum activation memory usage by 30.6% 34.4%. In Figure [] we the chunk size 64 for the expert-designed chunk it's an effective configuration and align the memory cost of them. AutoChunk achieves a speedup ranging from 9.2% to 14.6% when compared to expert-designed chunk. The results demonstrate that AutoChunk surpasses expert-designed chunk strategies in terms of both speed and memory efficiency.
4.2 BREAKING THE MEMORY WALL FOR LONG SEQUENCE INFERENCE
The memory wall has consistently posed a significant challenge for applications involving the processing of long sequences like images and documents. This challenge manifests as a barrier that restricts the execution of models more economical hardware, characterized by limited GPU DRAM, and on edge devices, which typically possess even scarcer computational resources. As illustrated in Figure I our research endeavors have yielded substantial reductions in activation memory usage, achieving reductions spanning several orders magnitude. Consequently. for inputs of those encountered in models like GPT, method permits remarkable 11.7-fold extension in the max inference length. For 2D inputs, such as those encountered in AlphaFold, VIT, and UNet, AutoChunk yields average 3.2-fold extension in max inference length. This marked improvement significantly mitigates the memory overhead associated with the execution of long sequence inference tasks.
4.3 ABLATION STUDY
As illustrated in Table T we evaluate the influence of the chunk selection strategy and the graph optimization on system performance. The observed speed metrics represent the average across a range of models and various memory limits. Our findings demonstrate that all considered strategies notably contribute to the system performance, underscoring the sensitivity of speed to these configurations. And graph optimization is effective in reducing redundant computation.
5 CONCLUSION
We present AutoChunk, an automatic and adaptive compiler system designed to significantly reduce activation memory usage for long sequence inference through the utilization of chunk strategies. AutoChunk efficiently reduces the activation memory requirements with minor speed loss, while offering practical solution for deploying models with long sequence more economical hardware even edge devices. By effectively reducing memory consumption, AutoChunk enhances the efficiency and accessibility of deep learning applications for long sequence, expanding their applicability in real-world scenarios. future, AutoChunk can also be adapted training to reduce activation memory along with checkpointing.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
Yang You's research group is being sponsored by NUS startup grant (Presidential Young Professor- ship). Singapore MOE Tier-I grant, ByteDance grant, ARCTIC grant, SMI grant (WBS number: A-8001 104-00-00), Alibaba grant, and Google grant for TPU usage.
REFERENCES
Gustaf Ahdritz, Nazim Bouatta, Sachin Kadyan, Qinghui Xia, William Gerecke, Timothy J. O'Donnell, Daniel Berenberg, Ian Fisk, Niccolò Zanichelli, Bo Zhang, Arkadiusz Nowaczynski, Bei Wang, Marta M. Stepniewska- Datisssaaaiiuuisska... Shang Zhang, Adegoke Ojewole, Murat Efe Guney, Stella Biderman, Andrew M. Watkins, Stephen Ra, Pablo Ribalta Lorenzo, Lucas Nivon, Brian Weitzner, Yih-En Andrew Ban, Peter K. Sorger, Emad Mostaque, Zhao Zhang, Richard Bonneau, and Mohammed AlQuraishi. OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization, November 2022. URL https: / www. biorxiv org/ cwwttonntttt//ccc wttp://1..110/2022.11111 /tt. 510.5 Pages:
Section: New Results.
Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li,
Du Li, Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, and Yuxiong He. DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale, June 2022.
URL htth:httpaxixii/ org/ abs 00032 arXiv:2207.0000 [cs].
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, and et al. Language Models are Few-Shot Learners, July 2020. URL http: //arxiv.. wwwpxxiv.rrr/e/ee argg/bs//200 1cXv:2005514165 [es].
14165
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training Deep Nets with Sublinear Mem- ory Cost, April 2016. URL http: http///axivv./m org/ absp//0ss/11000 wttpb//..rr//b//11....... 06174 urXiv:1604.0611
[cs].
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen
Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, October 2018. URLhtta
04799 arXiv:1802.0479: [cs].
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and Efficient Exact Attention with 10-Awareness, June 2022. URL http://arxxvv.org/ abs 0tss/2225.14133 arXiv:2205.14135 [cs].
Mefiir-Efffiient
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, June 2021. URL
/ttpxv////xiv./// org/ abs/ 2000. 11 1929. arXiv:2010.1 1999 [es].
Song Han, Huizi Mao, and William J. Dally. Deep Compression: Compressing Deep Neural Networks with Pruning. Trained Quantization and Huffman Coding, February 2016. URL http: //arxiv. org. abs wwtpab////111. 00149 arXiv:1510.0014. [cs].
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recog- wttp:///////1....../. 03385 arXiv:1512.0335.
nition, December 2015. URL httlttttp/aaa..ttt [cs].
Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Kurt Keutzer, Ion Stoica, and Joseph E. Gonzalez. Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerial- ization, May 2020. UtL/http://awet wrw.:oo.gbbabao... /ttpa//s///...//... 02653 arXiv:1910.0266: [es, stat].
John Jumper. Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tuunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, and et al. Highly
10

---

Published as a conference paper at ICLR 2024
accurate protein structure prediction with AlphaFold. Nature, 596(7873):583-
August
2021. ISSN 0028-0836, 1476-4687. doi: 011100088441186-02110088888
URL https
7/www. nwwurattuen..mm
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya Reformer: The Efficient Transformer, February 2020. URL http://///aaxi....ro//////20.0..
S0Xiv:2001100445 [cs, stat].
Raghuraman Krishnamoorthi. Quantizing deep convolaao networ/ for efficient inference: A MiXiv:186600844
whitepaper, June 2018. URL http:ppp//arruu.rrrr/////6666.. [cs, stat].
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, and Baining Guo. Swin Transformer V2: Scaling Up Capacity and Resolution, April 2022. [cs].
wtt.:oorggabbbgl/ooeee 09883 arXiv:2111.0988.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach De Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, Miigh-Pefformnne Deep Learning Library, December 2019, http://:://ann [cs, stat].
org/ www.aass1191 /1912. 01703 arXiv:1912.01700
Markus N. Rabe and Charles Staats, Self-attention Does Not Need $O(n Memory, October 2022.
URL /tt/tttp////aar....rrr///////2 05682 arXiv:21 12.05682 [es].
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, May 2020. URL httt:://arxiv.org//
1910. 02054 arXiv:1910.0200 [cs, stat].
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical Text- Conditional Image Generation with CLIP Latents, April 2022. URL httt:::/aaxiv.org org
abs/2204. 06125 arXiv:2204.0612. [cs].
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminahadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. ZeRO-Offload: Democratizing Billion-Scale Model Training, January 2021. URL http:: /tRp:////aaxxi wttp://..rgaaab///////oo
0684C arXiv:2101.0682 [cs].
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Hsser, and Björn Ommer. High- Resolution Image Synthesis with Latent Diffusion Models, April 2022. URL http:///rxii
org/abs 12. 10752 arXiv:2112.1077 [cs].
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation, May 2015. 1rXiv:1505.0455 0.497 [es].
04597
Amit Sabne. Xla Compiling machine learning for peak performance, 2020.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, March 2020. URL /ttL::/tt///axtt http:ppp///aaa.ooor////////...00
wttp:/ooggga///.oeeee. 08053 arXiv:1909.0803
[es].
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention All You Need, December 2017. URL http: aarxiv. wttp://////aaaaaaaaaba 03762 (xXi::7006.0766 [cs].
11


Reviews:
[{"review_id": "H5NEiWjDPB", "reviewer": "Reviewer_Tqhc", "summary": "The paper considers the memory consumption during inference of large deep neural networks on long input sequences. To reduce the activation memory, the paper proposes an adaptive compiler to automatically configure chunking strategies.", "questions": "How does the splitting across points inside a batch work?\nHow exactly does the dynamic programming approach work to solve Equation 11?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper tackles an important issue that is becoming increasingly relevant as model sizes continue to grow.\n\n- The empirical evaluation of the proposed method appears thorough.", "weaknesses": "- The paper employs substantial jargon and undefined terms. For readers who are not deeply familiar with the topic, sections of the paper are somewhat difficult to comprehend. For instance, it is unclear what portion of activation memory is contained within a chunk. \n\n- The ablation study is arguably somewhat restricted in scope.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AdQo0SYmJs", "reviewer": "Reviewer_Ypeq", "summary": "The authors define a formal model for breaking up neural net computations into chunks and sequentially executing them to save on memory footprint. They then formulate it as a search space optimization problem and provide an efficient search algorithm. They show that sequential execution of chunks on a limited set of ops is sufficient to provide good memory use gains while keeping overhead low.", "questions": "- While activation memory is generally correlated with model complexity, chunkable activations seem heavily dependent on the model type. Obviously the models chosen for evaluation in the paper are amenable (which is not a strike against the work---these models are relevant and important). Can the authors give some intuition or generalizations on the classes of neural net architectures that fail to chunk nicely (vs those that do)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "- By formalizing the definition of legal chunk flows and providing a cost function, AutoChunk turns a programmer intuition (\"certain computations can be sequentialized to save space\") into an computational optimization problem. This breaks down of a lot of barriers to entry. Wrapping everything up into a single function call that statically optimizes a compute graph is a testament to just how end-to-end the authors have made their solution.\n- Dimension slicing and tiling normally has a stupidly large optimization space. The authors provide several straightforward and effective means for reducing that space to a tractable size and then show that DP is sufficient to get good results.\n- The observations on the need for sequential chunking across all operators (fig 4) is useful in understanding the intuition behind why overhead can be kept low. This is generally helpful beyond just its applicability to chunking (even if it has been observed before).\n- I appreciated the measure of effectiveness in the presence of other memory optimization (i.e.- fused kernels). Often times, memory optimizations (sparsification, pruning, compression, etc.) partially or fully cannibalize each others benefits when used in conjunction. Good to see these play nicely together.", "weaknesses": "- The paper uses a *lot* of bespoke jargon and sometimes uses terms before they are formally introduced. For reference, the following terms are used with the form \"chunk ___\": flow, region, dimension, size, search, space, setting, selection, formulation, strategy, plan, candidate. If I don't read the word \"chunk\" again for a while, I'll be happier for it.\n- The benefits of AutoChunk vs. expert plans are a bit middling. This is less a weakness of AutoChunk's algorithm and more an observation that the expected benefits from AutoChunk will come from *unchunked* models rather than those already using a chunking strategy. In other words, AutoChunk is more useful when spreading the benefits of chunking to a broader set of models rather than improving on those that already use it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aYM3ZMIpwW", "reviewer": "Reviewer_m8Q2", "summary": "When serving large-scale deep learning models, their memory requirements are one of the major hurdles. Unlike the parameter memory, optimizations for the activation memory have not been much studied. Since the activation memory is variable depending on the context length, it is important to reduce the activation memory pressure for long context inference. In this research, the authors propose AutoChunk, an automatic compiler system that finds an efficient execution plan with low activation memory pressure. Their evaluation results show that AutoChunk can reduce 50~60% activation memory without speed loss, or 80% activation memory while maintaining speed loss within 10%.", "questions": "- Could you explain the reason why AutoChunk can even accelerate inference for AlphaFold (B=4, S=512) and UNet (B=1, S=192) cases?\n- Is batch dimension also considered as a candidate for chunking? If so, should we run the search algorithm for every execution? It might incur runtime overhead.\n- How long does AutoChunk take to search chunking strategy?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper suggests an important problem; optimizing the activation memory because the context length is rapidly increasing.\n- Unlike the existence of DL compilers related to parallel execution, the paper presents a new type of DL compiler.", "weaknesses": "- Little bit unclear how \"activation memory\" is measured. Unlike training, we can reuse memory in inference. For example, the MLP module of the Transformer layer has the following structure (not assuming gated linear).\n  ```\n  Y = UP_PROJ(X)\n  Z = DOWN_PROJ(Y)\n  ```\n  In this case, X and Z can use the same memory region. Did the paper consider such a characteristic? It is confusing because Figure 4 shows the activation memory distribution of each node.\n- More analysis for experiments will be helpful. For example, what is a chunking strategy that AutoChunk finds for the GPT model in Figure 5? For now, it is just a black-box compiler.\n- For the GPT model, if AutoChunk can reduce the activation memory by half, we can allocate more memory for the key-value cache. It will lead to an end-to-end throughput increase. Are there any results about this? The first paragraph of Section 4 says that the prefill stage is assumed for the GPT case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)