SpeedLoader: An I/O efficient scheme for heterogeneous and distributed LLM operation

Response to Question 1

Thank you for your question. We have prepared a detailed explanation of the one-shot hyperparameter tuning process below. Please review it and let us know if you have any questions, so we can integrate it into the camera-ready version if the paper is accepted.

The primary objective of the one-shot hyperparameter tuning process is to maximize high-bandwidth memory (HBM) and host memory utilization. Given a target context length, the tuner runs two iterations of the training script with different batch sizes, recording the maximum allocated host and device memory during these runs. Since memory consumption is highly linear with respect to the number of input tokens (as discussed in Section 5.1), the tuner calculates the memory increment per input entry. Using this information, it determines the maximum sub-batch size and the number of sub-batches that can be accommodated within the available device and host memory. This approach allows for efficient utilization of memory resources, ensuring optimal performance.

Response to Question 2

Thank you for your feedback and insightful advice. While SpeedLoader addresses optimization under restrictive computational resources, it also has broader applications. Even with ample hardware resources, the model states for training can exceed the capacity of a single device, making model parallelism essential in most training scenarios.

Recent advances like PagedAttention are primarily focused on optimizing inference rather than training. As of now, ZeRO and FSDP are still the leading paradigms for large language model training involving model offloading and sharding, which highlights SpeedLoader’s solid advantages in these contexts. For reference, we also conducted a trial with Paged Attention and offloading (Fig. S4). However, vLLM is not specifically optimized for offloaded inference, leading to a lower throughput than DeepSpeed.

Additionally, SpeedLoader is compatible with other recent innovations like Flash Attention. Following your suggestion, we integrated Flash Attention into our framework and observed a performance improvement of approximately 5%. We have conducted an ablation study testing the peak performances with this enhancement (Fig. S5). The updated results will be incorporated into the manuscript if the paper is accepted. Experiment details can be found in Author Rebuttal.

Response to Question 3 and Clarification on Weakness 4

We apologize for any confusion and thank you for your suggestion. While offloading is a significant factor in addressing HBM limitations, SpeedLoader also addresses challenges in sharded model training, such as those encountered in ZeRO and FSDP. Therefore, distribution is another major factor in alleviating HBM pressures.

In Figure 8, we present the per-device throughput rather than the overall throughput. This means that while the per-device throughput trend stabilizes after 16 GPUs, the overall training throughput continues to scale linearly with the number of GPUs. For example, training with 32 GPUs can achieve double the throughput compared to a 16-GPU setup. To prevent further confusion, we highlighted per-device nature of data on the y axis label in revised version (Fig. S6).

Therefore, steady per-device throughput does not imply a lack of scalability. This phenomenon occurs because more GPUs are physically distant from each other, leading to communication over lower bandwidth interconnections like Ethernet. Given these circumstances, maintaining linear scalability is actually a positive outcome.

We understand that our evaluation with up to 64 GPUs might seem limited compared to some large-scale training setups. Unfortunately, as a research institute, access to larger-scale resources is not always readily available or affordable. We will continue to explore opportunities for testing on a larger scale in the future.

Please let us know if there are any further concerns or if additional clarifications are needed.

Clarification on Weakness 1

Thank you for your valuable feedback. Below is a revised explanation of the design and implementation details for sub-batching and gradient accumulation in SpeedLoader:

Here we propose SpeedLoader, which processes multiple sub-batches, collectively referred to as an effective batch, with only two model loading and gradient synchronization (Fig. 1, right). For any incoming batches shaped ($N,l$), they are evenly split to sub-batches along batch size axis ($n_{sb},n,l$), where $N,n_{sb},n,$ and $l$ are effective batch size, number of sub-batches, sub-batch size and sequence length, respectively. …

…

To address this issue, SpeedLoader redesigned the computation graph. We re-route … switching activations. In this way, each module computes and accumulates the gradient with respect to input and loss of every individual sub-batch. This approach enables the automatic differentiation mechanism to function correctly without causing excessive I/O. This method achieves gradient accumulation with an altered computation graph, which has significantly reduced …

We will integrate these details into the camera-ready version of the paper upon acceptance. If there are any further clarifications needed, please feel free to ask during the discussion.

Comments on minor issues

Thank you for pointing out the minor issues. Regarding the confusion between "flops" and "FLOPs," we apologize for not clarifying the difference. The lowercase "flops" refers to "floating point operations," while the capitalized "FLOPs" refers to "floating point operations per second." We will include an explanation of these abbreviations, along with other necessary fixes, in the camera-ready version if the paper is accepted.

Dear Reviewer,

Thank you for your thorough review of our manuscript. We appreciate the opportunity to clarify and address your concerns. Below is a summary of the key points in our rebuttal:

• We have revised the manuscript to clarify the tuning and sub-batching processes.
• Our method achieves the same computational outcomes as gradient accumulation, but with a different order of operations.
• The primary innovation of SpeedLoader lies in its enhanced efficiency for distributed and offloaded training scenarios. This optimization is not limited to offloading; both ZeRO and FSDP without offload can also benefit significantly.
• While PagedAttention was tested, it was not optimized for offloading, resulting in inference throughput lower than DeepSpeed. Additionally, it is not suitable for training.
• There may have been some miscommunication regarding Figure 8. The figure depicts per-device throughput, not overall throughput. As a result, you can still observe doubled throughput on 32 GPUs compared to 16.
• Additional experiments have been added to examine compatibility and further enhance performance with SOTA techniques. Please refer to the Author Rebuttal for details.

Thank you for your time and effort in reviewing our paper. We look forward to your feedback and hope for a productive discussion.

Best regards

Dear Reviewer,

Thank you for your time and effort in reviewing manuscripts for NeurIPS this year. We greatly appreciate your contributions.

I am writing to respectfully request a reassessment of our manuscript in light of the clarifications, additions, and results provided in our rebuttal. We believe that our response, along with the global “Author Rebuttal,” addresses the concerns you raised and better clarified some points that might have cause misunderstanding, (e.g., per-device throughput).

Once again, thank you for your invaluable contributions to making NeurIPS 2024 a success.

Best regards

Response to Question

Thank you for your valuable feedback. While the diagrams may appear similar, SpeedLoader is fundamentally different from FlexGen. FlexGen is a significant work that substantially improves the inference throughput of LLMs by efficiently managing activation and cache tensors. In contrast, our work focuses on the training scenario, which requires carefully designed handling of gradients and the computation graph. For simplicity, Fig. 1 only illustrates the forward propagation computation to highlight the throughput advantages. The detailed mechanism of SpeedLoader, including its differences from FlexGen, is demonstrated and explained in Fig. 2.

Additional Comments

Thank you for your insightful comment on other parallelisms. Indeed, both tensor and pipeline parallelism (TP and PP) are prevalent model parallelisms. However, different parallelisms come with distinct challenges. For instance, despite its minimal communication overhead, PP often suffers from bubbles in the computation schedule. On the other hand, tensor parallelism involves heavy peer communication to exchange results. Like FlexGen, SpeedLoader is inherently compatible with these parallelisms. But the primary goal of SpeedLoader is to reduce the heavy burden associated with sharded and offloaded model training, rather than to optimize or alleviate the particular issues of TP and PP.

Response to Question 1

Thank you very much for the thoughtful insights. The key concept behind SpeedLoader is to achieve equivalent computation of a large batch with a small memory consumption. An important feature of SpeedLoader is its management of activations. For instance, processing a batch of 256 sequences, each with a length of 2k, can result in activation data occupying 4GB between every two layers in the LLaMA-2-7B model, not including the additional space needed for tensors saved for backpropagation. SpeedLoader mitigates this issue by breaking down the workload into sub-batches of 16 samples each, ensuring that only 256MB of HBM is in active use at any given time, while unused tensors are temporarily moved out of the GPU. Therefore, the computation can be done in a constant small HBM consumption.

Response to Question 2

We are sorry for the confusion and thanks for pointing it out. The captions might have been not informative in this submission. We have modified it in revised version as below:

Figure 5 Profiling Results. Panels (a-c) were conducted on a single device with offloading; panels (d-e) were distributed without offloading. (a) DeepSpeed (w/ offloading); (b) SpeedLoader (unsaturated-computing, w/offloading); (c) SpeedLoader (saturated-computing, w/ offloading); (d) DeepSpeed (w/o offloading); (e) SpeedLoader (w/o offloading)

The input for unsaturated case has a sub-batch size of 32 and effective batch size of 128; Saturated case has 256, and 2048 for sub-batch size and effective batch size, respectively. Each sample in the batches has a context length of 256.

We implemented meticulous overlapping schedules for communication and computation. Our profiling confirmed that these schedules performed as expected (see Fig. S1 in Author Rebuttal). Figure 5 does not reflect the overlapping schedules, as the pie chart shows the function calling wall time.

The left panels (a, b, c) and the right panels (d, e) were profiled under different settings. The memcpy in the left panels corresponds to the all_gather and reduce_scatter operations in the right panels as part of the communication payload. The 10.4% all_gather accounts for only a portion of the total communication in sharded training, with an additional 23.5% spent on reduce_scatter for gradients during backpropagation, contributing to the overall communication operations.

Dear Reviewer,

Thank you for your hard work in reviewing the manuscript and for your valuable insights. We hope that our rebuttal addresses your concerns. Below is a brief summary:

• SpeedLoader can compute over a large batch with significantly lower memory requirements.
• All_gather is not the only component of communication; reduce_scatter also accounts for a significant portion. Memcpy (in offloaded training) corresponds to the two operations (in non-offloaded training) together. SpeedLoader significantly reduces the reduce_scatter ratio in total wall time.
• Sequence length has a marginal impact on MFU. Regardless of sequence length, SpeedLoader enables a larger GPU batch size, allowing for more efficient computation.
• Additional experiments have been added to examine compatibility and further enhance performance with SOTA techniques. Please refer to the Author Rebuttal for details.

Thank you for your time and effort in reviewing our paper. We look forward to your feedback and hope for a productive discussion.

Best regards

Dear Reviewer,

We sincerely appreciate your valuable contributions to NeurIPS this year and thank you for your time and insights.

We kindly request a re-evaluation of our work based on the detailed clarifications and results provided in our rebuttal. We have addressed your concerns with additional experiments and texts, and hope these clarifications will resolve any misunderstandings, particularly regarding communication breakdowns.

Once again, we extend our gratitude for your hard work in making NeurIPS a success this year.

Best regards,

Response to Question 1

Thank you very much for the kind and insightful review. Currently, 7B models at bfloat16 precision is very challenging to be trained on a 40GB A100. The model and optimizer states alone can take up to 56GB memory. One of the methods to deal with it is to shard the model and optimizer across multiple GPUs, as we demonstrated in Fig. 5d,e and Fig. 6 right panel.

Following the advice, we also tested the training MFU of LLaMA-2-7B without sharding and offloading in an 80GB A100 GPU, which is 60.86% MFU under a batch size of 8. We have found that our proposed method (54.59% on the same machine) could achieve a similar MFU as the ideal scenario.

Response to Question 2

Sorry for the confusion and thank you very much for this valuable feedback. In our experiments, we compared the performance of SpeedLoader and DeepSpeed under the same maximum token number, which is constrained by the capacity of HBM. This ensures that our benchmarks reflect the best performance for each method. Specifically, we tested 128 entries with a sequence length of 256 per sub-batch (totaling 32k tokens) for SpeedLoader, and 32 entries with a sequence length of 256 per batch (totaling 8k tokens) for DeepSpeed. Given the same HBM size and context length, our memory-efficient approach allows for more effective computation. Given any sequence length, SpeedLoader can support a larger on-device batch size (sub-batch size) than original DeepSpeed, enabling more efficient kernel execution.

To address concerns about the impact of sequence length, we conducted additional tests varying the sequence length on a single A100 40GB with offloading using the LLaMA-2-7B model (Fig. S2 in Author Rebuttal). The results showed that increasing sequence length does not significantly impact both SpeedLoader and DeepSpeed performance, with a variance of no more than 5%. This variation is marginal compared to the 3.29x speedup achieved by SpeedLoader, demonstrating its efficiency and effectiveness.

Response to Question 3

Thank you for your suggestion. We highly appreciate your suggestion on the importance of clearly outlining the limitations. We add on a Limitation section with the contents as below and this section will be integrated into the camera-ready version of the paper if it is accepted.

Limitation

SpeedLoader provides distributed deep learning settings with an I/O efficient option. By fully utilizing the available host memory, SpeedLoader can significantly reduce the redundant model state exchange among the device, peer and host memory. However, this approach may make host memory a new bottleneck for efficient model operation. Additionally, the implementation of memory allocation in PyTorch only allows sizes that are powers of two, which can lead to inefficient use of memory. For example, if a tensor requires 33 GB, PyTorch allocates 64 GB, potentially leading to out-of-memory issues. Currently, there is no explicit option to modify this behaviour for CPU memory allocation.

Furthermore, while SpeedLoader is theoretically compatible with other parallelism strategies like tensor parallelism and pipeline parallelism, it has not yet been tested in conjunction with these methods. The potential benefits of SpeedLoader in larger-scale hybrid parallel settings remain to be explored. While data parallelisms like FSDP and ZeRO have low communication overhead in exchanging activations, tensor parallelism has lower communication overhead in exchanging parameters. Thus, exploring the dynamic trade-off with speedloader between these parallelisms is a valuable area for future research.

Additional Comments

Following your advice on compatibility across generations, to explore the compatibility of SpeedLoader on other platforms, we examined single-device speedups of SpeedLoader on NVIDIA V100S, RTX A6000, and H100 with PCIe Gen 3 to Gen 5 respectively (Fig. S3). The results are 1.37x, 3.56x, and 3.01x, respectively. This proves that our proposed method can have positive effects on different platforms with mixed I/O capability.