## 4 Experimental Testing
### Theoretical Analysis
#### 4.2.2 DeepSpeed ZeRO-Offload  
The ZeRO-offload [50], when used with 'DeepSpeed' [17; 52], enables parallel processing of training components across multiple GPU streams including optimiser states, gradients, and model weights. This provides significant benefits for training models on multiple GPUs, such as increasing training speed and minimizing resource utilisation. Additionally, ZeRO-offload assists in balancing the load between GPUs during the model training process ensuring that each GPU is utilised effectively and does not experience slower performance than other GPUs.  
ZeRO-offload collaborates with ZeRO to extend DL training across multiple GPUs. ZeRO comprises three stages, ZeRO-1, ZeRO-2, and ZeRO-3, each handling different aspects of model partitioning, including optimiser states, gradients, and parameters. While ZeRO-1 partitions only optimiser states, ZeRO-2 partitions gradients alongside optimiser states and ZeRO-3 partitions all model states. ZeRO-offload synergises with ZeRO-2. In ZeRO-2, every GPU retains a replica of all parameters but updates only its designated portion during each training step. As a result, each GPU stores only the optimiser states and gradients necessary for its update. Following the update, each GPU transmits its updated parameter subset to all other GPUs through an all-gather communication collective. The computation and communication schedule of ZeRO-2 are outlined as follows: During the forward pass, each GPU computes loss concerning a distinct mini-batch. During backward propagation, gradients are computed and then averaged using a reduce operator at the GPU (s) responsible for the gradient or its segment. Subsequently, each GPU updates its parameter subset and optimiser states using the averaged gradients. Finally, an all-gather operation is performed to obtain the remaining parameter updates computed on the other GPUs.  
In summary, ZeRO-offload is an optimisation mechanism for training models on multiple GPUs that increases training speed and optimises resource utilisation. Moreover, with the added offloading mechanism, the CPU can load additional amounts of gradients, optimiser states, or parameters to reduce the burden on the GPU and free up VRAM.  
As discussed in this paper, in our experimental testing and evaluation of our _Expert-B_ model, following the use of LoRA the model training time is now only 4 h/1 epoch, with 39/40 GB of VRAM NVIDIA A100 serving the training process. Following the use of 'DeepSpeed' and the offload mechanism for optimiser states, the current training configuration with a batch size of 1 now only takes up 36/40 GB of VRAM. As there is still 4 GB of VRAM available, the batch size can be increased to 2. When increasing the batch size to 2, the training time for one epoch is reduced to 3 h/1 epoch on 1 NVIDIA A100s. We have conducted tests using the NVIDIA A100 40 GB GPU. If the VRAM is not fully utilised during training, it suggests that the hardware is not optimised to its full potential. Specifically, employing DeepSpeed with a batch_size of 1 can lower VRAMusage, indicating inefficient GPU utilisation. Consequently, we raised the batch_size to 2, allowing for the optimal utilisation of the GPU's capabilities.