Memory Safe Computations with XLA Compiler
Abstract
Software packages like TensorFlow and PyTorch are designed to support linear algebra operations, and their speed and usability determine their success.
However, by prioritising speed, they often neglect memory requirements.
As a consequence, the implementations of memory-intensive algorithms that are convenient in terms of software design can often not be run for large problems due to memory overﬂows.
Memory-efﬁcient solutions require complex programming approaches with signiﬁcant logic outside the computational framework.
This impairs the adoption and use of such algorithms.
To address this, we developed an XLA compiler extension 1 that adjusts the computational data-ﬂow representation of an algorithm according to a user-speciﬁed memory limit.
We show that k-nearest neighbour, sparse Gaussian process regression methods and Transformers can be run on a single device at a much larger scale, where standard implementations would have failed.
Our approach leads to better use of hardware resources.
We believe that further focus on removing memory constraints at a compiler level will widen the range of machine learning methods that can be developed in the future.
Introduction
Progress in science is inextricably linked with advances in scientiﬁc computing, in terms of both software and hardware.
This is particularly noticeable in machine learning through the huge impact of numerical software packages supporting automatic differentiation (Baydin et al., 2018).
Packages such as TensorFlow (Abadi et al., 2016), PyTorch (Paszke et al., 2019) or JAX (Bradbury et al., 2018) greatly accelerated 1) the implementation of gradient-based optimisation procedures by eliminating error-prone manual differentiation, and 2) the execution of code by leveraging modern and heterogeneous hardware (e.g. GPU, TPU or IPU).
A large portion of this impact is attributable to the accessible and user-friendly form that these features were delivered in.
This contributed to the growth in the machine learning community, in terms of methodological researchers, as well as the wider scientiﬁc audience and practitioners.
The aforementioned software frameworks work by chaining together efﬁcient implementations of mathematical operations (known as kernels ).
By providing implementations that are tailored to various types of hardware, a speed-optimised implementation can be obtained.
While speed is certainly important to pursue, many algorithms face a different challenge: hardware memory constraints.
Often, these have a larger impact, as memory constraint violations can lead to the execution terminating before an answer is obtained.
This make-or-break property is particularly noticeable on GPUs, where allocating more memory than is physically available leads to an immediate termination of execution, and larger amounts of physical memory comes at a signiﬁcant cost.
Now that numerical computation frameworks are widely used, they strongly inﬂuence what machine learning algorithms are adopted.
This happens through hard limitations, as well as usability considerations through what is easily implementable.
Currently, the emphasis on optimising runtime causes many algorithms to be severely memory limited, or too cumbersome to implement.
This is particularly noticeable in methods that rely on linear algebra computations, e.g. kernel methods, nearest neighbour methods for geometric deep learning (Bronstein et al., 2017), or transformers.
In this work, we aim to remove these limitations, by developing a tool that optimises code to be more memory efﬁcient, with a particular focus on linear algebra operations.
This optimisation is transparent to the user, and therefore allows many algorithms to be run at scales that were previously impossible, while leaving implementations as simple as before.
This allows a wider range of algorithms to take advantage of “the bitter lesson”— “General methods that leverage computation are ultimately the most effective” (Sutton, 2019)—while making them more accessible to the wider community, as computational frameworks have sought to do all along.
Our method is implemented as an extension to the XLA compiler (Leary and Wang, 2017), which we chose due to its wide use and support for optimising computations speciﬁed in TensorFlow and JAX.
We demonstrate the beneﬁts of our method by scaling algorithms where simple implementations do not scale due to memory bottlenecks, such as k-nearest-neighbours, sparse Gaussian process regression (Titsias, 2009), and Transformer (Vaswani et al., 2017).
With our extensions, these methods scale to far larger problems, without changing a single line of their implementation in Python.
Our
Gaussian process experiment shows that simply scaling up a 13 year old method can outperform much more recent methods, indicating that older methods may be undervalued in recent literature.
Motivation: Memory-Constrained Machine Learning
Since memory overﬂows cause the execution of code to be immediately halted without producing any result, memory constraints form the key obstacle for scaling many machine learning algorithms.
Researchers and ML practitioners often reach this limit, as memory is a scarce resource that comes at a considerable cost, particularly in GPUs.
This is particularly noticeable in algorithms where minibatching is undesirable, or that rely on pairwise distances, like k-Nearest Neighbours (kNN) or Gaussian processes (Rasmussen and Williams, 2006).
Even in modern deep learning, memory constraints cause problems by limiting batch sizes, layer widths, or sizes of attention mechanisms (Vaswani et al., 2017).
In all of these examples, linear algebra operations cause the bottleneck.
For kNN, kernel methods, and transformers the root of the problem is a pairwise matrix needs to be computed between inputs, giving a quadratic memory cost.
Often, more memory efﬁcient implementations can be programmed at the cost of increased software complexity.
This ranges from minor annoyances, for example accumulating minibatch gradients in an outer loop for large-batch training, to complex engineering efforts that have been published as scientiﬁc contributions in their own right, for example in scaling Gaussian processes to > 10datapoints (Gal et al., 2014;
Wang et al., 2019a; Meanti et al., 2020).
Our goal is to provide a tool that ﬁnds memory-efﬁcient ways to execute algorithms, without the need for increasing software complexity.
This will allow scientists and practitioners to access the beneﬁts of scale in existing methods more easily, and without incurring the cost of expensive large-memory hardware.
For the main demonstration of our approach, we will automatically obtain a memoryefﬁcient implementation of sparse Gaussian process regression (Titsias, 2009), which was previously implemented with considerable difﬁculty (Gal et al., 2014).
The increase in scale makes the method competitive in comparisons where it was previously dismissed as not scalable enough (Wang et al., 2019a), showing the value of reducing the barriers to scaling.
Related Work
A popular approach to address memory issues is distributing computation across multiple resources like a group of GPUs or a computer cluster with network protocol connectivity between machines (Buyya, 1999; Dean and Ghemawat, 2008).
More speciﬁcally, sharding allows large tensors to be split up and distributed across multiple devices, which increases the total amount of memory available for an algorithm, but comes at the cost of requiring more hardware resources.
Most computational frameworks 2 (Abadi et al., 2016; Shazeer et al., 2018; Bradbury et al., 2018; Paszke et al., 2019) support sharding, although some manual speciﬁcation of how to distribute computation is required.
This complicates an implementation and requires the user to have a wide engineering skill set.
Automatic sharding tools such as Tofu (Wang et al., 2019b) simplify implementation somewhat, although the speciﬁcation of computations through a custom interface is still required.
While sharding approach does allow scaling of certain implementations, it remains wasteful for algorithms that can be implemented in a more memory-efﬁcient way, but where it is simply cumbersome to do so.
Compilers have been introduced to allow humans to express programs in an elegant way, while generating programs that actually run well on speciﬁc hardware Aho et al.
Our goal of obtaining memory-efﬁcient implementations, while keeping code convenient for humans, is therefore suited to be addressed by adding memory optimisations to a compiler.
Compilers are already being used to optimise computational graphs, notably in JAX, TensorFlow and PyTorch by XLA (Leary and Wang, 2017), TVM (Chen et al., 2018), Glow (Rotem et al., 2018) for PyTorch only.
TVM performs similar optimisations to XLA, but unlike XLA, it is not seamlessly integrated into popular frameworks and requires additional user effort.
The optimisations in XLA mainly focus on increasing code speed, for example through common sub-expression elimination (CSE), dead code elimination (DCE), operations fusion , and other more speciﬁc modiﬁcations.
The main advantage of XLA is that it optimises computations in a way that is completely transparent to the user who speciﬁes the computational graph.
Although XLA and TVM implement low-level memory optimisations, they do not adapt code handling large tensors to satisfy memory constraints.
For the matrix and linear algebra tasks that we consider, KeOps (Feydy et al., 2020; Charlier et al., 2021) currently provides the most efﬁcient memory management.
To achieve any beneﬁts, a user must specify a series of computations using KeOps classes, which form a layer above the PyTorch framework.
KeOps works similarly to a compiler, by ﬁrst building a symbolic representation of the computation, which allows the computation to be broken into memory-efﬁcient sections, that are then run with custom CUDA kernels.
In terms of prior work, KeOps is closest in aim and achievement to ours.
We aim to address three of its limitations.
Firstly, KeOps requires users to reimplement their algorithms using KeOps classes.
While the programming interface is elegant, needing to mix KeOps and other computational frameworks does add complexity.
Secondly, for KeOps to be able to optimise an operation, it has to be reimplemented within KeOps, which signiﬁcantly duplicates effort.
Finally, because of the former drawback, KeOps does not inherit the support for a wide range of hardware from e.g. JAX/TensorFlow.
Memory Efﬁcient Matrix and Linear Algebra Operations in XLA
Compilers are a promising way for improving runtime properties of code, without requiring user intervention and while leaving code elegant.
The speciﬁc matrix and linear algebra optimisations that we consider have not yet been implemented in any of the frameworks discussed above.
They could be implemented in any of TVM, KeOps, or XLA.
We choose to extend XLA over TVM, due to XLA’s better integration with common computational frameworks.
In addition, we choose to extend XLA over KeOps, because it 1) does not require algorithms to be rewritten in a separate framework, 2) can optimise computational graphs in their entirety, rather than just what is implemented in the separate framework, and 3) can take advantage of the full capabilities that already exist in JAX/TensorFlow.
We introduce several optimisation strategies (known as optimisation passes in the XLA codebase) into the XLA optimisation pipeline.
We aim to constrain the program’s memory footprint with minimal sacriﬁces in the execution speed.
The optimisation passes examine the entire computational data-ﬂow graph (High Level Optimiser Internal Representation, or HLO IR), search for weak spots, and try to eliminate them.
Abstractions at a similar level to HLO IR have been shown to be convenient for optimising linear algebra operations (Barthels et al., 2021).
We add match-and-replace operations, e.g. to introduce a more efﬁcient distance computation, reshufﬂing operations for expressions that are invariant to evaluation order, and splitting with large tensors to reduce memory usage.
mul(A v):
C =
A @ @ v The match and replace optimisation pass, also known as peephole optimisation (Aho et al., 2006), searches for expressions in a data-ﬂow graph for which we know in advance that an equivalent and more efﬁcient version exists.
For example, we search for expressions that compute Euclidean distance in naive form between vectors of length n and m with a dimension d .
The naive Euclidean distance computation uses broadcasting over the dimension d and creates a temporary tensor with entries ( x nd − y md ) 2 of size n × m × d .
This can be replaced with (cid:80) d x 2 nd + y 2 md − 2 x nd y md , where the largest tensor has size n × m .
Replacing sub-parts of the graph is a standard procedure in compilers like XLA, although many linear algebra tricks have been missing.
While the Euclidean distance is the only match-and-replace optimisation we implement, other operations can easily be added, for example, efﬁciently adding diagonals to matrices without allocating a dense square tensor where only the diagonal is non-zero.
4.2 Reordering
A computational data-ﬂow graph is an ordered sequence of operations, with the order of operations inﬂuencing the memory usage.
In some cases, reordering sub-parts of the data-ﬂow graph can lead to reductions in the memory footprint.
The classical example of reordering is the optimisation of matrix chain multiplications.
For example, consider the matrix expression C = AB v for matrices
A , B ∈ R n × n , and v ∈ R n .
In the listing 1, the order of operations determines that the matrix multiplications are performed from left to right, i.e. C = ( AB ) v , which gives the most inefﬁcient execution order with runtime complexity O ( n 3 ) and memory complexity O ( n 2 ) .
Changing the order to C =
A ( B v ) improves time complexity to O ( n 2 ) and practical memory complexity because the intermediate multiplication result of B v is a vector not a matrix as in the case of AB multiplication.
The optimisation of matrix chain multiplication is possible due to the associativity of matrix multiplication, such that the result of the matrix multiplication chain does not depend on where parentheses are placed.
There are many efﬁcient and sophisticated algorithms for addressing this task (Chin, 1978;
Czumaj, 1996; Barthels et al., 2018; Schwartz and Weiss, 2019).
We implement a simple matrix chain multiplication optimisation algorithm using dynamic programming in a bottom-up fashion (Barthels et al., 2018, §2).
Our version of matrix chain multiplication algorithm supports expressions with matrix multiplications, sum reductions and transposes.
4.3 Data-ﬂow graph splitting
Often, a part of a computational data-ﬂow graph can be divided into multiple independent copies, such that each copy of the data-ﬂow graph or its part act on a slice of the input tensor, and the results are combined afterwards in some fashion.
This splitting approach is also known as a MapReduce technique (Dean and Ghemawat, 2008), where a computation is divided into smaller and less expensive parts (map) and then combined into the ﬁnal result (reduce).
The splitting technique is common for distributing the computational load.
The focus of existing solutions is on exploiting hardware parallelism or utilising multiple devices.
Instead, we use the same techniques for reducing total memory consumption, which is possible because the memory for individual map operations can be freed before the whole result is computed.
An optimisation pass starts by running a depth-ﬁrst search from the ﬁnal result of the computation.
The operations dot or reduce_* are special, as they often indicate that a computation involving a large tensor can give a smaller result.
Once a dot or reduce_* operation is marked as fully traversed, we recursively search the traversed paths for operands that are impractically large tensors, until we reach operands that are deemed small enough.
Along the way, we keep track of which operations are applied, and along which axes they are trivially parallelisable.
The result is a collection of sub-graphs,
1: procedure H ANDLE D OT ( dot : HloInstruction ) 2: if output_size ( dot ) ≥ tensor_size_threshold , i.e. dot is splittable then 3:
Exit and continue traversing succeeding operations in the data-ﬂow graph and search for size-reducing operation for the output tensor of dot .
4: if dot.rhs is not splittable and dot.lhs is not splittable then 5:
Exit and continue traversing the data-ﬂow graph.
6: if dot.rhs (cid:54) = dot.lhs and both operands are splittable then 7:
Exit and continue traversing the data-ﬂow graph.
Let operand_to_split = dot.rhs or operand_to_split = dot.lhs depending on previous splittability conditions.
Let split_dims =
{ d 1 , . . .
, d n } , split_producers = { op 1 , . .
op m } , s.t.
10: ◦ ∀ op ∈ split_producers exists a path op (cid:32)
operand_to_split 11: ◦ ∀ op ∈ split_producers , input_size(op) ≤ tensor_size_threshold 12: ◦ ∀ op ∈ split_producers , ∃ d ∈ split_dims which is splittable on the path op (cid:32) operand_to_split 13: if split_dims = ∅ or split_producers = ∅ then 14:
Exit and continue traversing the data-ﬂow graph.
Let best_split_dim
=
d ∈ split_dims , and ops ⊆ split_producers , s.t.
16: ◦ min d ∈ split_dims (cid:98)
d ÷ split_size(operand_to_split, tensor_split_size) (cid:99) 17: ◦ ∀ op ∈ ops , the path op (cid:32) operand_to_split is splittable on best_split_dim 18:
Let split_size =
(cid:98) best_split_dim
÷ split_size(operand_to_split, tensor_split_size) (cid:99) 19:
Create while loop HloInstruction , s.t.
20: ◦
The loop iterates splits of
size split_
size at best_split_dim of paths ops (cid:32) operand_to_split 21:
◦ The loop applies dot reduction operation on the slice of operand_to_split
◦
The slice result of the dot reduction operation is put into the replica of the original dot.
result 23:
Replace ops (cid:32) dot.result instructions with created while loop
that start at operations that produce large tensors, and end at operations that reduce them again, together with candidate axes that they can be split across.
According to some heuristics which ensure appropriate sizes for intermediate results, we then turn this entire computation into a while loop, where each iteration computes a manageable part of the ﬁnal result (ﬁg. 1).
Checking if the axis is splittable is necessary as not all operations act independently on each dimension.
For example, element-wise operations can be split on any axis, whereas the triangular solve operation can be split on “batching” dimensions only.
Next, the data-ﬂow graph splitting procedure selects the dimension of the largest size which contributes the most to memory.
As we discussed earlier, the decision about where to split the graph depends on the tensor size.
We offer two XLA options to the user for out-of-memory mitigation: tensor size threshold and tensor split size upper bound .
Tensor size threshold is a criterion designed for detecting which operations should be marked as candidates for splitting.
Tensor split size upper bound serves as a threshold on the largest allowed chunk size for splitting.
These options are set equal by default.
The command-line snippet at listing 2 shows how a user would use these options by passing them via an environment variable, and the snippet is indifferent to the machine learning framework used by the script.
Minimal user effort is required for using our XLA compiler extension.
The user is involved only in deﬁning what the suitable threshold and splitting sizes are.
Listing 2 Example of how a user can set options for the extended XLA using the environment variable.
XLA
[Equation]
FLAGS
[Equation]
tensor
[Equation]
size
[Equation]
500MB
[Equation]
python train.py
I[a]
G
A[..., i:j, ...]
O[...]
I’[i:j]
A’[..., i:j, ...]
O’[...] I[a]
Loop
R
S
R’ G’
One strong beneﬁt of our compiler-based solution, is that the computational graph represents the whole pipeline of computations, including forward and backward propagation.
Our splitting procedure will be applied automatically, regardless of how many derivatives need to be computed.
In addition, our procedure encompasses two splitting schemes that the machine learning literature distinguishes: model-based and data-based splitting schemes of the data-ﬂow graph.
Model-based splitting schemes involve partitioning the model over its parameters, whereas data-based splitting schemes batch over inputs.
The proposed splitting approach supports both schemes.
However, our implementation of model-based splitting supports only splitting memory intensive expressions involving model parameters.
It is not able to split model parameters themselves into smaller chunks.
4.4 XLA limitations
While we still believe that XLA is the right framework for our extensions, several limitations came to light during implementation.
One limitation that is shared with all current frameworks, is that they only have a weak linear algebra type system, where matrices are represented as arrays without additional properties.
Solutions that support stronger type systems (Bezanson et al., 2017; Barthels et al., 2021) may be able to implement a wider variety of match-and-replace optimisations.
Another limitation comes from the default memory allocation manager not being aware of memory limits.
Its current behaviour is to execute nodes in the computational graph, and therefore allocate any required memory, as soon as the required inputs have been computed.
This means that even if tensors are split to manageable sizes, memory overﬂows can still occur if several are executed simultaneously.
To prevent this from happening, we had to use limits that were smaller than our total GPU memory.
Experiments
This section shows how existing software packages take advantage of our extension to XLA (eXLA).
We demonstrate our optimisations on non-parametric k-nearest neighbours (kNN) and sparse Gaussian process regression (SGPR) models, and Transformers.
5.1 Matrix-Vector Multiplication
We start by demonstrating the improved efﬁciency that eXLA offers to large-scale matrix-vector multiplications of the form y
= K v , where K is an n × n kernel matrix, and y, v ∈ R n .
Such computations are common in Conjugate-Gradients-based Gaussian process approximations (Gibbs and Mackay, 1997;
Wang et al., 2019a; Artemev et al., 2021), where K ij = k
( x i , y j ) and k is some kernel function.
We choose the common Squared Exponential.
We implement this equation using GPﬂow (Matthews et al., 2017), a TensorFlow-based package that provides a convenient software interface for Gaussian processes and kernel functions.
Without eXLA, the entire K would be stored in memory, leading to a n 2 memory cost.
This makes running on
[Figure]
Dataset
Distance n d
KeOps eJAX eTF
JAX
TF
Random
L1e
Random
L1e  3
Random
L1e  24  ∅ ∅
Random
L1e  123  ∅ ∅
MNIST
L6e
MNIST
L6e
Fashion
L6e
Fashion
L6e
GloveCosine 1 . 18e  3 464  ∅ ∅
GloveCosine 1 .
18e
[Table]
large datasets infeasible, where e.g. n = 10 6 would lead to a memory requirement of 8 TB, which is impractical even for the largest of modern GPUs with 40 GB of memory.
A memory efﬁcient split/implicit implementation is necessary to scale to large datasets, as was impressively done by
Wang et al.
(2019a)
, but is cumbersome.
We ran our implementation with eXLA enabled, which allows a user to control the memory of an algorithm.
We evaluated the expression in double precision on a Tesla V100 GPU with 32 GB of memory, and applied a range of memory limits.
In ﬁg. 2 we report the peak memory consumption and execution time of evaluating the kernel matrix-vector product for different sizes, with different memory limits applied.
We see that the memory constraints are not violated, and that dataset sizes are used that are far beyond the 32 GB memory capacity.
5.2 K-Nearest Neighbours
K-nearest neighbours is a fundamental machine learning algorithm, with a similar large memory cost.
A kNN query selects k closest data points in the dataset to each query point.
Brute-force implementations compute pairwise distances between m query points and n data points, resulting in the distance matrix of size m × n .
This is followed by a topk operation, which is often naively implemented using column-wise sort operation on the distance matrix.
Our benchmarks show that eXLA scales the brute-force approach and does not fail for large problems, i.e. large n and m .
We compare TensorFlow and JAX implementations with and without eXLA optimisations, and a KeOps implementation.
We use randomly generated data, common benchmarks like MNIST and Fashion-MNIST, and Glove-50, Glove-100 and Glove-200 from the ANN-benchmark toolkit
Aumüller et al.
[Table] [Table]
Our results are listed in table 3 (see the appendix for a full table that reproduces Feydy et al.
(2020, table 3)).
In all benchmarks, we set the tensor size threshold for eXLA to 100 MB for simplicity, even though this may not be optimal for performance.
We observe that eXLA prevents memory overﬂows in JAX and TensorFlow.
In addition, performance is comparable or higher.
We acknowledge that
KeOps performs signiﬁcantly better than any JAX or TensorFlow implementation.
This is explained by 1) JAX/TF not having efﬁcient implementations for certain functions (e.g. topk runs a full sorting algorithm), and 2) KeOps having implemented additional optimisations, which could also be added to XLA.
However, we note that we also achieved our goal of improving the memory and time performance of a JAX/TensorFlow implementation without changing the code .
5.3 Sparse Gaussian Process Regression
Gaussian processes (Rasmussen and Williams, 2006) are considered the gold standard method for performing regression with uncertainty estimates.
A straightforward implementation requires taking a matrix decomposition of an n × n kernel matrix (like those considered in section 5.1), which leads to an O ( n 3 ) time cost, and an O ( n 2 ) memory cost.
Scaling Gaussian process is challenging, which is often attributed to the time cost.
In reality however, large datasets cause memory overﬂows far before long runtimes become an obstacle.
[Figure]
Approximate methods have been introduced to deal with both the time and space issues.
While there are many variants (Quiñonero-Candela and Rasmussen, 2005), we consider the sparse variational approximation (Titsias, 2009) for which a naive implementation has O ( nm 2 + m 3 ) time cost, and
O ( nm + m 2 ) memory cost.
Here, m denotes the number of inducing variables , which controls the quality of the approximation.
Under certain conditions, the method provides reliable hyperparameter selection (Bauer et al., 2016), and very accurate posterior approximations (Burt et al., 2019, 2020) while using m (cid:28) n .
In practice, these standard implementations may still have their performance limited by how large m can become before a memory overﬂow occurs.
A more memory-efﬁcient implementation with a memory cost of O ( m 2 ) does exist (Gal et al., 2014), but is so cumbersome to implement that it is not widely used or compared against.
Fortunately, the splitting optimisation we implemented in eXLA can discover the same procedure that was engineered by Gal et al.
Moreover, since eXLA operates on the entire computation graph, it optimises gradients as well as the optimisation objective function with no additional effort, while
Gal et al.
(2014) needed to implement gradients manually.
We demonstrate the utility of eXLA by scaling the GPﬂow (Matthews et al., 2017, 2.3.1 release version) implementation of Sparse Gaussian process regression (SGPR, Titsias, 2009), without any modiﬁcations of the code.
With our eXLA optimisations, SGPR was able to scale to much larger datasets, with more inducing points.
We conduct experiments on a Tesla V100 GPU with 32 GB of memory, and run on two of the largest UCI datasets that are commonly considered in Gaussian process research: 3droad and houseelectric with total number of data points 434 , 874 and 2 , 049 , 280 respectively.
We set the tensor size threshold and the tensor split size in eXLA to 1 GB.
We primarily compare to Wang et al. (2019a), who use a Conjugate Gradients approximation (Gibbs and Mackay, 1997) to achieve the most impressive scaling of a Gaussian process approximation to date, using an impressively engineered implementation that manually splits and distributes parts of the computation.
In ﬁg. 3 we compare GPﬂow’s SGPR implementation with and without eXLA as we increase the number of inducing points.
We see that until about 800 inducing points, the normal and eXLA runs result in the same predictive metrics, as desired.
After 800 inducing points, runs without XLA fail with an “out of memory” error, while with eXLA we scaled to 10 4 inducing points.
Simply scaling the method in this way leads to signiﬁcant performance improvements.
We compare predictive accuracies directly with the scalable Conjugate Gradients implementation of
Wang et al. (2019a).
In that paper, SGPR was discussed as a method that would not scale, probably due to the difﬁculty of implementing it in a memory-efﬁcient way as in Gal et al.
Tableshows that simply scaling SGPR using eXLA makes it competitive with the Conjugate Gradients implementation of Wang et al.
(2019a)
, without needing additional hardware.
5.4 Transformers
Neural network architectures based on the attention mechanism, such as Transformers (Vaswani et al., 2017), have become a common tool for solving natural language processing tasks.
Transformers have also been applied successfully in computer vision classiﬁcation tasks (Dosovitskiy et al., 2020) and image generation (Ramesh et al., 2022).
The attention mechanism Att ( Q, K, V ) = softmax (cid:16) QK T √ d (cid:17) V , with Q, K ∈ R b × s × d k , and V ∈ R
b × s × d v , is an expression that lies in the core of Transformer models.
The memory consumption by the intermediate result QK T grows quadratically with dimension s , i.e. O ( bs 2 ) .
Rabe and Staats (2021) investigated ways of reducing the memory overhead by sequentially computing parts of the attention operation with large s .
Their technique is very similar to the eXLA’s splitting optimisation pass from section 4.3.
The advantage of eXLA is the automatic application of such schemes, depending on sizes of the input dimensions b , s and d as well as user’s eXLA setting for the tensor size limits.
We test eXLA on a sequence-to-sequence language translation Transformer.
We use the conﬁguration of the translation Transformer from the TensorFlow website 3 .
We increase the size of the attention mechanism by increasing the size of the randomly generated input sequence until the code fails with a memory overﬂow; we then show that running the same implementation with eXLA avoids this failure.
In the experiment with the language translation Transformer this corresponds to increasing the input sequence.
We run experiments on a single Nvidia V100 GPU with 32GB memory.
Figure
[Figure]
l a p
[Figure]
e d t i
[Figure]
e , s e c n shows that out of the box implementations of translation Transformer using TensorFlow with its default compiler (TF) and TensorFlow with the XLA backend fail to run when the sequence length larger than 2000 .
The Transformer model compiled with eXLA optimisations managed to run with sequences up to 7000 with tensor limit set to 10 GB and the tensor split size set to 1 GB.
The drop in the ﬁg. 4 after 2000 is explained by eXLA’s splitting optimiser modifying the graph with tensor splits equal to 1 GB.
After eXLA optimisation applied, there are no tensors in the graph larger than 1 GB which can be allocated at the same time together on the GPU.
In these situations a practitioner might use larger split size if the hardware allows it, which in turn will boost the computational performance.
Discussion
We showed that our XLA compiler optimisations (eXLA) prevented memory overﬂows in algorithms with large tensor or linear algebra operations.
Our compiler extension automatically adjusts computational data-ﬂow graphs to control memory utilisation.
As demonstrated in the experiments section, machine learning models compiled with eXLA could run at a greater scale, whereas their out-of-the-box implementations failed with Out of Memory errors.
Crucially, we used existing software packages without modifying any code.
In addition to showing that our compiler extensions work as intended, our experiments also provide directly useful empirical results for Gaussian processes.
We managed to run an “old” method (SGPR,
Titsias, 2009), with unchanged code, to obtain empirical results that outperformed a state-of-the art method (Wang et al., 2019a).
This shows that existing methods may simply need to be scaled in order to provide high-quality solutions, which are even guaranteed by theory in certain cases (Burt et al., 2020).
We also provide proof-of-concept experiments that show that eXLA can help with running
Transformers (Vaswani et al., 2017) which are notoriously memory hungry, on cheaper hardware without increasing software complexity.
The exciting possiblity of eXLA is that it opens up the possibility to probe behaviour of machine learning models in regimes that were previously infeasible, and on cheap hardware.
For example, one could train very wide neural networks, to empirically compare to behaviour predicted by NTK theory (Lee et al., 2018; Matthews et al., 2018;
Jacot et al., 2018; Novak et al., 2020).
The current implementation of eXLA is still only a demonstration of what compiler optimisations could achieve, and many more optimisations can be added.
We believe that increasing the capability of compilers like XLA will greatly increase the efﬁciency of researchers and practitioners.
We hope that community-driven compiler projects will contribute to the community in a similar way to how existing numerical frameworks already do.
Acknowledgements
Thanks to David R. Burt and Sebastian W. Ober for the feedback on the draft of this paper, Fergus
Simpson and Alan Saul for their advice on Transformer models.
We also would like to thank Lev
Walkin and Renat Idrisov for discussions about compilers in the beginning of this project.
A Appendix A.1 Code
The code for benchmarks and experiments is available at https://github.com/awav/gambit , and the fork of TensorFlow repository with extension to XLA compiler (eXLA) is available at https://github.com/awav/tensorflow .
A.2 Additional experiments
Dataset
Distance n d
KeOps eJAX eTF
JAX
TF
Random
L
[Table]
Random
L
Random
L
Random
L     ∅ ∅
Random
L     ∅ ∅
Random
L     ∅ ∅
Random
L
[Table]
∅ ∅
Random
L
[Table]
∅ ∅
Random
L     ∅ ∅
Random
L     ∅ ∅
Random
L     ∅ ∅
Random
Cosine     ∅ ∅
Random
Cosine     ∅ ∅
MNIST
L
MNIST
L
Fashion
L
[Table]
Figures 5 and 6 depict XLA HLO graphs for kernel matrix-vector multiplication before and after splitting optimisation in eXLA optimisation pipeline (section 4.3), respectively.
The same conﬁguration of the kernel is used as in section 5.1, i.e. squared exponential kernel from Matthews et al.
By kernel matrix-vector multiplication expression we mean the function g ( x , y , v ) =
k ( x , y ) v , where k ( x , y )
=
σ 2 exp (cid:16) − 1 / 2 (cid:107) x − y (cid:107) 2 /l 2 (cid:17) is the kernel with σ 2 and l hyperparameters.
The size of 1-dimensional input vectors v , x and y is 1e − 6 .
In turn, the size of the corresponding kernel matrix is 1e − 6 × 1e − 6 , and in double precision would require to allocate 8TB.
The tensor size threshold was set to 1GB, and eXLA splitting optimisation pass divided the expression of the kernel matrix-vector multiplication into smaller chunks, such that the maximum tensor size in the graph is 125 × 1e − 6 .
[Figure]
[Table]
[Figure]