SpeedLoader: An I/O efficient scheme for heterogeneous and distributed LLM operation
Authors: Yiqi Zhang
Yang You
Venue: NeurIPS 2024 poster (2024)
Keywords: Heterogenous Computing
Large Language Model
ZeRO
FSDP
Offload
Primary area: infrastructure
TL;DR: SpeedLoader can compute multiple batches within one Forward-Backward pass.

Abstract: With the surging growth of model parameters, foundation models pose unprecedented challenges to traditional computational infrastructures. These large models inherently require substantial accelerator memory to accommodate massive tensors during pre-training, fine-tuning, and even inference stages, making it even more challenging to deploy a model with restricted computational resources. Given this challenge, distribution and offloading the model states are two major solutions. Partitioning the required states to participating workers, and storing them in lower speed media, such as host DRAM and block devices, largely alleviate the accelerator memory pressure. However, the prohibitive costs of tensor communication render it a theoretically plausible yet practically inefficient solution. Previous efforts to improve efficiency include maximizing rematerialization and employing chunk-based tensor management to reduce host-device communication. Despite these efforts, the reported training throughput only achieves 36.54% of model FLOPs utilization (MFUs), still not comparable to full on-device training. In this work, we redesign the data flow of heterogeneous hardware and sharded model training to minimize the excessive communication overhead. Our proposed scheme significantly enhances training and inference throughput of large language models under restrictive computational resources. We confirmed a large leap in effective compute time by looking into the kernel-level runtime behavior of our trials, where the MFUs can achieve up to 51%. Compared to the state-of-the-art approach, our framework robustly achieves remarkable speedups from 3x to 30x in multiple distributed heterogeneous training setups and inference speedups of 1.5x to 2.35x without compromising arithmetic precision.

SpeedLoader: An I/O efficient scheme for heterogeneous and distributed LLM operation
Yiqi Zhang Institute of Data Science yiqi. zhanguu..n.c mus edu
Yang You School of Computing youyθcomp nus edu sg
National University of Singapore Singapore, 119077
Abstract
With the surging growth of model parameters, foundation models pose unprece- dented challenges to traditional computational infrastructures. These large models inherently require substantial accelerator memory to accommodate massive tensors during pre-training, fine-tuning, and even inference stages, making even more challenging to deploy model with restricted computational resources. Given this challenge, distribution and offloading the model states are two major solu- tions. Partitioning the required states to participating workers, and storing them in lower speed media, such as host DRAM and block devices, largely alleviate the accelerator memory pressure. However, the prohibitive costs tensor com- munication render it theoretically plausible yet practically inefficient solution. Previous efforts to improve efficiency include maximizing rematerialization and employing chunk-based tensor management to reduce host-device communication. Despite these efforts, the reported training throughput only achieves 36.54% of model FLOOPs utilization (MFUs), still not comparable to full on-device training. In this work, we redesign the data flow of heterogencous hardware and sharded model training to minimize the excessive communication overhead. Our proposed scheme significantly enhances training and inference throughput of large language models under restrictive computational resources. We confirmed a large leap effective compute time by looking into the kernel-level runtime behavior of trials, where the MFUs achieve up 51 31%. Compared the cate-of-the-aat approach, our framework robustly achieves remarkable speedups from 3x to 30x in multiple distributed heterogeneous training setups and inference speedups of 1 1.5x to 2.35x without compromising arithmetic precision.
Introduction
The trend of increasing transformer-based model sizes marks a paradigm shift natural language processing deep neural networks. Research on model scales suggests that expanding transformer sizes significantly benefits aspects such as performance, universality, and transferability 222. In accordance with scaling laws [I] large language models (LLMs) have surged in recent years. For instance, models have grown from the 175 billion parameters in GPT-3 to over 400 billion parameters
in LLaMA-3 B 4. the size of LLMs underwent a substantial leap. However, the continually growing size these models places mounting pressure on computational infrastructures. For example, training
38th Conference Neural Information Processing Systems (NeurIPS 2024).

---

GPT-3 175B requires approximately 3.14 X 1023 floating point operations (flops equivalent to around 280,000 GPU hours on an NVIDIA A100 at its peak performance, which is not pragmatic without large scale parallelism.
The scalability of LLMS heavily depends on parallelism and distribution, with significant constraint being the heavy memory footprint. Various techniques, such as tensor parallelism [6, pipeline parallelism[7 8 9 and quantization [O0 [T], can help alleviate this pressure. Model parallelism, in particular, involves sharding the model states (parameters, optimizer states, gradients) across accelerators to overcome limitations imposed by finite on-device high-bandwidth memories (HBMs) [2]. Unfortunately, HBMs remain searce and costly, imposing significant infrastructural constraints on training and using larger models. For example, a half salf-precision 1.8 trillion parameter model demands over 180 units of 80GB GPUs to hold the model and optimizer states, not accounting for the memory needed for essential activations. Relying solely HBMs for state storage is becoming impractical with current accelerator technology models continue to scale.
Offloading has emerged as a flexible solution for both training and inference with LLMs across various infrastructure scales. It allows model states to be stored on larger media like host DRAMs and NVMe drives, loading them to the accelerator only when required for computation [ However, the merits of offloading come the expense of even higher 1/O overhead. Model weights must frequently shuttle between accelerator and host memory during computation. Additionally, during inference, large volumes of key value caches (KV Cache) [16], need to be moved to external storage and subsequently brought back to the GPU for the next token. The overhead caused by data movement often surpasses computation time due to bandwidth bottlenecks between accelerator, CPU, and NVMe storage, resulting in high task latency.
14 15].
To minimize redundant communication and maximize hardware utilization, we propose SpeedLoader, an highly optimized scheme to eliminate I/O redundancies through computation rescheduling during LLM operation with model sharding and offloading. By meticulously feeding and offloading acti- vations, LodedLoader computes multiple batches with only two full-model loading. We evaluated SpeedLoader's performance with LLaMA-2 and OPT [7] [8] at different sizes. Results showed that SpeedLooder can robustly achieve a training speedup of 3.5x to 30x and over 50% model FLOOPs utilization (MFU) on multiple platforms compared to state-of- -ftthe-arr approaches.
The major contributions of this work are: (1) We proposed a compute strategy that minimizes peer and device-host I/O in heterogeneous and sharded LLM training; (2) We implemented high-efficiency tensor exchange manager that transparently shuttles tensors between device and host, minimizing fragmentation and redundancy; (3) Our optimized tensor management enables higher inference efficiency than previously state-of-the-aar (SOTA) approaches.
Related Work
LLMs are resource-intensive and typically require specialized strategies to alleviate hardware stress. One common approach is to optimize distribution and parallelism across multi-GPU clusters. For instance, Megatron-LMM [5] leverages advanced model parallelism techniques to achieve scalability, and Zero Redundancy Optimizer (ZeRO) partitions and distributes all model states (i.e., parameters, gradients, optimizer states) among GPUs, recollecting them only when the layer needs to be computed [1]. Similarly, Fully Sharded Data Parallelism (FSDP) shards both model parameters and input data across multiple GPUs 19]. These techniques usually gather full parameters only when needed and reduce scatter the newly computed gradients across peers, significantly alleviating memory pressure on the accelerators while introducing considerable communication overhead.
While distribution frameworks still require substantial hardware, offloading approaches address memory shortages in a different way. Work including ZeRO-Offoad, Capuchin, and SuperNeurons [13 20 21] enables large model training by offloading data and computation to host RAM, mitigating the memory and computational demands of training large models with fewer GPUs. ZeRO- Infinity
A list of abbreviations in this paper can be found in Tab.

---

14 extends this by offloading more data to high-speed block devices like NVMes, overlapping computation and communication for higher bandwidth utilization. During inference, offloading can significantly reduce the computing resources required for pre-trained models of unprecedented scale. Unlike early works, PatrickStar [22] organizes the model data memory chunks and dynamically distributes them in the heterogeneous memory. FlexGen [5] offers wider range of batch size choices with an efficient offloading strategy, and thus significantly increases maximum throughput, demonstrating that high-throughput generative inference of 175 billion parameter LLM is feasible support prefetching parameters after the initial forwardbbackward pass, and the granularity of overlapping can be precisely controlled (23). As an update, ZeRO++ implements hierarchical weight partition, which maintains full model copy on each node largely reduce inter-node tensor exchange [24]
with single GPU. Current offloading solutions 13 22
In addition these primary solutions, other approaches have been explored. Rematerialization techniques [25] 26 27 recomputes activations on-the-fly during backpropagation rather than storing all intermediate activations in memory. Model compression methods such as quantization are also commonly used to reduce memory requirements 10] [1]. However, compression usually comes at the cost of numerical inaccuracy and retraining of the model.
SpeedLoadee
3.1 Overview
In canonical solutions (Fig. [ left), each mini-batch processed sequentially, requiring the full model be loaded twice per batch. In distributed settings, the accelerator must read local partitions of model states from host DRAM and gather other partitions from peer ranks. Additionally, to properly partition the accumulated gradients, reduce-scatter operations across all ranks are inevitable, introducing communication equivalent to one copy of the model weight during each backward propagation.
Outputs
Outputs
Layer
Hest
(D)
Ussed
Layer
Inputs
Inputs
Figure 1: Mechanism comparison between canonical approach and SpeedLoader. Left, canonical approach ① load the whole model every time 2 loads single batch; Right, SpeedLoader feed every layer with multiple batches and exchange activations to and from host memory (0). Therefore, for one full model loading(@). (2), speedLoodeer process multiple batches.
Here propose LpdedLoader, which processes multiple sub- batches, collectively referred to effective batch (i.e., number of sub-batehe sub-batches), with only two model loading and gradient synchroniza- tion (Fig I right). For any incoming batches shaped (N, (), they are evenly split to sub-batches along batch size axis (n.b, n. 1), where b,n, and I are effective batch size, number of sub-batches, sub-ban batch size and sequence length, respectively. Each time a layer is loaded, new imputs are fed into the model while the activations from the previous sub-batch are offloaded memory spikes. After processing all sub-batches with the current layer, the accelerator relcases the current layer and proceeds to the next layer. Subsequently, the previously saved activations reloaded. This method achieves gradient accumulation with an altered computation graph, with
host buffer to prevent

---

significantly reduced frequency of model state communication. A quantitative analysis is discussed in Sec. A.3
The backward propagation process also needs be redesigned. By default, each sub-batch generates individual loss, and the automatic differentiation package conducts backward propagation on each sub-batch sequentially. This process loads the model and partitions gradients once per sub-batch, leaving the prefetching of associated gradients and activations out of our control. address this issue, SpeedLoader redesigned the computation graph. We re-route the computation graph by connecting activations from the same layer with a no-op function. This modification allows hackward propagation to trace back in the same manner as forward propagation. Additionally, the connecting function buffers and passes correct gradients when switching activations, ensuring scheme correctness. this way, each module computes and accumulates the gradient with respect input and loss of every individual sub-batch. This approach enables the automatic differentiation mechanism function correctly without causing excessive I/O.
3.2 Tensor Exchange Manager
The core mechanism of SpeedL lies in novel scheduling of tensor computation and exchange, particularly in managing and minimizing communication overheads. Heavy data exchanges can lead memory fragmentation, resulting in memory waste and reduced hardware efficiency. To address this, carefully examine the details and propose simple yet effective technique optimize performance.
a2
D2
M (7)a,
va,
(?),,
(V/c,
VC2
Compute Ho-pp 1D
a₂
by
C2
Desior
a,
vag
b,
rbg
C₁
VC2
Adtivation
Activition Laayr
va
7b,
VCc
a,
b,
Forward Propagation
Backward Propagation
Figure 2: Device-host communication breakdown of SpeedLoader.
Transparent data exchange. The implementation of SpeedLoader provides plug-n-play function wrapper. Users can wrap the forward function of each layer in the model to iteratively load and compute multiple sub-batcher in one explicit forward call. Taking closer look at the wrapper demonstrated in Fig. 2 and Algorithm 1, performs the following activities before and after the forward/backward calls:
Here, pinned_ x is relocalce host memory for activation and gradient storage. Before each forward pass, the activation of previous sub-batch is offloaded while prefetching the input of the next sub-batch simultaneously. This communication is also overlapped with the computation of current sub-batch using two separate CUDA streams in our implementation. During each backward propagation, the gradient of the previous sub-batch activation is offloaded, and the next required activation and gradient are prefetched. Similar to the forward propagation, these two processes are overlapped with the re-computation of checkpointing and computation of gradients. In inference scenarios, the communication is much simpler: assuming KV caching is enabled, exchange the KV cache instead of the entire activation.
Fragmentation-free Memory Pool. Without prior information the model, PyTorch can only adopt an on-demand approach to allocate host memory for offloaded tensors However, allocating lengthy and blocking operation. Meanwhile,
extensive continuous memory in this manner excessive memory allocations can result in significant memory fragmentation, and even highly optimized LLM serving systems suffer from heavy fragmentation 22) Given the configuration of transformer-based LLM. is possible to pre-allocate reusable memory pool. In this study,

---

allocate a minimal page-locked memory pool for caching offloaded activations or KV caches based on the given hyperparameters of the model.
token 0
sample
Reshape
Sample &
Token 1,
Figure Reorganizing buffer memory to ensure continuity.
Additionally, this pinned memory pool is highly optimized for sparse KV caching under inference workloads. KV cache within the same sub-batch is preserved for each layer throughout the generative steps to maintain computational consistency. In our approach, only the KV state of the new token is written to the host memory. ity of the host pinned memory buffer, the buffer pool is organized in the format of (n (n_sub batches, layer, sequence_lengie -1). Given the augmenting nature of se- quence length through generation steps, it is permuted to the third dimension from its original position in the sixth dimension (Fig Therefore, contiguous memory chunk is available for asynchronous bulk copy for all tokens in the sub-batch, eliminating the need for blocking operations additional loops.
Algorithm 1 Pssudd-code for tensor exchange.
< embedding( iiput ids)
for i = 1 to len batches) do
Offload x to pinned
pinned_x[ldd][= x[lid]
<- buffer
register_hook(4
Fetch pinned_
siinnddx(liddddd x(lid - T2 to buffer
4 layer(x)
ensure the continu-
end for output_oogit <- output_embeddingg.
procedure BACKWARD HOOCEKAWAAD_OO((_,,_____ i, lid)
Offload z.grad pinned baannddx[lid___ - 1]
x <- act_buffer
x.-grad grad buffer
Fetch pinned_ siieed_((iddd_e 1] 1] to act act__aaffee buffer Fetch [onned_x[[id]]] 1] to grad_ _uffer backward(zz)
end procedure
3.3 Hyperparameter Tuning Strategies and Communication Analysis
Hyperparameter selection is critical aspect of both the training and inference phases of LLMs. Our proposed method expands the search space for hyperparameter tuning, highlighting the importance of a swift tuning strategy. Based on observations in Section 5.1 we have developed one-shot hyperpa- rameter tuning strategy that not only addresses these new dimensions (i.e., sub-batch size, effective batch size and number of on-device layers) but is also compatible with existing framework-providee tuning tools. Details those hyperparameters can be found appendix A.2 Our observations show that the performance of LpeedLoader significantly benefits from the exhaustive usage of available memory, including accelerator HBM and host DRAM.
A.2
To achieve optimal training and inference throughput, we aim to minimize data movement and maximizing arithmetic density iterating through the same amount of training data. In this context, we will discuss scenarios with the assumption that gradient checkpointing is enabled for training. Given P, A, N refering to parameter, activation size and number of sub-batches, SpeedI.oader Loader has 5 3P and 3P local and remote communication, compared both 3NP for method (see quantitative details in appendix A.3). As result, the I/O advantage of Speedl Loader unmodified ZeRO-Offlon can be asymptotically enhanced by increasing effective batch size of SpeedLoader.
The one-shot hyperparameter tuning primarily aims to maximize the utilization of HBMs and host memory. Given 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 the number of input tokens (as discussed in Section 5.1). the tuner calculates the memory increment per input entry. Using this information,
5

---

determines the maximum sub- bub-batch size and the number of wub-batches that can be accommodated within the available device and host memory. This approach allows for efficient utilization of memory resources, ensuring optimal performance.
4 Evaluation
Implementation. For any transformer-famil model, we provide wrapper function that iteratively feeds the model with prefetched activations. This approach allows the model to compute and offload activations of several sub- bab-batches without exceeding device memory limits. The implementation also includes utilities that automatically tune hyperrarameeers. Our current implementations support [8]. Other transformer-base models can be integrated with minimal effort. For stringent comparison, we benchmarked the -LLMA-2--7B, 13B, and 70B models. To establish fair comparison with FlexGen[[5], we also benchmarked OPT-6.7B 30B, and 175B. For optimal offloading and scalability, Speedl Loader is implemented based on DeepSpeed ZeRO++ with hicrarchical weight partitioning (24). Additionally, 16-bit brain floating-poins and half precisions are employed to minimize memory pressure and maximally utilize tensor cores on NVIDIA GPUs.
both OPT and LLaMA-2 [7][8]
Platform specifications. Our experiments were performed VMs from Infrastructure as Service (laaS) provider and computing nodes from high-performance cluster (HPC). Our benchmarks were conducted on VMs with 16 NVIDIA A100-40GB GPUs with NVLink. Each VM is equipped with 96 cores of vCPU and 1360GB RAM. Functionality tests were conducted on HPC nodes with NVIDIA A A100-40GB GPUs and HPE Slingshot Interconnection in Dragonfly topology. Platform specifications can be found in Tab.
Table 1: Platform specifications for each experiment
Experiments
Internode Connection
Intranode Connection
Accelerator
Functionality Benchmark
HPE Slingshot
100Gbps Ethernet
PCIe Gen 4
NVIDIA A100-40GB
Profiling Scalability
100Gbps Ethernet
100Gbps Ethernet
PCIe Gen 5
NVIDIA H100-96GB
Compatibility
PCIe Gen 3
N/A
NVIDIA V100S-32GB
PCIe Gen 4
NVIDIA A6000
Results
5.1 Impacts of Hyperparameters
We first analyzed the impact of changing hyperparameters (Fig with model with LLaMA-2-7B on single NVIDIA A100 GPU.
Linear resource usage. Our analysis indicates that both HBM and host DRAM usage exhibit a linear pattern influenced by major hyperparameters. Using ordinary least squares regression, we found a highly linear relationship (R2 > 0.999) between sub-batch size/number of on-device layers and peak GPU memory usage, as well as total tokens and host residence memory usage, as shown in Fig. 4 Intriguingly, the tensor exchange memory pool allocates more page-locked memory than expected. This is rooted in the memory allocation behavior of PyTorch CPU tensors, which automatically rounds up the reserved memory to the nearest power of 2. Despite this linearity and predictable host memory allocation, predicting memory consumption on the accelerator remains challenging and can vary significantly across different models. To maximize the generalizability of our proposed approach, we decided to adaptively fit the linear resource model on the fly.
Asymptotic compute efficiency. We identified that the MFU is strongly affected by the hyperparam- eters including sub-batch size, number of sub-batches, and number of on-device layers. Increasing

---

Throughput (tokens/S))
2500
5000
2000
GPU Memory (GB)
0.25
1
16
200
Hozt Merorry (GB)
10
200
Sub-Batch Size
Number of Sub-Batches Number Live Blocks
FP
Tokens (K)
Figure Hyperparameter Analysis. a, the dynamics between throughput, GPU memory consumption and hyperparameters, where FP and BP denotes forward and backward propagation, respectively; Ratio in panel refers to the ratio between sub-batch size and number of sub-batches; e, total tokens' effect on host memory usage.
each hyperparameter can result in enhanced arithmetic intensity during computation. Aligned with theoretical analysis, the increasing hyperparameters asymptotically saturates the computation to the practical peak performance of the model operations. Considering the host DRAM as a bottleneck, the maximum token capacity of a given system is fixed. Given this restriction, the tradeoff between sub-batch size and the number sub-batches becomes critical. Both hyperparameters are positively correlated with throughput, however, they have different degrees of impact. We conducted brief test to examine the effects in action. The results (see Fig. 4p) show clearly that the model achieves optimal performance when the sub-batch size is maximized to the limit allowed by HBM.
Performance degradation on long sequence. With the total token number saturating the HBM, we adjusted the sequence length and sub-batch size ratio test trade-off between these two MIBudbuuded properties. The results show decrease total throughput with growing sequence length and shrinking sub-batch size. However, the sequence length is usually determined by the data and is not tunable hyperparameter. This observation does not alter the objective of filling all available HBM during hyperparameter selection.
5.2 Enhanced Arithmetic Intensity
Detailed profiling results further demonstrate that our approach substantially enhances the utilization of the given hardware. By properly applying our scheme, the effective GPU kernel execution ratio rises from 8.9% to 83.7%, shown in Fig The baseline approach spent nearly 80.6% of the time exchanging tensors to and from host DRAM, while SpeedLoader controls the memory accessing ratio to under 15%.
For distributed training with ZeRO stage 3 without offloading, the baseline approach spent 80% of the time partitioning the gradients and 14% of the time gathering weights. contrast, Speedl achieves 66% GPU time executing arithmetic operations. The model operation exhibits distinct behaviors between saturated and unsaturated computations. When the sub-batch size is large enough, kernel executions are seamlessly enqueued in series, typically lasting longer than the host-side function calls of the active layers. Conversely, when the sub-batch size is not large enough to fill the CUDA stream, the accelerator often waits for host-side function calls. The backward computation follows similar trend (Fig. b.c).
Additionally, DeepSpeed backward computation incurs significant blocking communication load, which offloads the gradients to pageable host memory. For every checkpointed segments of model in trial, this communication operation causes 21 10 ms of hardware turnaround without any arithmetic

---

b
5.9%
TO
14. 375
16.1%
52 32.17
10.4%
S1
08.0%
81.7%
0%
CPU Exccution
HeducaScatter
Altiather
Aritmmatic
Figure 5: Profiling results. (a) Deepspeed (w/ offloading); (b) SpeedLoader (unsaturated- computtin.. w/ offloading); (c) SpeedLoader cettuurated-comutingg w/ offloading); (d) Deepspeed (w/o offloading); (e) SpeedLoader (w/o offloading)
operations. On the other hand, the corresponding backward propagation arithmetic operations take 6679 ms, indicating that nearly a quarter of the backward propagation time spent waiting for blocking gradient partitioning. In conclusion, the sub-batch size saturates computation, while the number of sub-batches hinders the synchronization overhead in the computation of every effective batch.
Another interesting observation is that the prefetching operation of DeepSpeed is module-wise asynchronous. Instead of overlapping the communication with all operations before the target tensor called, only one modulc's operation is allowed to be overlapped with single I/O operation. This behavior results in a bottleneck at the end of every active layer group, where parameters of current layer group are released and prefetching kicks off.
5.3 Peak Performance Benchmarking
To demonstrate the pragmatic performance enhancements, we examined the peak performance of different offloading and distribution schemes on the same nodes from an IaaS provider. The results are shown Fig. [
The first benchmark was conducted single device evaluate SpeedLoader's offloading perfor- mance. Under the best configuration for both approaches, our method showed speedups of 3.29 and 6.49 for LLaMA-2-7B and LLaMA-2-13B, respectively. In both cases, SpeedLoader exhibited MFU of %. The second test examined performance in ZeRO distributed training combined with offloading. We distributed the batches across 64 GPUs on nodes interconnected with Ethernet. Our approach attained speedups of 5.48, 5.49, and 30.33 for the 7B, 13B, and 70B I.LaMA-2 models, respectively. Notably, there was a MFU drop compared the single device case, likely reflecting inter-node bandwidth bottleneck. which can be significantly alleviated in clusters with more reliable interconnections. The third test examined SpeedLoader's effectiveness In this trial, we disabled offloading compared to the previous test. Our approach increased the performance gap between the baseline, with speedups of 5.34 and 12.28 on the 7B and 13B models, respectively. This result highlights the potential current sharded training paradigms, where the MFU can be elevated to unprecedented levels.
peer communication.
Additionally, our optimized inference scheme outperforms previous state-of -f----f-thhheeee the-art approaches (Fig. 7). On single NVIDIA A A100, the speedups compared to FlexGen were 1.52, 2.09, and 2.35 for the OPT-6.7B 7B, OPT-30B, and OPT-175B models, respectively. vllm showed consistently inferior performance due to its simplistic handling offloaded computing. Note that the OPT-175B was tested with NVMe offloading and without KV caching due to excessive DRAM consumption. For OPT-175B, the KV cache for each sample with 256 tokens can consume up to 1. 1.125GB of host memory, resulting in heavy I/O overhead and becoming significant bottleneck in maximizing effective batch size. In contrast, strategy without KV caching benefits from larger sub-batch size and gains enhanced arithmetic intensity.

---

Single Device
Distributed
No Officed
MLM DeepSipee
5.0%
100
Thruughput (tokerNA))
6.54
80
SpeedLoader
39%
2094
Model
75
130
76
1:38
708
78
139
Model Size
Model Size
Model Size
6.75
3011
1738
Model Size
Figure 7: Inference through- put benchmarking single device.
Figure 6: Training MFU comparison. Left, Single device benchmark- ing w/ offloading: middle, Distributed benchmarking w/ offloading; right, Distributed benchmarking w/o offloading.
5.4 Functionality Test and Scalability
Our methodology does not alter the model structure and numeric value. To show the performance of our proposed approach in real-world scenario, we also examine the functionality of our approach via pretraining experiment. We pretrained 7B and a 13B model following the corresponding configuration of LLaMA-2. The trials ran Wikipedia, OpenWebText and C4 datasets for cutoff time. This cutoff was set before the full dataset processed, establishing that training loss is valid metric for assessing convergence effectively.
The training results are shown in Tab. 2 These trials were conducted using 4x NVIDIA A100 GPUs, distributed across four nodes interconnected by Slingshot Interconnection. They operated with an effective batch size of 512 and context length of 2048. Given the cutoff time, our approach robustly processes more tokens than baseline implementation trials with three datasets and two model sizes without compromising convergence. The average speedup for 7B and 13B model are 4.08 and 2.68 respectively.
Table 2: Pretraining results of DeepSpeed and SpeedLoader
DeepSpeed
Ours
Dataset
Model Size
Loss
Tokens (M)
Loss
Tokens (M)
7B
2,896
26.2
2.507
113.2
Wikipedia
13B
5.112
14.7
2.461
53.5
7B
4.409
27.3
3.705
116.4
Open WehText
13B
5.008
18.9
4.813
34,6
7B
2.517
35.7
2.199
131.1
C4
13B
3.544
18.9
2.058
48.8
The stark difference in convergence efficiency highlight the practical advantages of our proposed approach, especially when considering I/O operations. In this case, the standard approach required 32 gradient accumulation steps for single parameter update, involving the model being loaded 64 times from the host. Additionally, to synchronize gradients and parameters, the amount of reduce-scatter and all-gather communications was equivalent to over 32 times the model size. Instead, our approach achieves same amount of computing with only two instances of parameter loading, together with only two full model all-gather and one reduce-scatter.
Our proposed scheme delivers excellent performance in distributed settings, especially for model sharding strategies like FSDP and ZeRO, Speed significantly boosts performance by largely reducing the parameter gathering cost with or without offloading. To examine the scalability of our approach, we conducted tests on 64 NVIDIA A100 GPUs on aforementioned cloud platform. The results are shown Fig.

---

In the weak scaling test, we scaled up the total number of accelerators with fixed workload per device with offloading enabled. We observed a superliner scaling pattern, where the throughput per device increased with the number of computing device. For LLaMA-2 70B, per device efficiency has over 3-fold speedup. This behavior is expected due to the bandwidth-centle design ZeRO-Infinity [14]. With the engagement of more accelerators, the PCIe bandwidth bottleneck is largely alleviated. Additionally, our scheme can further reduce the inter-device and inter-node communication overhead for gathering weights and scattering gradients.
6 Limitations and Insights
SpeedLoader provides distributed deep learning settings with an I/O efficient option. By fully utilizing the available host memory, SpeedLoader can signifi- cantly reduce the redundant model state exchange among the device, peer and host memory. However, this approach may make host memory new bottle- neck for efficient model operation. Ad- ditionally, the implementation of mem- ory allocation in Py Torch only allows sizes that are powers of two, which can lead to inefficient use of memory. For example, if tensor requires 33 GB, Py- Torch allocates 64 GB. potentially leading to out-of- issues. Currently, there is no explicit option to modify this behaviour for CPU memory allocation.
Weak Scaling (FP)
Weak Scalling (BP)
Per-Device Threughput
16 Number GPUs
64
16
64
Number GPUs
136
70B
Figure 8: Weak scaling of SpeedLoader. Showing per device throughput.
Furthermore, while SpeedLoader is theoretically compatible with other parallelism strategies like tensor parallelism and pipeline parallelism, 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 Speedl Loader between these parallelisms is valuable area for future research. Amidst an age of LLM blooming, we hope that, together with other heterogencous techniques, SpeedLoader can provide researchers with a new option to explore the frontiers heterogeneous and distributed computing.
Conclusion
SpeedLoader is highly optimized computing scheme for I/O-bounded distributed and heterogeneous LLM operations. By virtue of the layered nature of transformer models, our proposed approach retains layers in-situ process more batches. Through meticulously designed tensor exchange, Speedl Loader enables efficient handling of multiple batches within single reducing the amount of required I/O. This not only allows users to maintain computational accuracy but also ensures time and energy efficiency by minimizing communication overhead. Consequently, speedLLoader provides a practical, high-performance solution for I/O-constrainee distributed and heterogeneous setups, where efficient large-scale model training and inference are key. Our results showed that SpeedLoader exhibits substantially higher performance than previously SOTA approaches during both the training and inference scenarios.
pass, significantly
10

---

8 Acknowledgements
special shout out to Dr. Huang Qirui, Dr. Wei Yuying, Ms. Nie Ying, and Dr. Wang Qinyi for their invaluable insights into this project and their selfless assistance in the writing of this work. The computational work for this article was partially performed using resources from the National Supercomputing Centre, Singapore (httss//www.nscc.gg)
Yang You's research group is being sponsored by NUS startup grant (Presidential Young Professor- ship), Singapore MOE Tier-1 grant, ByteDance grant, ARCTIC grant, SMI grant (WBS number: A-8001104-00-00), Alibaba grant, and Google grant for TPU usage.
References
[1] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprini 1xXix2001.08366 2020.
[2] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas. Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022.
[3] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1900 2020.
[4] AI@ Meta. Introducing meta laaa 3: The most capable openly available 1lm to date. https
7 7/ai meta. wtt.coom/bbla httpllco/obocoeeoeeeeeeeceeeecee whoggg/meaa-aeee 1lama-3. 3/ 2024. Accessed:2022444118
[5] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model 2019. URL
parallelism. CoRR,
08053
[6] Qifan Xu and Yang You. An efficient 2d method for training super-large deep learning models. In 2023 IEEE International Parallel and Distributed Processing Symposium (IPDPS), pages 222-232. IEEE, 2023.
[7] Yanping Huang. Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
[8] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM symposium on operating systems principles, pages 1-15, 2019.
[9] Ziming Liu, Shenggan Cheng, Haotian Zhou, and Yang You. Hanayo: Harnessing wave- like pipeline parallelism for enhanced large model training efficiency. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-13, 2023.
[10] Zechun Liu, Barlas Oguz, Changsheng Zhao, Ernie Chang, Pierre Stock, Yashar Mehdad, Yangyang Shi, Raghuraman Krishnamoorthi, and Vikas Chandra. Llm-qat: Data-free quantiza- tion aware training for large language models. arXiv preprint arXiv:2305. 17888, 2023.
[11] Elias Frantar, Salch Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training guantization for generative pre-trained transformers. arXiv preprint arXiv:2210.1732. 2022.

---

[12] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory opti- mization towards training A trillion parameter models. CoRR, abs/1910.02054, 2019. URL
02054
[13] Jie Ren, Samyam Rajbhandari,, Reza Yazdani Aminabadi, Olatunji Ruwase. Shuangyan Yang, Minjia Zhang. Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training. CoRR, abs/2101 abs/2101.068400, 2021. httphttts://aaat wwwpaaiiv.og//bb...ee
.6840
[14] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero-infinity: Breaking the GPU memory wall for extreme scale deep learning. CoRR, abs/2104.07857, 2021. URL htttps///axiv, org/ www.aaas.....eeee 07857
[15] Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie. Beidi Chen, Clark Barrett, Joseph E. Gonzalez, Percy Liang, Christopher Ré, Ion Stoica, and Ce Zhang. Flexgen: High-throughput generative inference of large language models with a single gpu, 2023.
[16] Reiner Pope. Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems, 5, 2023.
[17] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01065 2022.
[18] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 099888, 2023.
[19] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, al. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint 2r41v:230..11277, 2023.
[20] Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma, Qian Xiong, Fan Yang. and Xuehai Qian. Capuchin: Tensor-based gpu memory management for deep learning. Proceed- ings of the Twenty-Fifhh International Conference on Architectural Support for Programming Languages and Operating Systems, pages 891-905, 2020.
[21] Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song, Zenglin Xu, and Tim Kraska. Superneurons: Dynamic gpu memory management for training deep neural networks. In Proceedings of the 23rd ACM SIGPLAN symposium on principles and practice of parallel programming, pages 41-53, 2018
[22] Jiarui Fang, Yang Yu, Shenggui Li, Yang You, and Jie Zhou. Patrickstar: Parallel training of pre-trained models via basek memory management. CoRR, abs/2108.05818 2021. URL btt.ttss:://rre http::s////xxvv..rrr/..
05818
[23] Xiaoyang Sun, Wei Wang, Shenghao Qiu, Renyu Yang, Songfang Huang, Jie Xu, and Zheng Wang. Stronghold: fast and affordable billion-scale deep learning model training. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-17. IEEE, 2022.
[24] Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Connor Holmes, Samyam Rajbhandari, Olatunji Ruwase, Feng Yan, Lei Yang, and Yuxiong He. Zero++: Extremely efficient collective commu- nication for giant model training. arXiv preprint arXiv:2306. 10209, 2023.
[25] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, abs/1604.06174, 2016. URL http: pectt://arxiis org/ //aaabbs/14/ 06174
12

---

[26] Xunyi Zhao, Théotime Le Hellard, Lionel Eyraud-Dubois Julia Gusak, and Olivier Beaumont. Rockmate: an efficient, fast, automatic and generic tool for re-materialization in pytorch. In International Conference on Machine Learning, pages 42045. PMLR, 2023.
[27] Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and Ion Stoica. Checkmate: Breaking the memory wall with optimal tensor rematerialization. Proceedings of Machine Learning and Systems, 2:497-511, 2020.
[28] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[29] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pages 61 1-626, 2023.
[30] Ahmad Abdelfattah, Azzam Haidar, Stanimire Tomov, and Jack Dongarra. Performance, design, and autotuning of batched gemm for gpus. In High Performance Computing: 31st International Conference, ISC High Performance 2016, Frankfurt, Germany, June 19-23, 2016, Proceedings, pages 21-38. Springer, 2016.
[31] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[32] Tri Dao. Flashattention--: 2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.0869) 2023.
[33] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08k0 08608, 2024.
13

---

A Appendix supplemental material
A.1 List of Abbreviations
Table 3: List of Abbreviations
Abbreviation
Definition
flops
floating-point operations
FLOPs
floating-poin operations per second
MFU
Model FLOPs Utilization
LLM
Large Language Model
HBM
Memory
ZeRO
Zero Redundancy Optimizer
FSDP
Fully Sharded Data Parallel
FP
Forward Propagation
BP
Backward Propagation
GEMM
General Matrix-Matrix Multiplication
A.2 Hyperparameters
Our approach introduces a few new dimensions of hyperparameters:
Sub-batch size. Due to the re-batching operation, the original batched data is parceled into smaller chunks, with the sub-batch size referring to the size of individual chunks. The aligned sub-batch size is a key hyperparameter that impacts computation. The major arithmetic pressure of transformers lies in linear layers, where batched general matrix-matrix multiplication (GEMM) is heavily used. Performance can greatly benefit from larger aligned batch [30). However, the total available on-accelerator HBM is the limiting factor for sub-batch size. The HBM usage mainly shared by activation tensors, parameters in use, and exchange buffers, which highly sensitive other hyperparameters such the prefetch range. This indicates that tuning of sub-bate sub-batch size must take its interaction with other hyperparameters into account.
Number of sub-batches. As discussed Section A.3 a larger effective batch size can result in asymptotic I/O reduction. A lower communication-compute ratio allows easier overlapping for optimal efficiency. However, offloading activations inevitably leads to significant page-locked memory consumption. Since ZeRO-Offload also applies enormous pressure host memory, the host DRAM space becomes major limiting factor for effective batch size.
Number of on-device layers. This hyperparameter specifies the number of model layers computed together. Instead of loading the model layer by layer, we compute multiple layers in one iteration. larger number will result in more HBM consumption for buffering and storing active model parameters. Although this hyperparameter does not significantly affect model training inference efficiency, it is essential for balancing host-device memory costs.
From our observations, we showed that the resource usage of our method has linear relationship to the hyperparameters. Therefore, we can estimate the best batch configurations within only trials by confirming the coefficient and intercepts.
A.3 Communication Analysis
Taking closer look at the communication, in single forward propagation and backward propagation. As listed in Tab, 4 In canonical approach, for every batch with tokens, the model must be loaded twice from host for forward pass and re-computation, and one copy of gradient will be written back to host. In distributed settings, the parameters and gradients are communicated with peer ranks. To sum it up, the amount local and peer tensor communications both 3N P for every N batches.
quantitatively analyze the communication amount
We illustrate the data operations of SpeedLoader in Fig 2 Take sub-batch b example. During forward propagation, for each sub-batch one layer, prefetch the input of next mini-batch (c)) and
14

---

offload the output of previous mini-batch (a2). Meanwhile, we also prefetch the model parameters of the next layer. This sums to 2N + P local and P peer communication.
For backward pass with gradient checkpointing, the input activation (a1) and the output gradient (Vaz) are fetched. Meanwhile, previous activation gradient (Ve)) is offloaddd to the pinned memory that previously storing the activation (c1). Simultaneously. the model parameter of next layer is also being fetched. Above results in 3NA A + 2P local and 2P remote communication. In total, our proposed method has 5NA+3P and 3P local and remote communication, respectively.
Consider that a typical transformer model having a total size of 12L [1 I and checkpointed activation sizes nh, the local and remote I/O ratios between SpeedLoader and ZeRO-Offload 5m + N and N. respectively. Based the quantitative analysis, it is shown that the total 3GLF communication during offloading in SpeedLoader is asymptotically reduced with increasing the number of sub batches.
Table 4: Communication during training, where P. A, N refer parameter, activation size and number of sub-batches respectively.
Collective Communication
Ours
ZeRO
Parameter Gathering (FP)
P
NP
Parameter Gathering (BP)
P
NP
Gradient Reduce-Scatter
P
NP
Total
3P
3NP
Local Communication
Ours
ZeRO
Parameter Loading(FP)
P
NP
Activation Loading(FP)
NA
Activation Offloading (FP)
NA
Parameter Loading (BP)
P
NP
Parameter Gradient Offloading
P
NP
Activation Gradient Loading
NA
Activation Gradient Offloading
NA
Activation Loading (BP)
NA
Total
5 + 3P
3NP
A.4 Enhanced Performance with FlashAttention-2
We explored the synergy between SpeedLoader and FlashAttention-2233 We conducted experiments under identical hyperparameters with FlashAttention-2 enabled and disabled, ensuring maximum hardware utilization. Results (Fig. 10 showed that in distributed offloading scenarios, FlashAttention- provides 1.13x, 1.13x, and 1.05x speedups for LLaMA-7B, -13B, and -70B, respectively. In distributed training without offloading, the speedups are 1.16x and 1.13x for 7B and 13B models. For these tests, used 1, 2, and 4 nodes with 16 NVIDIA A A100-40GB GPUs each for 7B, 13B, and 70B trials, respectively. With this optimization and the upcoming FlashAttention-3,, SpeedLoader can achieve throughput very close to non- non-sherded and non-offloaded training.
A.5 Compatibility Test
SpeedLoader is designed to facilitate computing in I/O bounded settings. Therefore, link speeds of benchmark platform have substantial impacts on the performance of our solution. To explore LpderLoader's performance on devices with various I/O capabilities and architecture, tested the speedups of SpeedLoader on multiple distinct platforms. Details of benchmark platforms can be found in Tab. [
15

---

50%
Distributed
No Officed
SpeedLoader
60%
DeepSpeed
Model FLD% Unlivation (KG
Model FLOPs Unlicaton
50%
90%
42%
20%
30%
20%
10%
0%
H100
A6000
V1008
0%
138
708
138
Model Size
Model Sioee
Figure 9: Single device peak performance test on devices with different bus speed.
Figure 10: Training MFU ablation study with FlashAttention-2.
With speedups of 1.37x, 3,56x, and 3.0Ix on VI00S, A6000, and H100 (Fig. 9). SpeedLoader exhibited consistent gains on platforms with PCIe Gen 4 and over, and lower speedups on V100S with PCTe Gen Overall, the positive speedups showease robust compatibility on platforms with mixed architectures and capabilities. Note that the MFU being lower on H100 than A100 is expected behaviour. As reported by Dao et al. (.3) without architecture-specin optimizations like FlashAttention-3, transformers struggle unleash the full potential of Hopper GPUs.
A.6 Reproducibility Statement
To facilitate reproduction of experimental results, we provide an public repository As an example setting, we recommend readers reproduce part of aforementioned results with instances of A2 VMs on Google Cloud. This SKU has similar specifications to the computational resources of this research.
Access here
16

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA]
[NA] means either that the question is Not Applicable for that particular paper or the
relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] provvdddd proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA]" is not grounds for rejection. While the questions are phrased a binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write a justification elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] a question, in the justification please point to the section(s) where related material for the question can be found.
the checklist one of the factors their evaluation.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The contributions of our work are clearly stated in the abstract and introduction, which are backed by solid experimental results.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: As mentioned in conclusion, our work is not tested with other parallelisms like pipeline or tensor parallel. Apart from this aspect, the paper is not based on strong assumptions and overall generalizable.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
17

---

Answer: [NA]
Justification: Our research does not include theoretical results.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: We described the working mechanism of proposed method in the paper in detail, together with clear pseudocode and schematics. All experiments are conducted on open-source datasets and models.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our code repository is publicly available. Users can follow the instructions in Sec. A.6
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We clarified the hyperparameter tuning strategies and disclosed static hyperpa- rameters for the experiments.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The experiments are very computationally exhaustive, therefore impractical to create enough data for meaningful statistical tests.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: The computing infrastructure specifications clearly stated in Sec.
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httphhghppsss/...........ppp/paa
Answer: [Yes]
Justification: We reviewed the NeurlPS Code of Ethics and this research conforms with the guidelines.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
18

---

Answer: [NA]
Justification: Our work does not impose any foreseeable societal impact.
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper poses no such risks.
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The assets used in this work are properly acquired and cited, in accordance with the guidelines.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our work does not release new assets.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our research does involve crowdsourcing nor research with human subjects.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our research does involve crowdsourcing nor research with human subjects.
19


Reviews:
[{"review_id": "ZtGXTcvJRV", "reviewer": "Reviewer_bDTF", "summary": "This paper introduces an innovative compute strategy that minimizes I/O between peers and devices in heterogeneous LLM training, and a high-efficiency tensor manager that optimizes transfers between device and host, reducing fragmentation and redundancy. These enhancements lead to superior inference efficiency compared to existing state-of-the-art methods.", "questions": "1. The authors claim SpeedLoader can process multiple batches for one full model loading. What’s the difference from using a larger batch size?\n2. What is your experiment configuration of Figure 5? It seems that SpeedLoader minimizes the ratio of Memcpy, yet the ratio of AllGather does not change too much. Do you consider the overlap between Memcpy and Kernel? Why the communication only have a ratio of about 10% in the training scene?", "limitations": "Yes, the authors have adequately addressed the limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The topic of this paper is current and timely.\n2. The evaluation makes sense. This paper pretrains a 7B and a 13B model following the corresponding configuration of LLaMA-2 and shows better performance than DeepSpeed.", "weaknesses": "1. The optimization technique seems to be incremental.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Response to Question 1\n\nThank 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**.\n\n### Response to Question 2\n\nWe 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:\n\n**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)\n\nThe 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.\n\nWe 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.\n\nThe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "apvkNGnN5s", "reviewer": "Reviewer_m7wP", "summary": "This paper proposes a new paradigm called SpeedLoader for large language models (LLMs) inference and training to offload and reload the layer weights and activations to and from CPU memory. Compared to prior works, SpeedLoader can process multiple batches with the current active layer, which thus able to reduce communication traffic in CPU-GPU transfer and distributed settings. To implement this, SpeedLoader comes with a tensor exchange manager that minimizes memory fragmentation and redundancy. The experimental results show that SpeedLoader can outperform all existing solutions.", "questions": "Thanks for submitting this nice work to NeurIPS, which focuses to solve an important and timely problem. In general, I liked the idea of putting multiple batch together. However, there are a few concerns and comments about the paper:\n\nFirst, in the experimental results regarding to the performance numbers, there are no comparison with the ideal cases, i.e. without offloading. This comparison is important to understand the positioning of the SpeedLoader, e.g. whether SpeedLoader still have room for optimizations.\n\nSecond, in the evaluation plots, the number of tokens in each experiment have not been clearly discussed. This discussion only happens for the functionality test and scalability test. However, showing the (practical) number of tokens in the performance plots are also important because as the number of tokens increases, it might be the case that kernel time dominates (as attention time complexity is quadratic to the number of tokens), which could make SpeedLoader less motivated.\n\nThird, it would be appreciated if a limitation section is presented. This not only helps the reader better positioning the paper, but also inspires future works that improves upon the SpeedLoader.", "limitations": "* This work have only evaluations on A100. It is thus unknown how it performs with other generation of GPUs, especially with other PCIe generations that could affect the trade-offs.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* SpeedLoader finds a novel way to offload layer weights and activations to CPU memory;\n* SpeedLoader achieves the state-of-the-art results compared to existing solutions;", "weaknesses": "* No comparison with ideal scenario (i.e. without offloading).\n* Experimental details are missing from the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Response to Question 1\n\nThank 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. \n\nFollowing 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.\n\n### Response to Question 2\n\nSorry 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.\n\nTo 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.\n\n### Response to Question 3\n\nThank 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. \n\n**Limitation**\n\nSpeedLoader 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.\n\nFurthermore, 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.\n\n## Additional Comments\n\nFollowing 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CH53NatIQD", "reviewer": "Reviewer_4D75", "summary": "This paper proposes SpeedLoader, a system for offloading parameters and activations under restricted resources for distributed training and inference.  It utilizes a tensor exchange manager to minimize the communication overheads.  Emperically, a larger proportion of time is spent on computation rather than communication or data movement compared with previous systems, leading to a higher MFU.", "questions": "FlexGen has the zig-zag block schedule, which looks similar to the schedule in Figure 1 right side.  Could the author explain their difference?", "limitations": "The limitations are acknowledged but not addressed.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The method does not require training and is only at inference cost. \n\n2. The system has a super linear scalability with the number of GPUs because of the alleviation for PCIe bandwidth.", "weaknesses": "Only parameter sharding is considered, other popular scenario like pipeline parallelism covered in FlexGen, tensor parallelism, are not explored here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Response to Question \n\nThank 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.\nFor 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. \n\nAdditional Comments\n\nThank 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "noFaHrzfUJ", "reviewer": "Reviewer_Dtr1", "summary": "This paper propose a new scheduling mechanism via effective sub-batch centric computation, to facilitate the tensor computation and exchange, which is proven to particularly minimize the excessive communication overhead.", "questions": "1. How does the one-shot hyperparameter tuning work?\n2. It would be more convincible with adding more frameworks for detail comparison, for example, PagedAttention, FlashAttention, OLLA.\n3. The evaluation is conducted over 16 GPUs, with results appearing to be steady and not demonstrating increased scalability. Could you please include additional GPUs for a more comprehensive evaluation?", "limitations": "The authors have adequately addressed the limitation of the potential negative societal impact of their work.\nNo discussion on limitations, I suggest the authors to add one.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. The idea of redesign the data flow and sharded model training under restrictive computational resources is interesting.\n2. The paper has good coherence, and is well-structured.\n3. The introduction section explains the concept in an easy and understandable way.\n4. The evaluation demonstrates a significant improvement in Model FLOPS utilization compared to DeepSpeed.", "weaknesses": "1. Many details about the design and implementation of SpeedLoader are missing. For example, how is the structure and management of sub-batches handled (L105)? The authors claim that \"the method allows the computation to be equivalent to gradient accumulation across several sub-batches,\" but no theoretical proof is provided.\n2. The discussion of the hyperparameter tuning strategy in Section 3.3 is unclear. How does the one-shot hyperparameter tuning work? What benefits does it bring? The paper needs to provide a much clearer explanation, as many important implementation details are currently missing.\n3. The novelty is marginal, given that much prior art exists on optimization under restrictive computational resources. The paper would benefit from a comparison with more recent methods, such as PagedAttention, to clearly demonstrate differences in performance, resource requirements, and other relevant metrics.\n4. Offloading is identified as the most significant factor in addressing the HBM limitations. However, the experimental results show that after 16 GPUs, the trend stabilizes, which does not convincingly demonstrate an enhancement in scalability.\n5. Minor issues: \n* Please ensure consistency in the capitalization of \"FLOPS”, rather than “flops”.\n* In Figure 4 where FP and BP denotes -> denote", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Response to Question 1\nThank 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.\n\nThe 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.\n### Response to Question 2\nThank 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.\n\nRecent 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.\n\nAdditionally, 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.\n### Response to Question 3 and Clarification on Weakness 4\nWe 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.\n\nIn 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).\n\nTherefore, 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.\n\nWe 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.\n\nPlease let us know if there are any further concerns or if additional clarifications are needed.\n### Clarification on Weakness 1\nThank you for your valuable feedback. Below is a revised explanation of the design and implementation details for sub-batching and gradient accumulation in SpeedLoader:\n\nHere 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.**  …\n\n…\n\nTo 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 …\n\nWe 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.\n### Comments on minor issues\nThank 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear Reviewers,

Thank you for your detailed and insightful feedback. We have conducted additional experiments to enhance the evaluation of SpeedLoader, making it more credible and comprehensive. Below, we address your concerns and clarify some points:

1. **Impact of Different Sequence Lengths Under Full HBM Utilization:** To examine the MFU fluctuation with varying context lengths, we conducted trials with a constant number of total tokens, altering the ratio between batch size and sequence length. The results (Fig. S2) indicate that increasing context length results in a marginal MFU increment (no more than 5% of maximum throughput). This does not diminish SpeedLoader’s advantages, as it consistently allows larger batch sizes, leading to more efficient kernel execution. This observation aligns with our manuscript.

2. **Compatibility Test on Devices with Different I/O Capabilities:** To explore SpeedLoader’s performance on devices with various I/O capabilities and architectures, we tested SpeedLoader on multiple distinct platforms. Due to insufficient collection of GPU options on previous IaaS provider, we sourced the following platforms from different providers:

    | GPU           | PCIe Specification | Source       |
    |---------------|--------------------|--------------|
    | NVIDIA V100S  | Gen 3              | IaaS Provider|
    | NVIDIA RTX A6000 | Gen 4           | Personal     |
    | NVIDIA H100   | Gen 5              | IaaS Provider|

    The results showed 1.37x, 3.56x, and 3.01x speedups on V100S, A6000, and H100, respectively (Fig. S3). The positive speedups showcase robust compatibility across platforms with mixed architectures and capabilities. Note that the MFU being lower on H100 than A100 is expected behavior. As reported by Dao et al., without architecture-specific optimizations like FlashAttention-3, transformers struggle to unleash the full potential of Hopper GPUs. However, the FlashAttention-3 codes are not yet available.

3. **PagedAttention with Offloading:** To provide a more recent reference, we tested the vLLM’s (with PagedAttention) inference throughput with SpeedLoader (Fig. S4). Despite vLLM implementing naïve offloading strategies, it is not highly optimized for this feature. It achieved 22.0 and 4.49 tokens/s throughput for OPT-6.7B and OPT-30B, respectively, and encountered out-of-memory issues when attempting the OPT-175B trial.

4. **Enhanced Performance with FlashAttention-2:** We explored the synergy between SpeedLoader and FlashAttention-2 (Dao et al., 2023). We conducted experiments under identical tuned hyperparameters with FlashAttention-2 enabled and disabled, ensuring maximum hardware utilization. Results (Fig. S5) showed that in distributed offloading scenarios, FlashAttention-2 provides 1.13x, 1.13x, and 1.05x speedups for LLaMA-7B, -13B, and -70B, respectively. In distributed training without offloading, the speedups are 1.16x and 1.13x for 7B and 13B models. For these tests, we used 1, 2, and 4 nodes with 16 NVIDIA A100-40GB GPUs each for 7B, 13B, and 70B trials, respectively. With this optimization and the upcoming FlashAttention-3, SpeedLoader can achieve throughput very close to non-sharded and non-offloaded training.

To prevent misunderstanding and confusion, we made some adjustments and additions to the text (details are presented in individual responses):

1. **Revised Section 3:** We have provided more information on the design of SpeedLoader to help readers better understand its mechanism.

2. **Added Limitation Section:** To help readers better position the paper and inspire future work, we have added a limitation section:
   
   **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.


3. **Abbreviation Explanation:** We have added an explanation to help readers differentiate between floating-point operations (flops) and floating-point operations per second (FLOPs).

4. **Grammatical Errors:** We have corrected various grammatical errors throughout the manuscript.

The above changes and additional results will be integrated into the camera-ready version of the manuscript if the paper is accepted. Thank you again for reviewing this paper and for the valuable feedback provided.

Best Regards

Meta-review: Reviewers agree this paper addresses an important and timely problem.

Reviewers initial criticisms (e.g. lack of comparison on other platforms, lack of discussion about FlexGen and PagedAttention) are well addressed in the rebuttal, and we expect to be well addressed in a final version.
Decision: Accept (poster)