Reading MLP-Offload: Muti-Level, Multi-Path Offloading for LLM Pre-training to Break the GPU Memory Wall

Overall Introduction

Maurya et al. came up with MLP-Offload to speed up the pre-training process of LLMs, at the same time, break the GPU memory wall. Compared to state-of-the-art method, DeepSpeed ZeRO-3, their way realizes 2.5x speedup in iterations.

Background and Insights

Their MLP-Offload method is aimed at solving bottleneck caused by GPU memory limits in pre-training period. Details are as follows:

GPU Memory Wall: Growth of the size of LLMs is much faster than growth of GPU memory. At the same time, model parameters, gradient, optimizer states and activations all need to be stored in memory, making requirement of GPU memory several times that of parameter size when pre-training. All of the reasons above form what is called “GPU Memory Wall”.
Limits of Current Parallel and Offloading Technologies: Althrough there are several parallel strategies like data parallelism, pipeline parallelism and tensor parallelism and they can expand memory, communication overheads are large. ZeRO redundancy elimination can reduce redundacy by spliting optimizer states or gradients, but needs offload some into CPU or SSD. When training, the speed is acceptable if it in the host memory(CPU/GPU), but if the data size exceeds the memory, it must be offloaded into third-level storage such as SSD and PFS which however have much lower bandwidth than GPU or CPU memory. Meanwhile, update phase will read and write the optimizer states frequently, making I/O the main bottleneck of the whole pre-training.
What’s Worse —— Update Dominates Pre-training: In iterations of pre-training, most of the time is spent in updating, and most of the updating time is spent in I/O.
I/O of Third-Level Storage PFS isn’t Fully Utilized: Remote storage like PFS is usually in an idle state when training, but sometimes it may provide more throughputs than NVMe can, which makes it feasible to supplement local NVMe through PFS and form multi-path I/O.
I/O Bandwidth under Concurrency during Offloading: Every work process launches fetching/flushing operation for subgroups and thus making conflicts and competition among these processes in NVMe I/O. However, bandwidth of NVMe is limited.
Active Subgroups are Limited: Limited by host memory, there are only three subgroups in host memory at a time – one for prefetching, one for updating, and one for flushing back to SSD.
The Sampling and Offloading of Gradients are Redundant: Traditinal methods flush gradients to host memory and from there to third-level storage in FP32 format. The FP32 gradients and optimizer states are flushed back again when updating. Noticing that transformation of FP16 to FP32 is fast in modern CPU, we can only store FP16 gradients in CPU and transform them into FP32 when needed in update phase, avoiding redundant fetching and flushing of FP32 gradients.

Contributions and Design Principles

Firstly, they design unified multi-level, multi-path asynchronous offloading using virtual tiers. they conbined SSD, PFS and object stores into a virtual, third-level tier that can be used to offload model subgroups from the host memory. It allocates the number of subgroups according to the proportion of the I/O bandwidth of each type of storage. Formulas are as follows:

where is the number of storage tiers, is bandwidth of i-th storage tier, represents the total number of subgroups in one worker process, and represents how many subgroups the i-th storage tier should be allocated to.

It’ worth noting that will be adjusted after further experiments to fit bandwidth fluctuation of shared storage like PFS.

Secondly, the design of optimized virtual tier concurrency control for multi-path I/O allows each worker process to apply I/O load balancing on alternative storage independently of other workers. Their concurrency control strategy at the mode level is allowing only one worker process on each compute node to access a given alternative storage at a time. At the same time, the other worker processes are free to compute updates for prefetched groups using all the available CPU cores, and use other alternative storage(s) in parallel, resulting in a natural interleaving that achieves global I/O load balancing.

Thirdly, they have a cache-friendly ordering of model subgroup processing method. As mentioned above, memory limit makes only three subgroup activated in a time. To avoid memory trashing, this method start processing the subgroups in increasing order of IDs. After the update phase, the subgroups with a high ID will reside on fast (cache) tiers, while the subgroups with a low ID will reside on the slow tiers as they were evicted to slow tiers to make room for the subgroups with higher ID in the host memory. Then, in the second iteration, unlike state-of-the-art approaches, it reverses the order of processing the subgroups. This results in a significant acceleration of the update phase because a large number of subgroups with a high ID are already present on the fast tiers, thereby reducing I/O overheads due to memory thrashing.

Fourthly, delayed in-place mixed-precision gradient conversion during the update phase is an approach which stores FP16 gradients in host memory – which anyway needs to reserve enough room for the FP16 gradients of all subgroups to enable gradient accumulation, and during the update phase fetches the subgroup and converts the FP16 gradients to their FP32 variant on-the-fly. After updating, new FP16 gradients are still in host memory and used FP32 gradients are thrown away. In this approach, only optimizer states and parameters are stored in SSD, gradients are temporary in memory.

Performance Evaluation

Their baseline is DeepSpeed ZeRO-3, and used following two testbeds:

Models used for evaluations are as follows:

Their evaluation results can be summarized as following points: - MLP-Offload speeds up update phase 2.4x, and end-to-end iteration time 2.7x. - Update throughput of MLP-Offload is 1.8x to 2.4x of DeepSpeed, proving that bottleneck is I/O limits. - I/O throughput has almost doubled by using third-level virtual tiers. - Host memory trashing is reduced according to the experiments. - MLP has a greater advantage in scaling. - MLP-Offloading is more cost-effective than DeepSpeed. - In ablation experiments, semipath provides 1.6x speedup and multi-path finally provides 2.5x speedup.

Conclusion

Maurya et al. and his team mainly contributes that they use multi-level, multi-path asynchronous offloading method to break I/O bottleneck, and harvest success compared to previous approaches. Its expansion of DeepSeed ZeRO-3 leads to faster training in 40B to 280B models’ testing, which is also more cost-effective. Also, it reduces redundant I/O of gradients flusing in backward and update phase. Their insights about I/O are worth learning, which provides potential research perspectives. And in the future, the authors plan to extend this technology to CXL memory pools and other training frameworks, and further investigate bandwidth fluctuation mitigation strategies for the globally shared storage layer under I/O contention.