# RAM Offloading

RAM offloading is a method of using system RAM during training to reduce the need for VRAM. This document serves as a
technical documentation of the implementation in OneTrainer. The method described works best for diffusion
transformers (DiT), but can be applied to other models as well.

## Introduction

When training (or testing) a multi-layer neural network, only a single layer is ever needed at any given time. This fact
can be used to swap layers between RAM and VRAM as they are needed. A naive implementation could

1. Load the first layer of the model into VRAM
2. Execute that layer
3. Offload that layer back into RAM

Then repeat that process until all layers have been executed. While this method would work, it would significantly
impact performance. A more advanced method can use the asynchronous capabilities of modern GPUs to transfer layers of
the model between RAM and VRAM while other layers are executed. In PyTorch this can be done using the Stream feature in
combination with non_blocking=True tensor copies.

In addition to layer offloading, intermediate activations generated by gradient checkpointing can be offloaded to
further reduce VRAM consumption.

## Layer Offloading

To give the user more control over the amount of RAM they want to use for offloading, OneTrainer has a single offloading
fraction parameter. This value defines how many layers of the model are offloaded to RAM at any time.

The offloading process starts by moving only the first few layers of the model into VRAM. All remaining layers are
stored in RAM. After a layer was executed, it is immediately offloaded to RAM, and one additional layer is loaded into
VRAM. This process is done asynchronously while the model is executed. To do this, two things need to be done:

1. Layers in RAM are stored in pinned memory.
2. The transfer is executed on a separate CUDA stream, and the non_blocking=True flag is passed to every call to a copy
   function.

After each layer transfer, a CUDA event is recorded on the transfer stream. This event is used to let the default stream
wait for the transfer to be finished before executing the layers.

### Offloading scheme

Layers are cycled through from front to back (during a forward pass), or back to front (during a backward pass). This
method tries to load as many layers as possible ahead of time to reduce the wait time spent waiting on layers that are
still in transit.

This is best explained by an example. Let's imagine a model with 9 layers and an offloading fraction of 33%, meaning 3
layers are offloaded to RAM. The following diagram shows the location of layers after each execution of a layer. Keep in
mind that this is only an idealized case. In the real world, layers are constantly in transit, and timing differences
can mean that a layer transfer takes longer than the execution of a layer.

Each layer is represented by a column, each row represents a point in time.

<details><summary>Execution diagram during sampling</summary>

- `->` execution direction (forward or backward pass)
- `■` executed block
- `X` block in VRAM
- `_` block in RAM

```
->  ■ X X X X X _ _ _  (first step starts)
->  _ ■ X X X X X _ _
->  _ _ ■ X X X X X _
->  _ _ _ ■ X X X X X
->  X _ _ _ ■ X X X X
->  X X _ _ _ ■ X X X
->  X X X _ _ _ ■ X X
->  X X X X _ _ _ ■ X
->  X X X X X _ _ _ ■
->  ■ X X X X X _ _ _  (second step starts)
->  _ ■ X X X X X _ _
->  _ _ ■ X X X X X _
->  _ _ _ ■ X X X X X
->  X _ _ _ ■ X X X X
->  X X _ _ _ ■ X X X
->  X X X _ _ _ ■ X X
->  X X X X _ _ _ ■ X
->  X X X X X _ _ _ ■
```

</details>

<details><summary>Execution diagram during training</summary>

- `->` execution direction (forward or backward pass)
- `■` executed block
- `X` block in VRAM
- `_` block in RAM

```
->  ■ X X X X X _ _ _  (forward pass starts)
->  _ ■ X X X X X _ _
->  _ _ ■ X X X X X _
->  _ _ _ ■ X X X X X
->  _ _ _ X ■ X X X X
->  _ _ _ X X ■ X X X
->  _ _ _ X X X ■ X X
->  _ _ _ X X X X ■ X
->  _ _ _ X X X X X ■
<-  _ _ _ X X X X X ■  (backward pass starts)
<-  _ _ X X X X X ■ _
<-  _ X X X X X ■ _ _
<-  X X X X X ■ _ _ _
<-  X X X X ■ X _ _ _
<-  X X X ■ X X _ _ _
<-  X X ■ X X X _ _ _
<-  X ■ X X X X _ _ _
<-  ■ X X X X X _ _ _
```

</details>

### Optimizer steps

While this offloading scheme can greatly reduce VRAM usage, it has a pretty big downside. Optimizer steps usually
require all model parameters to be placed on the same device (like the GPU). To fix this issue, a technique called
"Fused Back Pass" can be used. Optimizer steps are executed immediately after the gradient calculations. That way, each
parameter is updated only when it's located in VRAM. This limitation only applies to full fine-tuning. In case of LoRA
training, all trainable parameters are stored in VRAM.

## Activation offloading

In addition to layer offloading, intermediate activations can also be offloaded to RAM. Gradient checkpointing already
reduces the amount of memory needed to store activations. But even then, each checkpoint needs to store the incoming
activations. Depending on the batch size, these activations can still take up a sizable amount of memory.

During the forward pass, activations are therefore offloaded when not needed. During the backward pass, they are loaded
back into VRAM only when needed.

## Memory Management

Efficient memory management plays a critical role in making this method possible. PyTorch employs a caching allocator
that tries to re-use allocations when creating or copying tensors. This works well in normal use cases where allocations
are very regular and always the same size. But when differently sized parts of a model are constantly moved around, this
leads to memory fragmentation and wasted space.

Therefore, OneTrainer implements two custom memory allocators, one for layer allocations, and one for activation
allocations. Memory is allocated by creating a one-dimensionally int8 tensor of the desired size. Allocating tensors
inside this block of memory first creates a slice of that memory, then uses a view into that slice to re-interpret it
into any data type and shape.

### Layer allocations

Layers in a model usually have a well defined execution order. This means we can know the allocation patterns ahead of
time, and it's possible to define an optimal allocation strategy. OneTrainer utilizes a ring buffer for allocations.
Tensors can be allocated in both directions depending on the execution direction of the model. Again, this is best
explained by an example. The following diagram shows the allocation pattern in VRAM for a 12 layer model with offloading
set to 50%, meaning 6 layers are offloaded at any time.

Each slot in the buffer is represented by a column, each row represents a point in time.

<details><summary>Allocation diagram during training</summary>

- `->` execution direction (forward or backward pass)
- `Number` the block index
- `*` executed block

```
->  *00* 01  02  03  04  05  (forward pass starts)
->   06 *01* 02  03  04  05
->   06  07 *02* 03  04  05
->   06  07  08 *03* 04  05
->   06  07  08  09 *04* 05
->   06  07  08  09  10 *05*
->  *06* 07  08  09  10  11
->   06 *07* 08  09  10  11
->   06  07 *08* 09  10  11
->   06  07  08 *09* 10  11
->   06  07  08  09 *10* 11
->   06  07  08  09  10 *11*
<-   06  07  08  09  10 *11* (backward pass starts)
<-   06  07  08  09 *10* 05
<-   06  07  08 *09* 04  05
<-   06  07 *08* 03  04  05
<-   06 *07* 02  03  04  05
<-  *06* 01  02  03  04  05
<-   00  01  02  03  04 *05*
<-   00  01  02  03 *04* 05
<-   00  01  02 *03* 04  05
<-   00  01 *02* 03  04  05
<-   00 *01* 02  03  04  05
<-  *00* 01  02  03  04  05
```

</details>

This allocation scheme works best if every block in a model is the same size. For modern models, this is not the case.
The same approach can still be used, we just need to allocate additional memory, and keep track of exact offsets in the
buffer instead of a slot number.

### Activation allocations

Activation allocation is simpler than layer allocation. During the forward pass, each tensor is added to the end of
allocated memory region. If the size of the region is exceeded, a new region is added to the end. After each forward
pass, all allocated buffers are combined into a single large memory buffer. With resolution bucketing, allocation sizes
can differ between each pass, which means some passes need more memory than would fit into the already allocated buffer.
This is solved by dynamically allocating more memory as needed.
