Deep Neural Networks (DNNs) have been dramatically successful over the past decade across many domains including computer vision [29], machine translation and language modeling [40], recommendation systems [30], speech [46] and image synthesis [47], and real-time strategy game control [43]. This success has in turn led practitioners to pursue larger, more expressive models. Today, state of the art models in language modeling and translation have 100s of billions of parameters [38] which requires 100s of GB of active working memory for training. For instance, large models such as BigGAN [6] found significant benefits from increasing both model size and training batch size, and Facebook's recent DLRM recommendation system [30] contains orders of magnitude more parameters than conventional networks. Additionally, to reach beyond human-level accuracy these models are expected to grow even larger with possibly 100× more parameters [22]. The large memory footprints of these models limits training to systems with large amounts of DRAM which incur high costs.
As the memory capacity demands of DNN training are growing, new high density memory devices are finally being produced. Specifically, Intel® Optane™ DC Persistent Memory Modules (PMM) [15,25] can now be purchased and are up to 2.1× lower price per capacity than DRAM. These devices are on the main memory bus, allowing applications direct access via load and store instructions and can be used as working memory. Thus, in this paper we ask the question "what are the design tradeoffs of using PMM in training large DNN models, and more specifically, can PMM be used as a DRAM replacement when training for large DNN models?" Figure 1 shows the training performance for three different memory systems: an all PMM system (lowest cost), an all DRAM system (highest cost), and a heterogeneous system (moderate cost). The all PMM bar shows that naively replacing DRAM with PMM results in poor performance (about 5× slowdown) for training large DNN models. The first-touch NUMA [27] bar shows that current system support for heterogeneous memory is lacking, providing only a small benefit over the all PMM case. However, AutoTM provides 3.7× speedup over the PMM case and is within 20% of the all DRAM system. Thus, we find that a small fraction of DRAM reduces the performance gap between PMM and DRAM, but only if we use smart data movement.
Use of heterogeneous memory to reduce DRAM has been studied in the past. Facebook has used SSDs to reduce the DRAM footprint of databases [13]. Bandana [14] uses SSD based persistent memory to store deep learning embedding tables [10] with DRAM as a small software cache. In the context of machine learning, vDNN [36], moDNN [8], and SuperNeurons [44] develope system-specific heuristics to tackle heterogeneous memory management between the GPU and CPU to overcome the low memory capacity of GPUs. Furthermore, future HPC systems will be increasingly heterogeneous with DRAM, PMM, and HBM [35], so we need a solution that is general and automatic.
In this paper we introduce AutoTM-a framework to automatically move DNN training data (tensors) between heterogeneous memory devices. AutoTM enables training models with 100s of billions of parameters and/or with large batch sizes efficiently on a single machine. We exploit the static nature of DNN training computation graphs to develop an Integer Linear Programming (ILP) [37] formulation which takes a profile driven approach to automatically optimize the location and movement of intermediate tensors between DRAM and PMM given a DRAM capacity constraint.
We evaluate the effectiveness of AutoTM on a real system with Optane PMM by implementing our approach in the nGraph compiler [11]. Our experiments show that naive use of PMM is not effective, but intelligent use of PMM and DRAM is required. Furthermore, using initial public pricing information, we evaluate the cost-performance benefits DRAM-PMM based systems. We show that ratios of 8 : 1 or 4 : 1 of PMM to DRAM can be more cost effective than only DRAM or only PMM. We also compare our approach to the existing hardware DRAM cache implemented in current Intel platforms [25] and find AutoTM offers up to 2× performance improvement over hardware-managed caching.
Finally, we demonstrate that AutoTM can be further generalized beyond PMM-DRAM heterogeneity by applying AutoTM to CPU-GPU systems. The approach taken by Au-toTM uses minimal problem specific heuristics and is thus a general approach toward memory management for many different heterogeneous systems.
The paper is organized as follows. In Section 2 we present a quick overview of training deep neural networks and Intel's Optane DC PMM. In Section 3 we will present the details of AutoTM and in Section 4 we will describe implementation details, followed by our evaluation methodology in Section 5, and the main results in Section 6. We will present extensions to AutoTM in Section 7 and conclude with related work and directions for future work.
Deep neural networks (DNNs) are often trained using a backward propagation algorithm [28] and an optimizer such as stochastic gradient descent. Popular deep learning frameworks such as Tensorflow [1] and nGraph [11] implement DNNs as a computation graph where each vertex or node in the computation graph represent some computational kernel. Common kernels include convolutions (CONV), pooling (POOL), matrix multiplication, and recurrent cells such as LSTM or GRU. Each kernel has its own characteristics such as number of inputs, number of outputs, computation time, and computational complexity. Directed edges in the computation graph between kernels denote data or control dependencies between kernels. An edge representing a data dependency is associated with a tensor, which we consider to be a contiguous region of memory with a known size. All operations were performed using AVX512 load and stores.
Copies between DRAM and PMM were done using streaming load and store intrinsics. Figure 2 shows a simple example computation graph with 5 kernels and 3 tensors. Nodes in the graph are compute kernels, each with zero or more inputs and outputs. The inputs and outputs of a kernel are immutable tensors. Each tensor is annotated with its producing kernel, each user of the tensor, and the last user of the tensor. After its last use, a tensor's memory may be freed for future tensors.
We focus on the case where the computation graph describing the training iteration is static. That is, the computation graph contains no data-dependent control behavior and the sizes of all intermediate data is known statically at compile time. While many DNN graphs can be expressed statically, there are some networks that exhibit data-dependent behavior [38]. We leave extending AutoTM to dynamic computation graphs as future work.
3D XPoint is a resistive memory technology developed by Intel [19] that was initially introduced in the form a SSD called Optane SSD. Recently, this technology as been made available in the form of a standard byte-addressable DDR4 DIMM on the CPU memory bus, just like DRAM DIMMs [25] with the new Cascade lake based chipsets and is called Optane DC PMM (different from Optane SSD). Optane DC PMMs are higher capacity than DRAM modules with up to 512 GB per module available today.
There are two operating modes for Optane DC PMM. In 2 Level Mode (2LM or cached) PMM act as system memory with DRAM as a direct mapped cache. This operating mode allows for transparent use of the PMM at the overhead of maintaining a DRAM cache. App Direct Mode allows users manage the PMM directly. The PMM are mounted on a system as direct access file systems. Files on the PMM devices are then memory mapped into an application. When using a direct access aware file system, loads and stores to addresses in this memory mapped file go directly to the underlying PMM. Note that in App Direct mode, the total available memory is the sum of DRAM and PMM while in 2LM only the PMM capacity is counted. In this work, we focus on using the PMM in App Direct mode, and make comparisons between our optimized data movement and 2LM.
Figure 3 shows the read, write, and copy bandwidth of DRAM and PMM on our test system with six interleaved 128 GB PMMs. The read, write, and copy operations were implemented by splitting a region of memory into contiguous blocks and assigning a thread to each chunk. AVX-512 streaming loads and stores were used to implement the copy operation as they provide significantly higher throughput between DRAM and PMM.
From Figure 3, we make the following observations about PMM: bandwidth is significantly lower than DRAM, read bandwidth scales with the number of threads, write bandwidth peaks at a low number of threads and diminishes with a higher number of threads, copy bandwidth from DRAM to PMM scales with the number of threads, copy bandwidth is chiefly limited by PMM write bandwidth, and there is significant read/write asymmetry. These findings agree with the performance evaluation of Optane PMM by other researchers [25].
The read/write asymmetry has implications on the performance of kernels with inputs and outputs in PMM or DRAM. Figure 4 demonstrates the performance impact on a single CONV kernel. We observe that when the input to the CONV kernel is in PMM and the output is in DRAM, the performance of the kernel is comparable to when both input and output are in DRAM. However, in the cases where the output is in PMM, the kernel runs over two times slower. Any system seeking an optimal runtime with a memory constraint must take these relative timings into consideration when making decision on where to assign data. In the next section, we will describe the details of AutoTM and how it manages these performance characteristics.
An overview of the proposed framework is shown in Figure 5. A DNN model is given to nGraph, which optimizes the network DAG according to the selected backend (e.g. CPU, GPU etc.). As part of the compilation process, our system inspects the nGraph DAG data structure to extract (1) the order and types nodes in the graph, (2) the tensors produced and consumed by each node, and (3) the specific kernels chosen by nGraph.
We then perform profiling on every kernel in the computation graph by varying its inputs and outputs between the different memory pools (i.e., DRAM or PMM) and recording the execution time of the kernel in each configuration. Since this step is potentially time consuming and DNNs typically contain many identical kernels, we keep a software cache of profiled kernels. By keeping a profile cache, profiling for a given DNN only needs to be performed once. Profiling and DAG information is then fed into a Memory Optimizer (described in Section 3.1) along with a DRAM capacity constraint, that mutates the nGraph data structure with the tensor assignments and data movement nodes.
A user of this system only needs a nGraph function, which is a collection of "Node" and "Tensor" data structures describing computations and data flow of the compute graph. These functions can be created by using one of the nGraph front ends, or directly using C++. Profiling, optimization, and code generation all happen as part of the nGraph compilation process and is transparent to the user.
In the rest of this section, we first give a high level introduction to the memory optimizer. Then we present the details of the optimizer's underlying ILP formulation.
The goal of the Memory Optimizer is to minimize execution time by optimizing intermediate tensor movement and placement. The inputs to the optimizer are (1) the types of kernels in the computation graph in topological order, (2) the set of all valid tensor input/output locations for each kernel as well as profiled execution time for each configuration, (3) the sizes of all intermediate tensors, as well as their producers, users, and final users, (4) synchronous copy bandwidths between DRAM and PMM, and (5) a DRAM limit. The output of the optimizer describes the location and date movement schedules for all intermediate tensors that will minimize the global execution time of the graph.
Since the Memory Optimizer is implemented as an ILP, we need to model tensor location and movement using integer or binary variables and constraints [32]. For each tensor t, we create a separate network flow graph G t = (V t , E t ) that traces the tensor's location during its lifetime. Examples of such graphs are given in Figure 6a and 6b. The structure of these graphs allows us to customize the semantics of possible tensor locations and movements.
Using this graph structure, we investigate two separate formulations, static and synchronous. The static formulation (Figure 6a) allows no tensor movement between memory pools. A tensor is assigned to either DRAM or PMM and remains there through its lifetime. The synchronous formulation (Figure 6b) allows tensors to be moved between memory pools but blocks program execution to perform this movement. We further generalize the ILP formulation to an asynchronous formulation that allows overlap between computation and data movement in Section 7.
Network flow constraints [16] are placed on each tensor flow graph G t so that flow out of the source vertex is 1, flow into the sink vertex is 1, and flow is conserved for each intermediate node. The solution to this network flow describes the movement of the tensor. For example, the bold path in Figure 6b implies the following schedule for tensor t 1 : (1) created by kernel k 1 in DRAM, (2) remains in DRAM until the execution of kernel k 2 , (3) after k 2 , synchronously moved t 1 into PMM, (4) prefetch t 1 into DRAM right before k 4 , (5) move t 1 out of DRAM after k 4 , (6) tensor t 1 is in PMM for the execution of kernel k 5 , (7) after k 5 , tensor t 1 is no longer needed and can be freed from all memory pools.
We wish to minimize the execution time of the computation graph under a DRAM constraint. In our framework, computation kernels are executed sequentially. Therefore, in the static formulation where there is no tensor movement, the objective function (expected execution time) is
where K is the set of all kernels k in the computation graph and ρ k is the expected execution time for kernel k. Note that ρ k depends on the locations input and output tensor for kernel k. The selection of input and output tensor locations is not trivial because of dependencies between kernels. For example, if tensor t 3 in Figure 2 is assigned to PMM, then kernel k 2 must produce t 3 into PMM and kernels k 3 and k 4 must reference t 3 in PMM, which has a performance impact.
Given the lower performance of PMM relative to DRAM, the cost of moving a tensor from DRAM to PMM may be amortized by a resulting faster kernel execution. In the synchronous formulation, tensor movement that blocks computation graph execution and may only happen between kernel executions. The objective function then becomes where T is the set of all intermediate tensors t in the computation graph and M sync t is the total amount of time spent moving tensor t. Note that a tensor t may be moved multiple times during its lifetime, so M sync t represents the sum of movement times of all individual moves of t.
As noted above, the execution time of a kernel depends on the locations of its input and output tensors. We must also keep track of all live tensors in DRAM to establish a constraint on the amount of DRAM used. Thus, we need machinery to describe for each kernel k whether the input and output tensors of k are in DRAM or PMEM and which tensors are in DRAM during the execution of k.
For each kernel k ∈ K and for each tensor t ∈ T where t is an input or output of k, we introduce a binary variable
In practice, this variable is implemented as t DRAM t ,k = 1 if and only if any of the incoming edges to the DRAM node in the component in the network flow graph G t for k are taken.
To determine tensor liveness, we introduce binary variables
for each kernel k ∈ K and for each tensor t ∈ T where t is an output or output of k. These variables describe whether a tensor is written into DRAM after the execution of a kernel, and if it remains in DRAM until the next time it is used. In practice, this is implemented as t DRAM t ,k = 1 if and only if the outgoing DRAM to DRAM edge is taken from the DRAM node in the component in network flow graph G t for k.
We make these two distinct class of variables to handle the case in the synchronous formulation where a tensor is prefetched from PMM to DRAM as an input to some kernel k and then moved back to PMM immediately after k.
Our main goal here is to establish a constraint on the amount of DRAM used by the computation graph. We must ensure that the sum of sizes of all live tensors in DRAM at any point is less than some limit L DRAM We use the DRAM variables discussed in the previous section. First, define a helper function ref(k, t) = k ′ where k, k ′ ∈ K and t ∈ T with k ′ defined as latest executing kernel earlier or equal to k in the topological order of the computation graph such that there exists DRAM node in G t for kernel k ′ . For example, in Figure 6a
We want to ensure that at the execution time for each kernel k ∈ K, the cumulative size of all live tensors resident in DRAM is with some limit L DRAM . Using the ref function, we add the following constraint for each k ∈ K:
where |t | is the allocation size of tensor t in bytes, IO(k) is the set of input and output tensors for k, and L(k) is the set of all non-input and non-output tensors that are "live" during the execution of k. We assign a separate limit L DRAM,k for each kernel k initialized to L DRAM to address the memory fragmentation issue discussed in Section 4.2
For each kernel k ∈ K, we use an integer variable ρ k for the expected execution time of k given the locations of its input and output tensors. First, we define a configuration c as a valid assignment of each of a kernel's input and output tensors into DRAM or PMM. For example, a kernel with one input and one output tensor may have up to four configurations, consisting of all combinations of its input and output in DRAM or PMM. The definition of ρ k is then
where C(k) is the set of all valid configurations c for kernel k, n k ,c is the profiled execution time of kernel k in configuration c, and d k ,c is a one-hot indicator with d k ,c = 1 if and only if kernel k's input and output tensors are in configuration c.
The movement cost of a tensor t is the size of the tensor |t | divided by bandwidth between memory pools. Since bandwidth may be asymmetric, we measure and apply each separately. For each tensor t ∈ T , the total synchronous movement time
is the sum of the number of taken edges in G t from DRAM to PMM multiplied by the DRAM to PMM bandwidth and the number of taken synchronous edges from PMM to DRAM multiplied by the PMM to DRAM bandwidth.
In our case where tensors are immutable, we may apply an optimization of only producing or moving a tensor into PMM once. Any future movements of this tensor into DRAM references the data that is already stored in PMM. Further movements from DRAM to PMM become no-ops.
In this section, we describe some of the implementation details which are not directly part of the ILP formulation. The memory optimizer itself was implemented in the Julia [5] programming language using the JuMP [12] package for ILP modeling. Gurobi [18] was used as the backend ILP solver. We chose nGraph [11] over other popular machine learning frameworks based on static computation graphs as our backend because it is optimized for the Intel hardware and is relatively easy to modify. However, AutoTM is a general technique that can be integrated into other frameworks with similar underlying semantics. 1
The nGraph compiler is an optimizing graph compiler and runtime developed by Nervana Systems/Intel for deep learning (DL) applications aiming to provide an intermediate representation (IR) between DL frameworks and hardware backends. An nGraph IR is a directed acyclic graph (DAG) of stateless operations nodes, each node with zero or more inputs, outputs, and constant attributes. Inputs and outputs of each node are multidimensional arrays called tensors with 1 All of the AutoTM code can be found on GitHub at https://github.com/ darchr/AutoTM. an arbitrary layout. The backend kernel used to implement a node is chosen based on the attributes of the node as well as the sizes, data types, and layouts of each of its inputs and outputs. nGraph will also apply generic and backend specific whole graph optimizations such as kernel fusion and algebraic simplification.
Memory location for intermediate tensors is performed using ahead-of-time heap allocation by traversing the function DAG and maintaining a list of live tensors. When tensors are last used, the memory space occupied by those tensors is freed and used for future tensors.
The ILP formulation presented thus far assumes perfect memory management, which means that if the sum of sizes of live tensors is under the memory limit, then all tensors will fit within memory. In practice, this is not always the case. The process of allocating and freeing tensors may fragment memory resulting in a larger memory requirement.
To manage this, we use an iterative process of reducing the DRAM limit for kernels where the the following limit is exceed and rerunning the ILP.
1. We initialize the kernel-wise DRAM limits L DRAM,k to the L DRAM . 2. We solve the ILP using the current values of L DRAM, K .
nGraph translates the resulting schedule and then executes its memory allocator pass. 3. We collect the set of kernels K frag where the total amount of memory allocated exceeds L DRAM due to fragmentation. If this set is empty, we are done. 4. Otherwise, we apply an update L DRAM,k = 0.98L DRAM,k for all k ∈ K frag and go back to step (2). Thus, the ILP solver may have to run multiple times before a valid solution is found. In practice, this process is usually only done 1 to 2 times with a maximum of 5 as discussed in Section 6.6.
Synchronous movement operations are integrated as new move nodes in the nGraph compiler, which are automatically inserted into the nGraph computation graph following memory optimization. The implementation of these move nodes uses a multithreaded memory copy with AVX-512 streaming load and store intrinsics followed by a fence.
Operation scheduling in nGraph consists of a simple topological sort of the nodes in the computation graph, beginning with the input parameters. This creates unnecessary memory usage with move nodes as they are scheduled ad hoc, resulting in tensor lifetimes that are longer than necessary. Thus, we extended the nGraph scheduler so that if a tensor is moved from DRAM to PMM after some kernel k, we ensure that this movement occurs immediately after the execution of k. Conversely, if a tensor is moved from PMM to DRAM Session 10A: Tensor computation and data orchestration -Playing musical chairs! ASPLOS '20, March 16-20, 2020, Lausanne, Switzerland to be used for kernel k, we ensure this occurs immediately before the execution of k.
Our experimental Optane DC system was a prototype dual socket Xeon Cascade-Lake server. Each socket had 6 × 32 GB of DRAM and 6 × 128 GB Intel Optane DIMMs. Each CPU had 24 hyperthreaded physical cores. In total, the system had 384 GB of DRAM and 1.5 TB NVDIMM storage.
NUMA policy was set to local by default. Unless specified otherwise, all experiments were conducted on a single socket with one thread per physical core. Each workload was run until execution time per iteration (traversal of the computation graph) was constant. Since these workloads contain no data dependent behavior, performance will be constant after the first couple of iterations. Checks were used to ensure no IEEE NaN or subnormal numbers occurred, which can have a significant impact on timing [3].
Our approach does not change the underlying computations performed during training; it is a transparent backend implementation optimization. Thus, the performance of our benchmarks across a few training iterations is sufficient to obtain performance metrics.
We chose to evaluate AutoTM with a multicore CPU platform because Optane PMMs are only available for CPU platforms. However, the ILP formulation of AutoTM should apply to any heterogeneous memory system. We explore one other example with CPU and GPU DRAM in Section 7.
We choose a selection of state of the art Deep Neural Networks for benchmarking our approach. A summary of the benchmarks and batch sizes used is given in Table 1. Conventional CNNs for the Optane DC system were Inception v4 [41], Resnet 200 [21], DenseNet 264 [24], and Vgg19 [39]. All but Vgg19 have complex dataflow patterns to stress test AutoTM. The batch sizes were chosen to provide a memory footprint of over 100 GB for each workload. These batch sizes, while larger than what is typically used, mimic future large networks while still fitting within the DRAM of a single CPU socket of our test system.
We also compare our approach against the native 2LM mode, which is a hardware solution to data management that uses PMM transparently with CPU DRAM as a cache. Since we can not change the physical amount of DRAM used by 2LM, we used very large neural networks that exceed the CPU DRAM and require the use of PMM to train. These very large networks include Vgg416 [36] (constructed by adding 20 additional convolution layers to each convolution block in Vgg16) and Inception v4 with a batch size of 6144.
We want to determine whether PMM is cost effective for training DNNs, and how AutoTM compares against existing solutions to use PMM.
For the conventional benchmarks, we consider the impact of performance with different ratios between PMM and DRAM. These ratios are given in the form a : b where a is the amount of PMM relative to b the amount of DRAM used to train the network. A ratio of 1 : 1 indicates that a network was trained with half PMM and half DRAM. For a network requiring 128GB total to train would have a split of 64 GB PMM and 64 GB DRAM. Setting a ratio such as this may lead to a larger total memory footprint in total due to memory fragmentation in both PMM and DRAM. However, in practice the total memory footprint expansion is minimal with an observed maximum observed value of 3.83% occuring in the static formulation for Inception v4. A ratio of 0 : 1 denotes a system where only DRAM is used while 1 : 0 is a system using only PMM.
We use a baseline of a first-touch NUMA allocation policy with DRAM as a near node and PMM as a far node for the conventional benchmarks. The NUMA policy was encoded in our framework by assigning intermediate tensors to DRAM as they are created until the modeled memory capacity of DRAM is reached. Future tensors can only reside in DRAM if existing tensors are freed.
For the large benchmarks, we compare our approach to the 2LM hardware managed DRAM cache to determine the effectiveness of AutoTM relative to an existing approach.
Figure 7 shows the speedup provided by our scheduling for over training solely with PMM (ratio 0 : 1). The horizontal axis is the ratio of PMM to DRAM used to train the network. We observe that when PMM is used as a direct substitute for DRAM, performance is poor with a 3x to 8x increase in training time (red horizontal line). However, with a minimal amount of DRAM such as an 8 : 1 PMM to DRAM ratio, we are able to dramatically improve performance without changing the overall memory footprint of the application. Further, adding more DRAM only marginally increases performance. This above performance gain does not occur using conventional first-touch NUMA. This is because first-touch NUMA [27] works by allocating tensors into DRAM as they are used by the computation graph until the DRAM capacity is reached. In the training of DNNs, tensors produced early in the forwards pass are used during the backwards pass and thus must be live for the majority of the graph's computation [36]. With first-touch NUMA, these long lived tensors are assigned to DRAM forcing future short-lived tensors into PMM.
AutoTM, on the other hand, is aware of the performance implication of these long lived tensors. The general strategy AutoTM takes is to prioritize short-lived tensors for DRAM placement (Section 6.4). These short-lived tensors mainly include intermediate tensors generated during the backwards pass. By prioritizing short-lived tensors, AutoTM ensures that more tensors overall may reside in DRAM.
Vgg is an outlier due to its extremely large second convolution layer. With small DRAM sizes, some or all of the input and output tensors of this large layer must be placed in PMM, incurring a performance penalty. Once these tensors can be placed in DRAM, we see a significant performance improvement as can be seen in the performance jump from the 4 : 1 ratio to the 1 : 1 ratio. Another interesting feature of this network is that the synchronous formulation performs slightly worse than the static formulation for an 8 : 1 ratio. This is caused by the interaction between the insertion of move nodes and the defragmentation procedure. The regions where the bars are higher than the dollar signs are regions where price-performance is lower.
We use very large networks to compare AutoTM to a hardwarecontrolled DRAM cache (2LM mode). The results from the large benchmarks Vgg416 and the large batchsize Inception v4 are shown in Figure 8. For our analysis, we use the price of the cheapest PMM at $7.85 per GB and the cheapest DRAM at $16.61 per GB. This means the cost-per-GB advantage of PMM over DRAM is about 2.1x. In Figure 9 we only include the cost of the memory actually used. Since Optane DC is a new technology, prices are still adjusting, and as the technology matures, price will likely decrease, improving its cost-effectiveness.
Figure 9 shows the relative performance of AutoTM for our workloads (bars, left axis) as well as the cost of memory used by the application relative to the case where all DRAM is used (dollars, right axis). The use of PMM can be cost effective if the performace lost by replacing some DRAM with PMM is less than the cost reduction. We observe that only using PMM directly is not cost effective, the performance loss caused by the slower devices is not offset by the lower price. However, for PMM to DRAM ratios of 4 : 1 and 1 : 1, AutoTM can provide a cost-performance benefit. This cost-performance benefit may be reduced when taking the whole system into account, but the cost of memory is usually the dominant cost in large systems.
In this section, we present some insight to how and why Au-toTM works using Figure 10. Figure 10a shows the slowdown of the static and synchronous relative to all DRAM. With a small amount of DRAM, performance improves rapidly. This trend continues until a critical threshold where adding DRAM yields diminishing returns.
To understand this behavior, we look at the input and output memory locations for each kernel as well as the amount of data moved. Figure 10b shows the percent by memory footprint of kernel input and output tensors in DRAM. We see a trend to assign as many kernel inputs and outputs into DRAM, with a slight priority on output tensors. This is consistent with the lower write bandwidth of PMM. Furthermore, the point where almost 100% of output/input tensors are in DRAM corresponds to the critical point in the performance graphs. This implies a general strategy to maximize kernel read and write memory accesses in DRAM, followed by data movement to PMM when DRAM capacity constrained.
This idea is reinforced by Figure 10c, which shows the total amount of memory moved between DRAM and PMM in the synchronous formulation. With a DRAM limit near zero, no data movement occurs since no data may be moved into DRAM. A small DRAM allowance, however, is followed by a dramatic increase in data movement, again with an emphasis on moving data from DRAM to PMM. Once the DRAM limit allows almost all tensor inputs/outputs to reside in DRAM, the amount of data movement decreases. The region of gradual slowdown seen in the performance plot is caused primarily by data movement rather than kernel slowdown from more memory accesses to PMM.
To evaluate the accuracy of our profile based approach, we show the error between the expected runtime and the measured runtime in Figure 11. The worst case error occurs for in the static formulations for DenseNet 264 (19%). This error is likely due to CPU caching. During profiling, move nodes are placed at the inputs of kernels under test to allow the inputs and outputs of the kernel to be varied between DRAM and PMM. Kernels cannot be directly profiled due to levels of indirection used in nGraph. Because move nodes are implemented using streaming instructions, no data is resident in CPU caches following these instructions. Hence, our profiling step is essentially measuring the cold-performance of these kernels. This results in an overestimation in run time for the static formulation since no move nodes are used. Vgg19 is less affected due to its very large intermediate layers.
The expected runtime for the synchronous formulation closely follow the predicted runtime because of the use of move nodes placed in the computation graph. The error in the 1 : 0 all PMM case exists for similar reasons.
It is important that the memory optimizer is able to run in a reasonable amount of time. Although ILP is inherently N Phard, recent solvers can find solutions to many problems quickly. Table 3 shows the total amount of time optimizing the ILP. The number of retries due to memory fragmentation is shown in parentheses. Solution time increases with model complexity. Since the optimized computation graph will run for days or weeks to fully train the DNN, this optimization overhead will be amortized. The worst case is the static formulation for DenseNet which takes a little less than an hour to fully solve.
In this section, we discuss two extensions to AutoTM: allowing asynchronous data movement and performing kernel implementation selection. We explain why these extensions were not included in the original formulation and demonstrate their viability on a CPU-GPU platform. These extensions and the GPU implementation of AutoTM show that it is a general and flexible framework for managing heterogeneous memory.
The first extension we investigate is asynchronous offloading and prefetching of intermediate tensors between memory pools. This allows data movement to be overlapped with computation, improving the throughput of the application as a whole. We implemented asynchronous data movement on the PMM system, but found it performed poorly on existing CPU only systems for a number of reasons. Neither a dedicated copy thread nor DMA provided sufficient performance to mitigate the overhead of these approaches. However, a PCIe connected GPU offers a high speed asynchronous data copy API, which is ideal for implementing this extension.
The second extension to the formulation is performing kernel implementation selection. The underlying library used by nGraph to perform forward and backward convolutions for the GPU backend is cuDNN [9], a deep learning library from Nvidia. This library exposes several different implementations for each convolution, each with performance and memory footprint tradeoffs. Generally, faster implementations require more memory. In a memory starved case, this larger memory footprint may require more offloading of previous tensors, resulting in a global slowdown. Since nGraph does not expose any kernel selection options for the CPU backend, we implement this on the GPU instead.
Since AutoTM is implemented using an ILP formulation, we can extend it to be aware of the performance and memory footprint of these different kernels and globally optimize tensor movement and implementation selection. Here, we provide a high level overview of the additions to the ILP formulation to express asynchronous data movement and kernel implementation selection.
In our formulation, we allow an arbitrary number of tensors to be moved between GPU and CPU DRAM concurrently
where ASYNC(k) = {t ∈ T : t can be move concurrently with k} and M async t ,k is the amount of time (if any) spent moving tensor t during the execution of k. The max operation is implemented using standard ILP techniques.
We must extend the tensor flow graphs G t to encode points of asynchronous tensor movement. We identify kernels that can be overlapped with data movement and add a component in each tensor's graph (like those shown in Figure 6b) for each kernel with which the tensor can be moved concurrently.
Asynchronous move times for tensor t must be generated for each kernel k across which t may be moved. This comes directly from the extended tensor graph
where e P →D (e D→P ) is the binary edge variable in E t corresponding to the asynchronous movement of t from PMM to DRAM (DRAM to PMM) across kernel k.
Let I(k) = {1, 2, . . . , n k } be an enumeration of the implementations for kernel k. We generate one-hot binary variables v i,k for all i ∈ I(k) where v i,k = 1 implies implementation i is to be used for kernel k.
Constraining DRAM is similar to the static and synchronous formulations, but now includes kernel memory footprints with
where s k ,i is the memory footprint of implementation i of k.
The expected runtime of a kernel is now dependent on which implementation of the kernel is chosen. Building on the example given in Section 3.5, assume that k has two implementations (i.e. I(k) = {1, 2}). The expected execution time for ρ k kernel k is modeled as
with n k ,c ,i is the profiled runtime of implementation i of kernel k in IO configuration c. This approach does not account for the performance impact of memory conflict between data movement and the computation kernel. However, the maximum memory bandwidth of our GPU is 616 GB/s while the maximum bandwidth of PCIe is 16 GB/s. Thus, the impact of asynchronous data movement is likely low.
We modified the GPU backend of nGraph to support synchronous and asynchronous tensor movement as well as to allow for kernel selection of forward and backward convolution kernels. All GPU kernels are profiled with inputs and outputs in GPU memory. When implementation selection is available, all possible implementations of a kernel are profiled as well. Asynchronous movement was implemented using two CUDA [31] streams: one for computation and the other for data movement via cudaMemcpyAsync. These streams are synchronized before and after an asynchronous movement/computation overlap ensure data integrity.
Our system used a Nvidia RTX 2080 Ti with 11 GB of GDDR6 using CUDA 10.1 and cuDNN 7.6. The host system was an Intel Core i9-9900X with 64 GB of DDR4 DRAM. We use the same convolutional neural networks used earlier. The networks and batch sizes used are given in Table 1. We compare the results of AutoTM with the performance of cudaMallocManaged, which is a memory virtualization layer offered by Nvidia for automatically data from the CPU the GPU the event of a page fault and moving unused pages from GPU DRAM to CPU DRAM.
The results for the GPU experiments are given in Figure 12. For networks that fit on the GPU, our approach has no overhead as the ILP optimizer realizes no data movement is needed. As the intermediate working set increases, we observe a several fold improvement with AutoTM over cud-aMallocManaged due to the lack of runtime overhead of our approach and its algorithm awareness. AutoTM provides considerable speedup when data movement between the CPU and GPU is required. The asynchronous extension outperforms the synchronous formulation with its ability to overlap data movement and computation. However, the asynchronous extension is limited to overlapping tensor movement with a single kernel at a time. Since the RTX 2080 Ti executes kernels faster than data movement, time must be spent to synchronize the two CUDA streams.
The synchronization overhead of overlapping tensor movement with a single kernel can be seen by comparing the achieved performance with the theoretical best performance, calculated by assuming infinite GPU DRAM capacity and using the fastest possible implementations for all kernels. As the memory requirement for training increases, AutoTM achieves a lower fraction of this best performance due to synchronization. We did not compare our results directly against vDNN [36] for two reasons. First, the RTX 2080 Ti GPU is much faster than the Titan X used in that work and thus we cannot compare results directly. Second, the code for vDNN is not available, making direct testing on our GPU difficult. However, while vDNN leverages the same characteristics as AutoTM (communication overlapping, kernel selection, and liveness analysis), AutoTM uses mathematical optimization rather than heuristics providing a more general solution.
As an emerging technology Intel Optane DC has been explored in several recent works. These include in depth performance analysis [25], large graph analytics [15], and database I/O primitives [42]. Research into using Optane PMM for virtual machines demonstrates that only a small amount of DRAM is needed [23]. Flash based SSDs have also been used to reduce the DRAM footprint in database [13] and ML [14] workloads. These approaches use a software managed DRAM cache to mitigate the slow performance and block level read/write granularity of NVM SSDs. Operating system support for managing heterogeneous memory [2,45] and support for transparent unified memory between GPU and CPU [26,33] have been studied extensively in the past. However, to the best of our knowledge, the proposed work is the first to explore the design space and cost-performance tradeoffs of large scale DNN training on systems with DRAM and PMM.
Previous works such as vDNN [36] exploit heterogeneous memory between GPUs and CPUs by recognizing that the structure of DNN training computation graphs has a pattern where intermediate tensors produced by early layers are not consumed until much later in the graph execution. The authors of vDNN exploit this to develop heuristics for moving these tensors between GPU and CPU DRAM during training to free GPU memory. SuperNeurons [44] and moDNN [8] build on vDNN. SuperNeurons introduces a runtime manager for offloading and prefetching tensors between GPU and CPU memory as well as a cost-aware method of applying recomputation of forward pass layers during the backward pass to reduce memory. Similar to our approach, moDNN allows tensors to be offloaded and uses profiling information of kernel runtime and expected transfer time to determine how it will overlap computation and communication. AutoTM differs from these previous approaches in that we use mathematical optimization rather than problem specific heuristics. AutoTM also generalizes the location of data across DRAM and PMM instead of requiring data to be in DRAM for computation.
Integer Linear Programming and profile guided optimization have been used widely to address similar problems in research literature. For example, work in the embedded system space [4] uses ILP in to optimize the allocation of heap and stack data between fast SRAM and slow DRAM. ILP has also been used in register allocation [17] and automatic program parallelization [20]. ILP has been used to optimize instruction set customization and spatial architecture scheduling [32]. Profile guided optimization has been used for dynamic binary parallelization [48], process placement on SMP clusters [7] and online autotuning of CPU and GPU algorithm selection [34]. AutoTM builds on these ideas to address the new problem of data movement in heterogeneous memory systems.
We present AutoTM, an ILP formulation for modeling and optimizing data location and movement in static computation graphs such as those used for training and inference of DNNs. AutoTM uses profile data to optimally assign kernel inputs and outputs into different memory pools and schedule data movement between the two pools to minimize execution time under a memory constraint. With AutoTM, we can obtain 2x performance improvement over hardware DRAM caching solutions. We further find Intel Optane PMM can reduce the DRAM footprint of DNN training by 50 to 80% without significant loss in performance. Given the lower cost of Optane PMM, this can yield a cost-performance benefit in systems with mixed DRAM and PMM over a system with only DRAM.
AutoTM uses minimal problem specific heuristics, making it generally applicable to different systems and networks. We demonstrate this flexibility by extending AutoTM to GPUs, and believe it can be further extended to further heterogeneous systems, such as those with multiple GPUs or multilevel systems with HBM, DRAM, and PMM.
This work is supported in part by the Intel corporation and by the National Science Foundation under Grant No. CNS-1850566.
https://www.lenovo.com/us/en/p/7X05A01TNA/customize?dcscGuid= f3dd16d0-96dd-4deb-9c48-9c6cec9578ba (accessedAugust 14, 2019)