|
| 1 | +<!-- |
| 2 | + Copyright 2025 Google LLC |
| 3 | +
|
| 4 | + Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | + you may not use this file except in compliance with the License. |
| 6 | + You may obtain a copy of the License at |
| 7 | +
|
| 8 | + https://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +
|
| 10 | + Unless required by applicable law or agreed to in writing, software |
| 11 | + distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | + See the License for the specific language governing permissions and |
| 14 | + limitations under the License. |
| 15 | + --> |
| 16 | + |
| 17 | +# Tiling |
| 18 | + |
| 19 | + |
| 20 | +Often in high-performance computing, there's a trade-off between memory usage and computation time. **Tiling** (also known as chunking) is an optimization technique that prioritizes reducing memory usage, but sometimes at a cost to running time. |
| 21 | + |
| 22 | +The core idea is to partition a large tensor into smaller, more manageable blocks called **tiles**. Instead of loading the entire tensor into memory, a program processes these tiles sequentially. This significantly lowers the peak memory required to perform an operation. |
| 23 | + |
| 24 | +While processing data in chunks can introduce minor computational overhead, an efficient tiling strategy minimizes this cost. By preventing out-of-memory errors, tiling enables programs to handle larger problems than would otherwise be possible. This often leads to better hardware utilization and improved end-to-end performance metrics, such as MFU. |
| 25 | + |
| 26 | + |
| 27 | +## The Concept of Tiling |
| 28 | + |
| 29 | +The effectiveness of tiling stems from the **linearity** of many operations in LLMs. Operations like matrix multiplication and gradient calculation can be broken down into smaller, independent sub-problems. |
| 30 | + |
| 31 | +For instance, consider the matrix multiplication operation: |
| 32 | + |
| 33 | +$$A[M, N] \times B[N, 1] = C[M, 1]$$ |
| 34 | + |
| 35 | +If the matrices `A` is too large to fit into memory, you can tile the operation. By splitting matrix `A` into $K$ smaller chunks along its `M` dimension ($A_0[M/K, N], \dots, A_{K-1}[M/K, N]$), you can load each chunk separately and compute a corresponding portion of the output matrix `C` ($C_0[M/K, 1], \dots, C_{K-1}[M/K, 1]$) follows |
| 36 | + |
| 37 | +$$A_i[M/K, N]\times B[N, 1] = C_i[M/K, 1] \quad \forall i=0, \dots, K-1.$$ |
| 38 | + |
| 39 | +Finally, you can concatenate the smaller `C` matrices to form the complete result. |
| 40 | + |
| 41 | +This principle extends to the backward pass as well. Instead of computing the full gradient `dA[M, N]`, which also exceeds memory capacity, you can compute the gradient for each tile individually: |
| 42 | + |
| 43 | +$$dC_i[M/K, 1] \times B^\intercal[1, N] = dA_i[M/K, N] \quad \forall i=0,\dots,K-1.$$ |
| 44 | + |
| 45 | +Similarly, the gradient on `B` is the accumulation |
| 46 | + |
| 47 | +$$\sum_{i=0}^{K-1}dC_i^\intercal[1, M/K] \times A_i[M/K, N] = dB[N, 1] \quad \forall i=0,\dots,K-1.$$ |
| 48 | + |
| 49 | +This tiling approach reduces the peak memory usage from $\mathcal{O}(MN)$ to $\mathcal{O}(MN/K)$, which facilitates model training with limited memory resources. |
| 50 | + |
| 51 | + |
| 52 | +## Tiling in MaxText |
| 53 | + |
| 54 | +### Gradient Accumulation |
| 55 | + |
| 56 | +Gradient accumulation (GA) is a form of tiling that splits a full data batch into smaller micro-batches. |
| 57 | + |
| 58 | +In MaxText, the `gradient_accumulation_steps` config determines the number of micro-batches. When this is set, `per_device_batch_size` effectively becomes the size of each micro-batch. During training, the model processes one micro-batch at a time, and the resulting gradients are accumulated. A full parameter update occurs only after all micro-batches have been processed. |
| 59 | + |
| 60 | +GA reduces the size of activations in memory at any given moment, which is crucial for training with large batch sizes or when using sharding strategies like pipeline parallelism. |
| 61 | + |
| 62 | + |
| 63 | +*Figure 1: Gradient accumulation tiles a global batch into smaller micro-batches.* |
| 64 | + |
| 65 | + |
| 66 | +### Vocabulary Tiling |
| 67 | + |
| 68 | +Vocabulary tiling is another memory-saving technique designed to handle the large vocabulary sizes in modern language models. |
| 69 | + |
| 70 | +The final output unembedding layer of a language model maps hidden states to logits—a large tensor of scores with a shape of `[per_device_batch_size, seq_length, vocab_size]`. This tensor can be enormously large and cause out-of-memory errors. |
| 71 | + |
| 72 | +Vocabulary tiling avoids materializing the full logits tensor. Instead, it tiles the input hidden states and computes the logits, loss, and gradients one tile at a time. Unlike GA, which is applied at the start of the model, vocabulary tiling is applied only to the input of the final layer. |
| 73 | + |
| 74 | +In MaxText, the `num_vocab_tiling` configuration controls the number of tiles. This technique is especially advantageous for models with large vocabularies (e.g., Gemma and Llama), particularly when training with long sequence lengths. By preventing out-of-memory errors, vocabulary tiling can enable simpler sharding strategies (like FSDP) and unlock better computational performance. |
| 75 | + |
| 76 | + |
| 77 | +*Figure 2: Vocabulary tiling processes hidden states in tiles to avoid generating the full logits tensor.* |
| 78 | + |
| 79 | +### Other Tiling Methods |
| 80 | + |
| 81 | +Tiling is also crucial for managing data movement across the memory hierarchy (HBM, VMEM, VREGs). It's applied at various levels: manually within low-level custom operations like Pallas kernels (e.g., Splash Attention) to precisely control HBM to VMEM transfers, and automatically by the XLA compiler for standard computations such as matrix multiplications. This hierarchical approach of breaking data into smaller, manageable tiles ensures efficient hardware utilization by pipelining data through stages. |
| 82 | + |
| 83 | +## Tiling vs. Sharding |
| 84 | + |
| 85 | +**Tiling** and **sharding** are independent concepts that do not conflict; in fact, they are often used together. Sharding distributes a tensor across multiple devices, while tiling processes a tensor in chunks on the same device. |
| 86 | + |
| 87 | +To learn more about sharding in MaxText, please refer to the [sharding documentation](https://maxtext.readthedocs.io/en/latest/explanations/sharding.html). |
0 commit comments