TURNIP: A “Nondeterministic” GPU Runtime with CPU RAM Offload

## Abstract

Abstract An obvious way to alleviate memory difficulties in GPU-based AI computing is via CPU offload , where data are moved between GPU and CPU RAM. While CPU offload is useful, it can greatly slow down a computation due to the relatively slow transfer rate between CPU RAM and GPU RAM. To address this, overlapping memory transfer and compute is a necessity, but this asynchronicity introduces nondeterminacy, and hence it is impossible to know beforehand what is the best order of operations. We describe Turnip , which is a system for running AI computations using CPU offload, designed to handle this nondeterminacy. The key innovation in Turnip is the compilation of the AI computation into a dependency graph that gives the Turnip runtime freedom to run operations such as GPU kernel calls in many different orders; at runtime, Turnip chooses the best order on-the-fly, in response to real-time events. We find that Turnip outperforms standard, PyTorch-based systems supporting constrained GPU RAM by a significant margin, and also avoids out-of-memory errors in a severely memory constrained environment.

## 1 Introduction

Memory management for modern AI computations is difficult. In the LLaMA large language model (Touvron et al., ), for example, the attention computation on a sequence of length $n$ produces an intermediate result with $128\times n^{2}$ floating point numbers. Thus, for a long input sequence of 100,000 tokens, the attention computation will result in 1.2 trillion numbers, which would require 2.4 terabytes to store at half-precision. It is for these reasons that out-of-memory errors plague AI programmers ().

CPU offload—where data are moved to CPU RAM for storage—can help. As CPU RAM is much more inexpensive than GPU RAM and it is possible to install many terabytes of CPU RAM in a GPU server essentially “for free”,(^1^11At current prices, a single state-of-the-art H100 GPU costs the same as approximately 10TB of CPU RAM—more than 100$\times$ the RAM available on a H100 GPU) it makes sense to leverage CPU RAM to temporarily store data. This idea has been explored in several systems such as pofo (Beaumont et al., ), AutoTM (Hildebrand et al., ), SwapAdvisor (Huang et al., ), Capuchin (Peng et al., ), and POET (Patil et al., ). These systems view a GPU computation as a dataflow graph, and plan how to fit the computation into GPU RAM by making use of CPU RAM offload.

While CPU offload is an obvious idea, it can greatly slow down a computation, due to the relatively slow transfer rate between CPU RAM and GPU RAM. Thus, any system for CPU offload must ensure that when such a transfer happens, no computation is blocked waiting for the transfer to finish.

In this paper, we propose Turnip (short for “nondeTerministic gpU RuNtime wIth cPu offload”) which is a runtime for multi-GPU servers, designed to systematically support CPU RAM offload. The key innovation of Turnip is its combination of a pre-computed memory access plan called a memgraph with a “nondeterministic,” event-driven system runtime. A memgraph is a dependency graph where vertices represent tasks (such as the execution of a GPU kernel to perform a small
part of attention computation in a layer of a large language model) and edges represent data or memory dependencies. Any execution order that respects the dependencies in the memgraph is
valid, and tasks are dispatched at any time that their dependencies have been met and the appropriate resources are free. Thus, two executions of the same memgraph may lead to different sequences of operations being executed on a GPU, or different sequences of tensors being paged to CPU RAM—hence the non-determinacy. However, the dependencies in the memgraph are such that the final output is always *correct*, no matter the execution order. Turnip’s event-driven, fully asynchronous runtime is unique.
Because operations can be dispatched whenever the dependencies are fulfilled and are not constrained to any specific ordering, it lowers the chance that any GPU will be stalled waiting for a memory transfer to complete. If one task cannot run due to an un-met dependency in the memgraph, it is possible that there is another task that *can* run.

The key technical challenge is how to effectively build a memgraph with as few dependencies as possible, to allow the runtime as much freedom as possible to dispatch operations so that it is never blocked, waiting for a memory transfer to complete. Turnip builds a memgraph by simulating an execution of the computation, mapping tensors to GPU memory locations and, when necessary adding edges that represent memory dependencies, as well as offload and reload operations.

## 2 Why Is Non-Determinacy of Execution Order Crucial?

The design of Turnip is based on a simple hypothesis: When running a GPU-based computation that utilizes CPU RAM, asynchronous operations such as offload and reload will have a seemingly nondeterministic running time that is difficult to pre-plan for. The system runtime must accommodate
the resulting non-determinacy, or else performance can suffer.

Figure: Figure 1: A simple memgraph.
Refer to caption: extracted/5899448/SimpleMemgraph.png

Consider Figure 1, which depicts a memgraph for a single GPU system with an offload (data movement from GPU RAM to CPU RAM) and a reload (from CPU RAM to GPU RAM). Vertices are operations (for example, GPU kernel calls) that produce data. Black edges indicate data or consumption dependencies; red edges indicate memory dependencies (memgraphs will be described in detail in Section 4). Note the memory dependency from $\texttt{offload}_{1}$ to 4. This exists because the output of 4 will be written to the location of the output of 1, and so 4 cannot execute until the $\texttt{offload}_{1}$ completes. The data dependency from $\texttt{reload}_{3}$ to 5 exists because the kernel associated with vertex 5 will consume the reloaded tensor.

Imagine that a system has executed GPU kernels associated with vertices 1 and 2. It is currently executing the $\texttt{offload}_{1}$ and the $\texttt{reload}_{3}$. At this point, *it is impossible to know which kernel should run next* (4 or 5) as this depends on which memory transfer finishes first. Ideally, *this decision will be made at runtime*. If the system deterministically decides to run 4 before 5 at compile time, and the $\texttt{reload}_{3}$ finishes first, the GPU will sit, idle, waiting for the $\texttt{offload}_{1}$ to complete. This is why a special-purpose, “nondeterministic” runtime is needed: to properly handle the non-determinacy induced via the addition of memory operations.

## 3 Related Work

There are two approaches taken by systems dealing with limited GPU memory. Some, like Turnip, accept an abstracted version of a generic GPU computation. Other systems are more specifically targeted to certain categories of models, optimization algorithms, or to specific tasks such as training or inference. Unlike Turnip, none of these existing systems consider the effect of the nondeterminism of offload and reload operations on system performance, nor do any focus on the system runtime.

Using the first, more general approach, are systems that accept a generic dataflow graph and, like Turnip, plan for execution in limited memory: pofo (Beaumont et al., ), AutoTM (Hildebrand et al., ), SwapAdvisor (Huang et al., ), Checkmate (Jain et al., ), Capuchin (Peng et al., ), and POET (Patil et al., ) all assume an input dataflow graph for a machine learning computation, and then plan for execution in limited memory. Checkmate considers only tensor re-materialization, whereas POET, pofo, and Capuchin consider re-materialization and offload; AutoTM and SwapAdvisor consider only offload.

The more targeted approach is taken by the DeepSpeed project () and the various ZeRO optimizations.
For transformers and other, similar models, DeepSpeed inference (which includes ZeRO-Inference) (Aminabadi et al., ) has two key ideas. First, DeepSpeed inference “offload[s] some activation from
GPU to CPU memory while not in use.” Second, DeepSpeed inference “pins the
model weights either in DRAM (if large enough) or NVMe,
and streams each layer into GPU memory for computation
when needed.” FlexGen (Sheng et al., ) seeks to use a variety of methods to speed transformer inference given limited hardware, including model weight offload to CPU, quantization (Yao et al., ; Frantar et al., ), and sparse attention (Child et al., ). The latter two ideas are orthogonal to the ideas in this paper. For CPU offload, FlexGen optimizes a “zig-zag” block scheduling that works through a transformer layers and sequences in the batch, offloading and reloading the KV-cache (Pope et al., ) and model weights. PagedAttention (Kwon et al., ) deals with low memory utilization in transformers, developing a paging system for the KV-cache.

Figure: Figure 2: A decomposition of matrix multiplication.
Refer to caption: extracted/5899448/Figs/Decomp.png

ZeRO-Offload (Ren et al., ) is a comprehensive solution for limited-memory training that can be seen as primarily using CPU RAM for running the ADAM optimizer, moving weights to GPU RAM on a carefully-controlled schedule. ZeRO-Offload is an enhancement on ZeRO (Rajbhandari et al., ), which is designed to be memory-efficient, partitioning both the optimizer and the data across multiple GPUs.
ZeRO-Infinity (Rajbhandari et al., ) is similar, and includes a CPU offload engine, as well as tiling of operators to utilize the RAM of multiple GPUs.

## 4 taskgraph s and memgraph s in Turnip

Figure: Figure 3: Example taskgraph consisting of six GPU kernel calls and three GPU-to-GPU transfers.
Refer to caption: extracted/5899448/Figs/TaskGraph.png

Turnip takes is input a taskgraph. A taskgraph is a dataflow graph (a directed, acyclic graph) that describes how to perform multi-GPU computations. In a taskgraph, edges represent data flow, and vertices represent operations over tensors. A vertex without any inputs (called an *input vertex*) is associated with an input tensor.
An operation associated with a non-input vertex may be either a kernel call that is to be executed on a specific GPU, or a GPU-to-GPU data transfer.

Turnip is agnostic as to how the taskgraph is created; it could, for example be created using a framework such as FlexFlow (Jia et al., ) or Alpa (Zheng et al., ).
Consider a matrix multiplication $\textbf{X}\times\textbf{Y}$, and assume we wish to execute this matrix multiplication on three GPUs. To produce a taskgraph, a framework such as FlexFlow may choose to decompose this matrix multiplication as depicted in Figure 2, perhaps corresponding to the taskgraph of Figure 4.

Given such a taskgraph, Turnip first compiles the taskgraph into a memgraph, which it will eventually execute. Like a taskgraph, a memgraph is also a directed acyclic graph. Every vertex in the original taskgraph will be present in a corresponding memgraph. Further, the compilation process may add additional offload and reload operations that move memory from GPU RAM to CPU RAM, and vice versa.
During the compilation process, the output associated with every vertex in the memgraph is mapped to a memory location. Unlike the input taskgraph, the memgraph is not a dataflow graph; it is a dependency graph. If there is an edge from $v_{1}$ to $v_{2}$, it means that $v_{2}$ depends on $v_{1}$ and $v_{2}$ may not execute until after $v_{1}$ has been executed. In a memgraph, there are two types of dependencies. One is a data dependency, which is inherited from the taskgraph (or is created via the addition of an offload or reload; see below). The second is a memory dependency, which is added to ensure that there are no race conditions in the graph. A *race condition* occurs when there is some vertex for which two valid executions of the graph may produce a different output. This can happen when two vertices write to the same memory location, and it is possible for a third vertex to read either output, depending upon the execution order.

Figure: Figure 5: memgraph corresponding to Figure 4.
Refer to caption: extracted/5899448/Figs/DepGraph1.png

Let us illustrate a possible compilation of the taskgraph of Figure 4 to a memgraph.
Imagine that our three GPUs each have five memory locations, and for simplicity, each tensor is the same size and occupies exactly one memory location. During compilation, the tensor associated with the output of each operation in the taskgraph is assigned to a memory location, as depicted in Figure 4. GPU 1 must deal with seven tensors total (two input tensors and five additional tensors that are created via the execution of some operation), and we cannot fit all seven of those tensors in memory, given our five locations. Thus, the tensors output by operations A and 4 are both mapped to GPU1-Loc1, and the tensors output by operations 1 and 8 are both mapped to GPU1-Loc3.

Figure: Figure 6: Possible mapping of tensors to GPU RAM.
Refer to caption: extracted/5899448/Figs/MemoryAllocation2.png

A corresponding memgraph is shown in Figure 5. Note that two new edges representing memory dependencies have been added. These edges guarantee that the graph is free of race conditions.
Specifically, a graph will be free of race conditions, if, whenever the outputs of vertices $v_{1}$ and $v_{2}$ have both been mapped to the same memory location, either $v_{1}$ *safely overwrites the result of* $v_{2}$, or $v_{2}$ *safely overwrites the result of* $v_{1}$.
We say that “$v_{1}$ safely over-writes the result of $v_{2}$” if and only if, for every $v_{3}$ that consumes the output of $v_{2}$, there is a memory dependency from $v_{3}$ (or some descendent of $v_{3}$) to $v_{1}$ (or to some ancestor of $v_{1}$).
Why? If $v_{1}$ is to safely over-write the result of $v_{2}$, we need to ensure that $v_{1}$ cannot execute until all of the consumers of $v_{2}$ have finished execution—such memory dependencies ensure this.

For example, from Figure 4 we see that the output of vertex 4 is mapped to the same location as the output of vertex A. In the associated memgraph of Figure 4, to ensure that 4 safely over-writes the result of A, we add a memory dependency from 3 (the only consumer of A) to 4.
From Figure 4 we also see that the results of 1 and 8 are mapped to the same location. To ensure that 8 safely over-writes the result of 1, we add a memory dependency from 3 (the only consumer of 1) to 8. Note that this memory dependency is shown as a dashed line; this indicates that it is *superfluous*, as there is already a data dependency from 3 to 8, so this memory dependency is not needed for correctness.

Figure: Figure 7: memgraph with less GPU RAM.
Refer to caption: extracted/5899448/Figs/DepGraph2.png

Things can become more intricate if the memory is more constrained. Consider the case where we have only four memory locations on each GPU, and we wish to compile the same taskgraph. One possible mapping of the vertices taskgraph of Figure 4 to memory locations for GPU 1 is shown in Figure 6; the associated taskgraph is shown in Figure 7.
Note in particular the addition of an offload-reload pair. Both the offload and the reload are new operations that are added to the memgraph during compilation, to facilitate execution in memory-constrained scenarios. We can always compile $v_{1}\rightarrow v_{2}$ in a taskgraph to $v_{1}\rightarrow\texttt{offload}_{v_{1}}\rightarrow\texttt{reload}_{v_{1}}
\rightarrow v_{2}$ in a memgraph. After the $\texttt{offload}_{v_{1}}$, the result of $v_{1}$ takes up no GPU memory, but it cannot be used until the $\texttt{reload}_{v_{1}}$, where it is again mapped to a GPU memory location. The reason for the inclusion of the offload-reload pair in this case is that it allows the result of A to be removed from GPU RAM for a time. Thus, vertex 4 can execute and write its result on top of the result of A, which is subsequently reloaded so that vertex 3 can be executed.

In Figure 6 we see that there are four pairs of vertices whose results are mapped to the same GPU memory locations, and so memory dependencies must been added to the memgraph to ensure that there are no race conditions. Consider A and 4, which are both mapped to GPU1-Loc1. To ensure that 4 safely over-writes the result of A, we have a memory dependency from the $\texttt{offload}_{v_{1}}$ (the only consumer of A) to 4. Or, consider $\texttt{reload}_{v_{1}}$ and 2, which are both mapped to GPU1-Loc4. To ensure that the $\texttt{reload}_{v_{1}}$ safely over-writes the result of 2, there is a memory dependency from the only consumer of 2 (vertex 4) to the $\texttt{reload}_{v_{1}}$.

## 5 The Turnip Execution Engine

Once a memgraph has been produced, it is executed by the Turnip engine using a nondeterministic, event-based framework. As soon as a GPU is unused or a tensor is ready to be offloaded to RAM, the Turnip runtime can immediately assign any available work to the GPU or begin the transfer, without regard to the overall state of the computation.
Also note that there are no calls to memory-management routines such as cudaMalloc or cudaFree during memgraph execution, as memory management is no longer dynamic. Tensor placement is pre-determined before execution, and if dependencies are respected, there can be no memory corruption due to race conditions.

To execute the memgraph, Turnip runs a central event processing loop, that repeatedly processes callback functions that are called response to completion of the work associated with a memgraph vertex (completion of a GPU-to-GPU transfer, completion of the GPU kernel, or completion of an offload or reload). When a vertex completes and a callback is invoked, the event loop checks to see if any other vertex can be executed. That is, it searches for a vertex $v_{1}$ where (a) all vertices $v_{2}$ with an edge $v_{2}\rightarrow v_{1}$ in the memgraph have also completed; (b) if $v_{1}$ is a kernel call, then the GPU $v_{1}$ is assigned is currently free. When the event loop finds such a vertex, it launches it, and searches for another such vertex. When it can find no such executable vertex, it goes to sleep until woken by another callback.

## 6 Building a memgraph

The key technical question we address in this paper is: How to construct a memgraph from a taskgraph? The primary requirement for the compilation process is *correctness*. Correctness requires that (a) every data dependency present in the taskgraph is also present in the memgraph, *or* is replaced with a sequence of offload-reload operations;(^2^22So, for example, if $v_{1}\rightarrow v_{2}$ is present in the taskgraph, we may have $v_{1}\rightarrow\texttt{offload}_{v_{1}}\rightarrow\texttt{reload}_{v_{1}} \rightarrow v_{2}$ in the memgraph) (b) there are no race conditions in the memgraph; (c) the memgraph has no cycles. In addition, it is desirable for the memgraph to be performant. A memgraph will not be performant if memory dependencies severely constrain the execution order of vertices. Such constraints may reduce parallelism and GPU utilization.

Figure: Figure 8: Building a memgraph via execution simulation.

Our basic tactic during compilation is to rely on a simulated execution of the taskgraph to generate the memgraph. Given a serial ordering of the vertices in the taskgraph that respects all dependencies (so that if $v_{1}\rightarrow v_{2}$ is in the taskgraph, $v_{1}$ is before $v_{2}$ in the ordering) we simulate its execution, making calls to special variants of malloc and free that do not actually allocate GPU RAM, but instead maintain a map of used and free RAM slots on the GPU that is the target of the compilation. These implementations also maintain a history of which tensors occupied which positions in simulated GPU RAM, to correctly generate memory dependencies. As the simulation runs, the memgraph is constructed. Calls to the special malloc implementations associate memgraph vertex outputs to GPU memory locations (effectively producing the mappings depicted in Figure 4 and Figure 6). Whenever a call to the malloc variant fails because there is not enough GPU RAM, an offload vertex must be added to the memgraph. Whenever it is time to simulate execution of a taskgraph vertex but one of the inputs is not in the simulated GPU RAM, then a memory location for the corresponding reload vertex is allocated, and a data dependency on that reload is added to the memgraph.

As the simulation runs, there are two *horizons*, or counters that mark progress through the serialized taskgraph. The first is the allocHzn. Every vertex in the taskgraph that is older than the allocHzn has had a space allocated for it. The second is the execHzn. Every vertex in the taskgraph that is older than the execHzn has been “run” according to the simulation. To ensure a high-quality memgraph, our compilation algorithm greedily tries to push the allocHzn as far as possible past the execHzn. Intuitively, this will produce fewer constraints in the resulting memgraph. A kernel associated with a vertex cannot run until it has GPU RAM to write its output. If this GPU RAM is available very early in the simulation, then it gives the Turnip event processing loop more freedom to choose a vertex execution order that does not exactly match the simulated ordering, generating more opportunities to run available kernels while waiting for memory transfers.

The overall algorithm, BuildMemgraph, is given above in Figure 8. Note that this variant of the algorithm assumes each tensor takes up exactly one slot in GPU RAM. In the ”real life” case where tensors are variably-sized, the algorithm does not change appreciably—specifically, in the variably-sized case, freeing space for a tensor can evict a variable number of tensors to CPU RAM—but assuming uniformly-sized tensors simplifies the presentation.

At the highest level, the algorithm operates by first checking to see if it can allocate space for the vertex at the current allocation horizon, $v_{\texttt{allocHzn}}$. If it cannot, the algorithm makes sure there is space available for the output of the next vertex to be executed (the only way there is not space is if $\texttt{allocHzn}=\texttt{execHzn}$ and the last allocation failed; this implies it is time to execute $v_{\texttt{execHzn}}$ and we just failed to allocate space for it). If there is space, the simulation “executes” $v_{\texttt{execHzn}}$.

Figure: Figure 9: simMalloc variants used in memgraph construction.

There are four memory management subroutines used by the algorithm: three variants on malloc (simMalloc, simMallocForceReld, and simMallocOffld) and one variant on free called simFree. Like a traditional malloc, simMalloc finds an open slot for the allocation, but it also adds the memory dependencies to the memgraph necessary to ensure that the vertex $v$ that will occupy the slot will safely overwrite the previous occupant of the slot. simMallocForceReld is like simMalloc, but it is used in the case when a vertex must be reloaded because it is going to be used immediately, and hence the allocation cannot fail. simMallocOffld is a variant of simMalloc that cannot fail, as it finds a victim to offload to ensure the success of the allocation for vertex $v$, adding the offload-reload sequence to the memgraph. Crucially, it renames all instances of the victim $v^{\prime}$ in the taskgraph to refer to $\texttt{reload}_{v^{\prime}}$. In this way, all “future” accesses to $v^{\prime}$ will refer, in fact, to its reloaded version. The routine also adds a memory dependency from the $\texttt{offload}_{v^{\prime}}$ to $v$, as we cannot execute $v$ until the $\texttt{offload}_{v^{\prime}}$ has taken place, and freeing GPU RAM for use.

## 7 Correctness of memgraph Construction

By construction, all edges present in the taskgraph are present in the memgraph, as when a vertex is “executed” during the simulation, all of its incoming data dependencies are added to the memgraph.

To ensure correctness with respect to race conditions and the absence of cycles, the algorithm relies on the total ordering of vertices in $V$, the list input into BuildMemgraph. Further, all new offload and reload vertices produced by the compilation process also have a consistent placement in this ordering. Consider the case when simMallocOffld or
simMallocForceReld are called to obtain data necessary to execute vertex $v$, and an offload or a reload is produced. All of those offloads and reloads take place just before $v$, with all of the offloads happening just before the reloads.

This ordering ensures that there can be no race conditions in the output memgraph: if the outputs of $v_{1}$ and $v_{2}$ are mapped to the same memory location and $v_{1}$ comes before $v_{2}$ in the ordering, then the BuildMemgraph algorithm ensures that $v_{2}$ safely overwrites the result of $v_{1}$. Consider the implementation of simMalloc. Whenever a tensor produced by $v$ is mapped to a memory slot previously occupied by the result of $v^{\prime}$, we add edges to ensure that every consumer of $v^{\prime}$ executes before $v$. Also, consider simMallocOffld, where $v^{\prime}$ is offloaded to CPU RAM to accommodate $v$. Here, memory dependencies are added from the $\texttt{offload}_{v^{\prime}}$ to $v$ and
from all consumers of $v^{\prime}$ to $v$ (when those consumers appeared before $v$ in the list $V$). Note that the taskgraph is modified so that all “future” consumers of $v^{\prime}$ will consume $\texttt{reload}_{v^{\prime}}$ rather than $v^{\prime}$, so they cannot induce a race condition over $v$ and $v^{\prime}$.

Also, consider why there can be no cycles in the output memgraph: all edges added to the memgraph point forward in the total ordering. Consider the
edges added by simMalloc. $v$ can only be mapped to the location used by $v^{\prime}$ if $v^{\prime}$ has been previously free’ed. This implies that any vertices using the output of $v^{\prime}$ have already been “executed”, and so come before $v$ in the total ordering. Thus, any edge from a consumer of $v^{\prime}$ to $v$ must point forward in $V$. Also consider the edges added by simMallocOffld. A similar argument holds here, as we explicitly only add edges that point forward in $V$.

## 8 Experiments

Figure: Figure 10: Time for LLaMA first token (prefill) inference, A100 server. “OOM” is out-of-memory.
Refer to caption: extracted/5899448/FigsNew/A100-1.png

Figure: Figure 11: Time for LLaMA first token (prefill) inference, P100 server. “OOM” is out-of-memory.
Refer to caption: extracted/5899448/FigsNew/P100-1.png

Our experiments evaluate the ability of Turnip to deal with Meta AI’s LLaMA large language model (LLM) (Touvron et al., ), with severely constrained memory. LLM training and inference are chosen as representative, challenging computational workloads encountered by modern ML systems, particularly difficult given the large memory footprint. We assessed Turnip’s performance on 7 billion and 65 billion parameter models.

Figure: Figure 12: Single-sequence inference times.
Refer to caption: extracted/5899448/FigsNew/Single-Seq-Inf.png

The system is implemented in C++, with most GPU kernels generated using Nvidia cuTensor.
Experiments were conducted on two machines: (i) an older 4 $\times$ P100 GPU server (16 GB RAM each) and 22, 64GB DDR4 2666MHz CPU RAM modules, for a total of 1.3TB of RAM, and (ii) an Amazon Web Services p4d.24xlarge instance, equipped with eight A100 GPUs (40 GB RAM each) and 1.15TB of RAM. We were particularly interested in seeing the ability of Turnip to operate in a difficult environment with extremely limited GPU RAM, hence the P100 GPUs, with only 64 GB of GPU RAM total on the server. Key quesitons are: Can software help bridge the gap—particularly the lack of GPU RAM—between older and newer hardware? Can Turnip facilitate model training and inference in a situation with limited RAM?

Figure: Figure 13: Comparing Turnip and ZeRO Infinity for LoRA training.
Refer to caption: extracted/5899448/FigsNew/A100-train.png

(1) LLaMA first token inference.
Our experiments target “first token” inference (also known as ”prefill”): How long does it take to produce the first output token, given an input prompt? We focus on prefill as it is exceedingly expensive in terms of the memory required, scaling quadratically with the size of the prompt. On both machines, we run Turnip, ZeRO Inference (Aminabadi et al., ) (using weight partitioning and model weight offload), and FlexGen (Sheng et al., ). Note that these are PyTorch-based systems, whereas Turnip is not. For FlexGen, we use full CPU offload for activations. All testing is done using batched input, as batching is required for FlexGen and ZeRO (as Turnip simply runs a dataflow graph, it is agnostic to batching). For the smallest batch sizes considered, we test prefill input sequence lengths: 1K, 2K, 4K, 8K, and 16K tokens. For larger batches we use 1K, 2K, 4K and 8K. For Turnip, all model weights and computations were performed using 16-bit floating points, though FlexGen uses very low precision arithmetic to save RAM and speed compute. Results for the A100 GPU server are given in Figure 10. Results for the P100 GPU server are Given in Figure 11.

One of the advantages of Turnip is that it executes arbitrary dataflow graphs in limited memory. Thus, as long as a computation is appropriately decomposed to run on multiple GPUs, Turnip can execute it. This means, for example, that Turnip does not need to perform inference over batches of input sequences and supports arbitrary combinations of model and data parallelism (unlike FlexGen and in ZeRO Inference). While batching tends to increase computational efficiency, the RAM used by a large batch means it is not possible to run inference over long sequences in limited memory (batching precludes that all 320GB of GPU RAM on be dedicated to prefill for a single long sequence). To investigate the ability of Turnip to perform inference over a single long sequence, we test sequences sizes of up to 32K tokens, on both GPU servers and on both the 7B and 65B parameter models. Results are shown in Figure 12.

Figure: Figure 14: Observed speedup due to nondeterministic ordering, compared to (partially) deterministic ordering.

(2) LoRA training for LLaMA.
We also experiment with LoRA training (Hu et al., ). We use a LoRA rank of 16, and train LoRA adaptors for the $K$, $V$, $Q$, and feedforward mapping matrices. Here we run Turnip and ZeRO Infinity (Rajbhandari et al., ); ZeRO is executed using all three “stages” (gradient partitioning, model weight partitioning, optimizer state partitioning) as well as CPU offload. Both Turnip and ZeRO use checkpointing during the forward pass to reduce the memory footprint. We measure the time it takes to run the forward and backward pass for one batch, with varying batch sizes and sequence lengths. All Turnip model weights are stored as single precision (32 bits). Results for training using both the P100 and the A100 server at 7B parameters are given in Figure 13. Both systems had a difficult time training the 65B parameter model. Turnip was faster for the one case it was able to run (1K length sequence, batch size eight took 58.5 seconds using Turnip and 72.9 using ZeRO Infinity) but Zero Infinity was more robust to larger batch sizes, where Turnip failed.

(3) The effect of nondeterminism. One of the key hypotheses of this paper is that nondeterminism can help performance, by allowing the system to react to the observed state of the computation. To test this, we implement a (semi-)deterministic version of Turnip and run the same set of experiments using this new version. Specifically, we add dependencies to fix the order of operations at each of the GPUs to be exactly the order that is input into BuildMemgraph. Unfortunately, this does not remove all nondeterminism from the system because of the way that reductions (summations) of tensors are handled by the system; fully removing all nondeterminism from Turnip is not easily possible. Still the experiment may be instructive. Figure 14 shows the speedup obtained by the nondeterministic version compared to the deterministic, over all experiments.

Discussion.
Throughout the first token inference experiments, Turnip typically performed the best in terms of latency, with ZeRO Inference generally much slower but outperforming FlexGen, except for larger sequence sizes. This would seem to validate the highly non-deterministic, dataflow-based approach advocated for in the paper, at least if the goal is low latency.

To be fair, we note that FlexGen is designed for high throughput, as opposed to low latency, and FlexGen utilizes multiple GPUs only via pipelined parallelism. Note that FlexGen does not seem to get any slower when moving from batch size of eight to 16 on the A100 server. This suggests that filling the pipeline leads to substantial latency. Further, pipelined parallelism is more effective with more work in each pipeline stage, due to the high synchronization overhead and the need to try to overlap communication with computation, perhaps explaining FlexGen’s better performance for larger sequences, which are more dense computationally.

ZeRO Inference takes a much different approach, but it uses a highly synchronized form of model parallelism as it traverses the levels in a transformer, also carefully trying to overlap communication and computation, which may simply be more effective when there is more work to at each level. Turnip, on the other hand, is radically different. It does not “understand” the levels in a transformer, does not need to synchronize processing of the various levels, and simply tries to asynchronously process kernels as fast as it can. If it is stuck waiting for communication, it simply tries to do something else.

For training, there were clear advantages of Turnip over ZeRO Infinity, especially for smaller sequence lengths. This was particularly true on the A100 server, where Turnip was often much faster. For batches of sequences of length 1K, Turnip often took less than 50% of the time to process each batch, compared to ZeRO (at a batch size of 8, the time to process 1K sequences was 5.6 seconds for Turnip and 12.5 seconds for ZeRO, for batch size of 32 it was 12.1 seconds for Turnip and 17.8 seconds for ZeRO). The differences in performance were much less pronounced on the P100 server, though there Turnip was still faster. Finally, we note that both systems suffered significant out-of-memory errors during training. Interestingly, Turnip seemed to have more problems with memory on the A100 server, whereas ZeRO Infinity had more problems with memory on the P100 server. We conjecture that some of that could be solved in Turnip with a better input dataflow graph, which cuts the input problem into smaller pieces.

## 9 Conclusions

We present Turnip, a system and runtime for executing memory-intensive computations on GPU servers, supporting CPU offload. The key innovation of Turnip is its reliance on a “nondeterministic” runtime where a dependency graph is used to guide execution of the computation; any order of operations is acceptable, if the dependencies are respected. We argue that this is necessary when CPU RAM offload is used, or else the system will often be stalled, waiting for CPU-to-GPU transfers.

Limitations. As currently implemented, the biggest limitation of Turnip is that the input computation (in the form of a taskgraph) must be static, and known completely beforehand, so that the memgraph can be constructed. This is not of consequence during model training, but can be an issue during any recursive, generative AI computation. This includes LLM inference, where the next token must repeatedly be generated and the $KV$-cache increases in size. There are some naive solutions to this (such as pre-compiling a memgraph for a specific number of token generations in the case of an LLM) but more work will be necessary to create a satisfactory solution to recursive computation.

Another limitation is that while our experiments did show that Turnip has certain performance advantages, it is not possible to be absolutely sure where those advantages come from. Our ablation shows that a fixed execution order slows down Turnip, suggesting that nondeterminism is important. But unlike ZeRO and FlexGen, Turnip was implemented from the ground up and does not rely on Python or PyTorch—that fact alone could account for some of the performance differences.

## Acknowledgments and Disclosure of Funding

This work was supported by an NIH CTSA award No. UL1TR003167, by the NSF under grant numbers 1918651, 2008240, 2131294, and 2212557, by the US DOT Tier-1 UTC CYBER-CARE grant #69A3552348332, and by a gift from the Roblox Corporation.

## Appendix A Turnip Engine Details

Our execution engine consists of a central event loop that “launches” each vertex in the memgraph.
A vertex can be launched when (1) all dependencies have been completed and (2) the required resources are obtained.
When a vertex is launched, the corresponding operation is executed and then a provided callback is called to notify the event loop that the vertex has completed. In turn, the event loop frees up the obtained resources and keeps track of when vertices complete execution so that subsequent vertices can be launched.
In practice, when launched, a vertex will execute one or more asynchronous CUDA operations on CUDA stream and will then call cudaStreamAddCallback.
As such, every vertex requires as a resource a stream, where a single stream can only be used by a single launched vertex at a time.
We use 5 streams per GPU.
offload, Reload and inter-GPU communication vertices will call cudaMemcpyAsync.
For CPU storage, we allocate a single, large contiguous block of memory with cudaHostAlloc
with flags cudaHostAllocPortable and cudaHostAllocWriteCombined.
When executing Offload vertices, we allocate into the CPU storage memory using our custom allocator; when executing Reload vertices, we free from our custom allocator.
All compute vertices are executed using either cuTensor functions or hand-written CUDA kernels.
An example where a hand-written CUDA kernel is beneficial is for executing portions of softmax so that less workspace memory and fewer vertices would be required.
Two additional resources may be required for computing vertices: workspace memory as required for executing multiple cuTensor functions and locks around write-protected memory.
As an example, we would execute a summation of $n$ tensors with $n$ calls to tensor increment sum-into kernels. However, the output memory would be protected by a resource so that only one sum-into can happen at a time. This implementation is designed to support non-determinism.
We use CUDA version 11.8.0 and cuTensor version 2.0.1. All other code is C++.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1) What is the effect of C++ vs Python?
2) Please elaborate on the effects of variable size allocations on the Memgraph generation
3) How does it compare to other existing work-conserving, priority-based algorithms?
soundness: 2
strengths: - The paper is well-written and the examples are helpful
- The problem addressed in this work is interesting and relevant
- Using LLaMA and LoRA in the experiments makes this work applicable
summary: This paper presents TURNIP, a GPU Runtime that incorporates CPU-RAM offload into the execution graph of AI models under memory constraints. Specifically, the authors create a MemGraph and GPU memory mapping by simulating task graphs. Comparing with PyTorch-based systems, they show faster runtimes of LLMs primarily for short sequence lengths.
weaknesses: - The idea of "nondeterministic" is not new. The described scheduler is simply a work-conserving dynamic scheduler.
- The simulation of the task graph to create the MemGraph is omitting any possible knowledge of execution times. Therefore, the resulting memgraph is only dependent on the topological sort and thus likely suboptimal.
- The authors only present memory allocation sizes of 1 unit each and claim that "In the 'real life' case where tensors are variably-sized, the algorithm does not change appreciably". Intuitively, the problem should get significantly harder with variable size allocations. The authors should elaborate on this.
- Section 7 could benefit from a more formal analysis.
- In Section 8: "Note that these are PyTorch-based systems, whereas TURNIP is not." Does that mean the authors are running their system in C++ while the others are running in Python? Does this explain the initial performance benefit for low sequence lengths? They should run a test of a model that has no memory problems. In that test all algorithms should be the same.
- Section 8, Nondeterministic vs deterministic: In the deterministic case precedence constraints are added based on a topological order which has many solutions. Therefore, it is unsurprisingly weaker. The authors should compare to any other scheduling algorithm, such as priority-based algorithms (e.g. priority-based list scheduling).

Minor issues:
- l147 takes is -> takes as
- Pacement of figures in Sec 4 is not ideal. Try placing them where they are discussed.
- Fig. 7: Why is there a memory dependency from 4 to 3. It should at least be dashed.
- Figure 8 -> Algorithm 1, Figure 9 -> Algorithm 2
- Figure 8: Assignment and comparison of v_allocHzn.loc in the first "if" clause is difficult to understand

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: How is the simulation obtained to generate the memgraph?

How does the system ensure that operations are executed in the best order at runtime?

Has the system been open-sourced?
soundness: 2
strengths: Provided detailed explanations of how to generate memgraph
 
Implemented a working system
summary: This paper presents a runtime system for executing AI applications via GPU-to-CPU offloading. The system dynamically determines whether to allocate AI computations to GPU or CPU memory. It begins by transforming a task graph into a memgraph and implementing the memgraph engine through simulations. This approach enables efficient execution of AI computations through best order of actions (offloading or reloading).
weaknesses: The paper lacks detailed discussions on how Turnip differs from existing systems mentioned in the related work section.

It is unclear how the simulation for memgraph is generated.

The simulation-based memgraph implementation raises questions on how the system guarantees optimal action sequencing during runtime.

Has the system been opensourced?

There are language issues here and there:
Page 3: Turnip takes [as] input ...
Three GPUs, each [having] five ...
With seven tensors [in] total ...

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: The reviewer is happy to raise the rating if some/all the experiment issues stated are addressed.

1. Do you already have any plan to address your paper's known limitations? Since you build a customized framework that doesn't involve Python or PyTorch, For the second one, I believe a separate experiment, when no CPU offloading is involved, comparing your framework with Pytorch on top of the existing experiment parameters can somewhat contribute to the understanding of the performance differences.
2. Do you have any comments on the experiments that I think are missing? (In the weakness section, Others, 1-4)
3. Do you have any explanation for the 5th point in the weakness section under the others sub-section?
4. Just out of curiosity, would it be possible to apply TURNIP to generic graphics or data processing? I feel like it's highly possible.
5. Last question that doesn't contribute to the paper's rating: Will you open-source the code once the paper gets accepted in the future?
soundness: 2
strengths: 1. Originality: TURNIP accepts an abstracted version of a generic GPU computation. Other systems are more specifically targeted to certain categories of models, optimization algorithms, or specific tasks such as training or inference. None of the previous work considers the effect of the non-determinism of offload and reload operations on system performance, nor does it focus on the system runtime.
2. Quality: The quality of the design is demonstrated by experiments under various hardware conditions (A100 and P100), using LLAMA 7B & 65B.
3. Clarity: The details of the idea are clearly demonstrated with graphs and explanations in section 4-7.
4. Significance: TURNIP provides a generic solution to democratize AI computing, especially LLMs, which extends the functionality of older GPUs and enables large-scale AI model operations in constrained environments, making it valuable to both the AI and systems communities.
summary: The paper presents TURNIP, a GPU runtime that supports offloading data to CPU RAM to remove GPU memory restrictions in AI computations. TURNIP addresses the challenges of nondeterministic execution introduced by the slow data transfer rate between CPU RAM and GPU RAM, which harms computational efficiency. TURNIP’s main contribution is its use of a dependency graph (MEMGRAPH) that allows for overlapping execution of GPU kernel operations and memory transfers. The system dynamically chooses the best order of execution based on real-time events, thus minimizing idle times caused by memory transfer delays. TURNIP outperforms existing systems like PyTorch-based ones in constrained GPU environments, contributing to efficient memory management in AI workloads.
weaknesses: 2 limitations that the author stated in the conclusion section of the paper, which potentially harms the significance of contribution (1) and soundness (2):
1. The biggest limitation of TURNIP is that the input computation (in the form of a TASKGRAPH) must be static and known completely beforehand so that the MEMGRAPH can be constructed. This is not of consequence during model training, but can be an issue during any recursive, generative AI computation. This includes LLM inference, where the next token must repeatedly be generated and the KV -cache increases in size. There are some naive solutions to this (such as pre-compiling a MEMGRAPH for a specific number of token generations in the case of an LLM) but more work will be necessary to create a satisfactory solution to recursive
computation.
2. Another limitation is that while the experiments showed that TURNIP has certain performance advantages, it is impossible to be sure where those advantages come from. Our ablation shows that a fixed execution order slows TURNIP, suggesting that non-determinism is important. But unlike ZeRO and FlexGen, TURNIP was implemented from the ground up and does not rely on
Python or PyTorch—that fact alone could account for some of the performance differences.

Others:
1. The experiments conducted in this paper use the LLAMA 1st generation model, while the newest one up to date is already LLAMA 3.2. The paper will benefit from using the SOTA LLMs to demonstrate the effectiveness of TURNIP.
2. As nowadays, the sequence length has grown rapidly (e.g., LLAMA 3.1 models have already supported 128K context window), the paper only did experiments on max 16K, so experiments on longer sequence length will be appreciated to demonstrate the effectiveness of TURNIP.
3. more diverse hardware, such as V100 and H100, can be used in the experiments if possible. Nowadays, people rarely use P100, while some still stick with V100. Meanwhile, H100 represents cutting-edge technology.
4. We already have the 405B model in LLAMA 3.1, so seeing how it performs with TURNIP can make the results more convincing.
5. In Figure 10: Time for LLaMA first token (prefill) inference, A100 server, TURNIP seems to show some scalability issues (e.g., 7B inference batch size 8, TURNIP starts to demonstrate the worst run time when the sequence length is 16K, compared to FlexGen and ZeRO Inference). Although the author mentioned that FlexGen uses very low precision arithmetic to save RAM and speed computing, this doesn't explain the issue here.

A minor issue:
1. Figures 8 & 9 are actually not figures but details of algorithms. You should call them Algorithms 1 & 2.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: What is the definition of the MemGraph referenced in Figure 5? Does it represent solely a data dependency graph with memory dependencies, or does it include additional information like control flow dependency?

In Figure 6, if the tensor locations on GPU devices are predetermined at the compilation stage, will different tensor sizes lead to GPU memory fragmentation? 

How does TURNIP compare to TensorRT or other GPU-only solutions for inference? This information will help us understand its overhead. 

ZeRO-Offload and ZeRO-Infinity adopt different strategies for allocating tensors between GPU devices and CPU memory. To provide a more comprehensive comparison, could you include GPU-only and ZeRO-Offload solutions in Figure 13?

Lastly, why was ZeRO-Infinity not configured with an NVMe as secondary storage, as shown in Figure 13? Doing so could potentially extend the input sequence length or support larger model sizes.
soundness: 2
strengths: An interesting approach for incorporating dependency analysis during the compilation phase to optimize offloading and reloading operations. The evaluation was conducted using a state-of-the-art LLM.
summary: This paper presents Turnip, a software framework to support CPU memory offloading for training and inferencing deep neural networks. The key idea is to introduce a MemGraph, generated at the compilation phase, which is then used to maximize computation and memory transfer overlaps. Turnip was evaluated on 7B and 65B LLaMA models and compared against ZeRO-Inference. The experimental results demonstrate that Turnip reduces first-token latency during inference and accelerates LoRA-based model fine-tuning.
weaknesses: A straightforward approach where the positions of tensors reloaded to GPUs are predefined at specific slots through a compilation process based on the data flow graph. This is fine for simple model architectures with a static execution flow but can struggle to deal with dynamic networks like MoEs where tensors/variables are activated based on the current input data?

The evaluation can be improved. For example, I would like to see a scalability evaluation by testing Turnip on various model sizes, as well as models with f16 during inference. 


This may speed up simple model architecture, but is difficult to work with or achieve better performance on complex/dynamic model structures. Second, the experiment design could introduce more workloads such as 13B, 175B, 400B models and more scenarios such as half precision full parameter training and even more solutions such as ZeRO offload, seen Questions.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: While this paper applies a nondeterministic MEMGRAPH approach to AI computing, similar techniques have been extensively explored in compiler research for general-purpose workloads. What are the unique challenges of applying such memory management strategies to AI-specific workloads?

The non-deterministic, event-driven scheduling works best for tasks with loose dependencies. However, heavily dependent tasks are common in AI workloads (e.g., transformer models with layer-by-layer computations). When dependencies are heavy, it may leave GPUs idle because tasks are queued. How does the proposed approach work under such a case? 

The paper does not seem to optimize GPU assignments or balance workloads effectively across GPUs, especially when some GPUs are busy while others are idle. This could lead to bottlenecks and inefficient use of resources in multi-GPU settings.

What is the time complexity of the MEMGRAPH construction process, particularly for very large TASKGRAPHS? How much overhead in this process?
soundness: 3
strengths: ++ The paper targets important problem of addressing GPU utilization
summary: The paper presents TURNIP, a GPU runtime system optimized for AI computations under constrained GPU memory by utilizing CPU RAM offloading. TURNIP addresses the challenge of nondeterministic execution by introducing a dependency graph -- MEMGRAPH, which allows for flexible operation sequencing. The approach dynamically decides the best operation order at runtime, minimizing GPU stalls caused by memory transfers.
weaknesses: -- The proposed approach may perform poorly when handling complex, dependency-heavy workloads

-- It is not clear what unique challenge in nondeterministic graph-based scheduling for AI computing

Meta-review: Summary
The paper proposes a framework to make the inference of large neural networks more memory efficient. The proposed framework, called Turnip, offloads CPU memory by creating MemGraph which is a dependency graph created at compilation time and MemGraph can be used to decide the best operations at runtime which reduces GPU idle times triggered primarily by memory transfers.

Strengths
1. Improving inference resources needed for large models is very important.
2. The method is tested on Llama 7B and 65B models and compared against ZeRO inference. The key result is that the latency for first token generation is reduced. The authors study this on different hardware too.
3. Since inference is improved, the authors also demonstrate speed ups for LoRA based training.

Weaknesses
1. The biggest limitation is the requirement of a static computation graph. This is a strict requirement and not true for most modern LLMs, especially MOE based models. This severely limits the impact of this paper.
2. The novelty of this work is limited as it can be viewed as compiling the model's execution into a memory access plan. While this is effective, it is a widely used technique in systems across computer science. 
3. Limitations of the evaluation setup in terms of model sizes, floating point precisions (fp8)

Justification
The paper received a majority negative review primarily because of its limitation using static compute graphs and limited novelty. The AC agrees with this majority recommendation from the reviewers.

Decision: Reject