AI At Home Part 2: Multi GPU Drifting

In the previous section, the author constructed a makeshift home server utilizing discarded GPUs to facilitate AI language model processing. This article will delve into the necessary steps to optimize performance within such a system. The article is derived from the existing code and techniques, with a focus on tweaking the llama.cpp settings, rather than developing new ROCm kernels.

The article will commence with a brief background; however, it will bypass the details of transformer models and multi-GPU parallelism, and instead concentrate on the aspects relevant to enhancing the performance of this particular setup. AI text generation software currently in use leverages the transformer model or large language model.

The transformer model was introduced in the highly influential paper, 'Attention is All You Need'. The basic concept of this model has been widely adopted, making it likely that every software engineer reading this article has familiarized themselves with the concept. The article will simplify the explanation, focusing on the perspective of someone attempting to run these models efficiently on subpar hardware, rather than delving into tensor mathematics or model training.

The large language model operates using tokens, which are essentially word fragments. Instead of dealing with individual letters, the model processes sequences of letters. This process is more efficient, as it allows the model to handle larger chunks of language at once. Each time the model generates a token, it generates a probability distribution, randomly selecting one from this distribution.

This is crucial as it allows the model to generate more coherent and contextually relevant output. The model is composed of layers, including an input layer, an output layer, and several hidden layers. The input layer receives the entire input prompt, while each subsequent layer performs calculations on the output of the previous layer.

This process is known as the attention mechanism. The attention mechanism allows the model to consider the relationships between different tokens throughout the input series. This is what sets AI language models apart from simpler models like Markov chains. The attention mechanism is crucial to understanding how AI language models generate their outputs.

Each time the model generates a token, it reads in the existing context and all the weights of the model, performing matrix multiplications for each layer. This process is limited by memory bandwidth, not compute. Therefore, the server built in the previous section contains multiple GPUs, each with a large amount of VRAM, to quickly access the necessary data.

The Gemma4-31B model, for example, has 31 billion weights. The author emphasizes that loading these weights takes significantly more time than the actual attention calculations. This is why the server features multiple GPUs with large amounts of VRAM, to ensure that the model weights and KV cache can be quickly accessed. The author also discusses the memory shortage in the industry, driven by the shift towards high-bandwidth memory for data center GPUs.

This memory limitation poses a significant challenge for scaling up AI language models. To address this issue, the article introduces the concept of Mixture of Experts (MoE) model architecture. The main idea behind MoE is to process the input embedding matrix using the first couple of layers every time, then routing this information to a subset of the model.

Only a small fraction of the weights, referred to as 'experts', need to be loaded for each token. For instance, Deepseek V4 Flash has 284 billion weights, but only 13 billion get activated for each token. This approach significantly reduces the amount of data that needs to be loaded and processed for each token, thereby improving the model's performance.