Less memory, more AI – Method development for smaller language labs

For a smaller research lab like AI Sweden's NLU group, this memory limitation becomes a critical bottleneck. Amaru's research therefore goes beyond the traditional escalation of resources and instead focuses on streamlining the processes and the underlying mathematics in the algorithms that form the foundation of language models. His research is driven by questions such as: How does one effectively handle large amounts of data? What can be done with limited memory? And which parts of the process can be changed to enable large-scale operation?

The role of GPU memory in LLM training

High-quality GPUs are fundamental to LLM development, and the same applies to strategies for utilizing their full capacity. The complex mathematical operations that the models are based on are dependent on the GPUs' memory. Bottlenecks arise in training if performance is insufficient, which creates delays – a type of unwelcome limitation for a group like NLU, which works in time-limited projects with funding from the European Union and Vinnova, such as the EU flagship project OpenEuroLLM.

The GPU's memory, or VRAM, differs from traditional RAM memory in that it sits directly on the GPU card, which is necessary for massive parallel data processing. During training, all of the model's parameters, along with the data being processed and all intermediate results (such as gradients in backpropagation), must be stored in this memory. When a language model grows to billions of parameters, the demands on the hardware increase enormously. The same applies when the model's activations grow faster than the data size, which happens, for example, with naive attention. In these situations, the limited VRAM quickly becomes a bottleneck that determines whether the calculation can be performed at all.

The pursuit of simplicity

The core of Amaru's research is about solving the problem of inefficient memory usage found in neural network layers like attention. A tenfold increase in a sequence results in a hundredfold increase in memory usage in a naive implementation. His work here involves reformulating the computational problem. It's about moving the method from keeping the entire dataset in memory to handling it incrementally through online calculations, similar to how one calculates an average of a million numbers without loading all of them at once. His theory is called FlashReduce and is reminiscent of MapReduce, as it involves solving subproblems and then combining the solutions into a final result.

The theoretical framework Amaru has developed formalizes this process. This is done by breaking down large, computationally intensive problems into small subproblems, a method called tiling (from the word "tiles"). Based on this, he has developed a proof of concept that streamlines memory usage. The underlying idea originated in FlashAttention, which can be seen as an example of this, more general, method. 

For the researcher, the beauty lies in finding a solution where only information about the final result and a single input value are needed to calculate the gradient, instead of considering all individual parts. This is possible as long as the merging function of the sub-solutions meets certain criteria, which Amaru describes as "not particularly hard," giving developers great freedom in design.

Reduced implementation burden

One of the goals of FlashReduce is to reduce the complex burden for other developers and researchers. Just as MapReduce handles the complex details of distributed computing, his framework could relieve the developer. The developer would then only need to define how to solve the subproblems and how the solutions are combined, and let the framework identify the best way to execute these components on the GPU. This makes it easier to develop and experiment with memory-efficient kernels and is a possible method for creating larger and more efficient open source models.

This approach demonstrates a perspective where "simplicity and elegance in mathematical theory can be translated into practical efficiency," he argues. Although a proof of concept shows progress, the remaining hurdle is how to create a sufficiently user-friendly and effective implementation on GPUs. Unlike the theoretical framework, which is simple and limits the researcher's search space by showing which layers can be made memory-efficient, the real challenge lies in the constantly shifting and complex hardware limitations. The goal is for the framework to handle these details completely automatically. This way, the programmer avoids handling the "degrees of freedom" otherwise required, ensuring optimal efficiency. "This is a possible way to create larger and more efficient models and also to reduce the implementation burden," says Amaru.

An enabler for open source in Europe

Amaru's research—on reducing the implementation burden and enabling larger, more effective models within the open source category—gains particular weight when considering the European AI landscape. OpenEuroLLM, a large initiative funded by the European Union with the goal of creating transparent AI for Europe by releasing large-scale language models under an open license, is dependent on technical breakthroughs like Amaru's. In Europe, where access to computational resources has historically been limited, the pursuit of efficiency becomes a strategic necessity, rather than merely an academic challenge.

Amaru's theoretical framework, which makes it easier to develop memory-efficient kernels and thereby reduces the implementation burden for developers, is the type of research that is crucial for European open source models to compete in size and performance with other international alternatives. The research thus addresses the central key issue: how can European actors create and maintain cutting-edge expertise in large-scale AI, without relying solely on the scaling up of resources? By researching new frameworks and processes, AI Sweden aims to contribute to creating solutions for problems such as the hard limitation in GPU memory, thereby strengthening the European capability to build large-scale and transparent AI under an open license.