1. Is ZeRO-Infinity the SOTA Approach

ZeRO-Infinity is the most widely adopted framework for tensor offloading, which has been integrated into the DeepSpeed framework. Recent tensor offloading studies also include FlashNeuron, Smart-Infinity, and others. FlashNeuron only considered offloading activation tensors, Smart-Infinity focused on using specialized near-storage accelerators to accelerate optimizer updates.

To the best of our knowledge, ZeRO-Infinity is still a well-acknowledged and practical SOTA approach for tensor offloading, which serves as the most competitive baseline for TeraIO.

2. Why TeraIO Outperforms ZeRO-Infinity

TeraIO outperforms ZeRO-Infinity due to its design differences in offloading granularity and migration strategy.

Offloading Granularity: ZeRO-Infinity plans the offloading of parameters and optimizer states at the layer granularity of ML models. This introduces bursts in I/O patterns, and can underutilize the migration bandwidth. In contrast, TeraIO considers all tensors and performs offloading at the GPU kernel granularity based on our precise tensor lifetime analysis. Our approach ensures high and consistent migration bandwidth utilization (see Fig. 9 (2)).

Migration Strategy: ZeRO-Infinity uses a heuristic-based policy to decide which tensors should be offloaded and when the offloading should start, regardless of the storage bandwidth usage. It lacks global optimization and produces unpredictable I/O patterns. In contrast, TeraIO quantifies the benefit and cost of tensor offloading via I/O-aware planning (Section 3.2), which generates globally optimized migration plans that maximize I/O bandwidth utilization.

3. Clarification of CPU Memory Requirements of ZeRO-Infinity

Since ZeRO-Infinity performs optimizer state updates on CPU, its design intentionally places gradients and optimizer states in CPU memory to reduce CPU-SSD I/O bandwidth consumption. This high CPU memory usage is fundamental to ZeRO-Infinity’s design.

4. Clarification of Speedup Breakdown

The speedup of TeraIO is mostly attributed to a more intelligent migration mechanism. To validate this, we conducted experiments when disabling GPUDirect Storage (GDS) and using regular host-based storage (first read data from SSD to host memory, and then transfer to GPU). We obtained similar end-to-end performance (see table below).

Table: Training throughput comparison with and without GDS (tokens/s)

Although GDS does not have a significant impact on the end-to-end performance in our current setting (2 GPUs, each connected to 4 SSDs), it provides significant scalability advantages by eliminating host resource contention.

Without GDS, tensor offloading frameworks will suffer from scalability bottlenecks. To validate this, we scale the number of GPUs and SSDs and measure the host resource usage. As shown in the table below, host resource usage increases linearly. With 2 GPUs and 4 SSDs per GPU, over 100 GB/s of host memory bandwidth is consumed, and more than 3 CPU cores are fully used, since the host memory is used as a bounce buffer for data transfers between GPUs and SSDs.

Table: Host resource utilization without GDS

*CPU Usage: Usage relative to a single CPU core.

It is worth noting that an important goal of TeraIO is to achieve low-cost LLM training, it is critical to minimize the use of host CPU and memory resources. Compared to ZeRO-Infinity, TeraIO significantly reduces the host CPU and memory requirements, while offering better scalability for offloading tensors to low-cost SSDs.

5. Novelty of Tensor Profiler

Compared to existing DNN profilers, our tensor lifetime profiler has two unique features:

(1) Kernel-level tensor lifetime tracking: Our profiler can precisely capture active/inactive periods of each tensor across GPU kernels. Existing profilers like Nsight Systems focuses on performance bottleneck identification, while PyTorch Profiler provides only memory snapshots and kernel timing. But none of them can provide the tensor lifetime information needed for our optimized migration planning.

(2) Lightweight profiling with PyTorch integration: We instrumented PyTorch's automatic operator generator to capture all tensor activities at runtime with minimal overhead. Because of this, our profiler can capture all tensors and their operators in a comprehensive manner, while existing tools like torch.compile() cannot achieve this, when the model includes custom operators or autograd hooks.

6. Theoretical Analysis of Expected Performance

TeraIO provides a theoretical performance model for estimating the training time in the tensor migration algorithm. As described in Section 2.2, our performance model takes tensor sizes, profiled kernel execution time, the tensor migration plan, available I/O bandwidth, and GPU memory capacity as inputs to predict the training time. With each kernel’s execution time collected in our profiling, our performance model simulates tensor migration at the designated bandwidth and checks whether the tensors needed by the kernel are already in the GPU memory. If not, these tensors’ migration overhead will be added to the total training time. This enables us to examine different migration plans with different performance-critical parameters, and check whether the migration overhead can be fully hidden with the migration plan, as demonstrated in the roofline analysis (Fig. 3).

7. Activation Recomputation

We choose activation offloading over recomputation for two major reasons: (a) Our roofline analysis shows that, with 32GB/s migration bandwidth, TeraIO can achieve near-ideal training performance. This bandwidth can be easily achieved by aggregating multiple commodity SSDs. In contrast, recomputation will consume precious GPU resources and could introduce extra computation overhead. (b) It is much easier to scale SSD resources at a lower cost by adding more SSD devices, compared to the approach of scaling expensive GPU resources. This reflects our “low-cost” design philosophy.

Moreover, activation recomputation is orthogonal to offloading. As having both activation recomputation and offloading significantly increases the complexity of the framework design and implementation, it requires extensive investigation and analysis. We wish to explore this as future work.

8. Clarification of Settings in the Characterization Study

We list the settings in the table below. For diversity, we used different batch sizes across different models to study their impact on the memory capacity requirements.

As we vary the parameters shown in the table, we reach the same conclusion. This is because of the inherent LLM characteristics. As most LLMs use transformer architectures, their sparse tensor access patterns and high compute intensity stem from the fundamental structure of attention mechanisms and multi-layer transformer architecture.

9. Clarification of Performance Gap

The performance gap still exists when the memory requirement of an LLM significantly exceeds the available GPU capacity. For large models like Llama3-70B (>900% memory usage), the volume of tensor migrations overwhelms the available I/O bandwidth even with optimal tensor migration plans. This causes GPU kernel stalls when the required tensors are still being migrated. The performance gap could be eliminated by scaling the migration bandwidth (e.g., adding more SSDs), as demonstrated in our roofline analysis (Section 2.2).

10. Kernel Stalls

To elaborate on kernel stalls and their impact on performance, we show the breakdown of end-to-end execution time in the table below. It includes three components: (a) the time when tensor migrations stall GPU computation, (b) the time when tensor migrations overlap with GPU computation, and (c) the remaining GPU computation time. On average, kernel stalls account for less than 9% of the total execution time, thanks to our optimized tensor migration schedule.

These kernel stalls can be predicted and minimized using our performance model as discussed in Section 2.2 and Question 6).

Table: Breakdown of training iteration time (kernel stall%, overlap%, computation only%)

11. Emerging Memory Technologies

Although NVLink-C2C in NVIDIA Grace systems provides high GPU-CPU bandwidth, the memory capacity is still hard to scale, due to the fundamental DRAM scaling walls (physical limitations, power wall, and architecture constraints). As the memory capacity is insufficient for large models, expanding the GPU/host memory with low-cost and scalable SSDs ($0.2/GB for SSDs vs.$4/GB for DRAM on average) has become a practical and promising solution.

Even if the cost is not a concern, our lifetime-aware tensor offloading approach can be applied to new and emerging memory technologies such as CXL memory. By expanding the GPU/host memory with external memory devices via CXL, new performance tradeoffs need to be considered, as the bandwidth and latency of accessing CXL memory vs. SSDs are different. Unfortunately, commodity CXL memory devices are not widely available on the market yet. We wish to extend TeraIO with new memory technologies as future work.

Dear Reviewer 9SW7,

Please verify whether the authors have adequately addressed your comments, post your response amd submit the “Mandatory Acknowledgement” accordingly. Kindly note that you must engage in discussions with the authors and then submitting the “Mandatory Acknowledgement” to avoid being marked as a non-participating reviewer, which may result in penalties.

Thanks for your support!

1. Generalizability Issue of NVIDIA GDS

We used GPUDirect Storage (GDS) in TeraIO, because it provides significant scalability advantages by bypassing the performance bottlenecks caused by the host CPU. Without GDS, the conventional host-based approach (CPU reads data from SSDs to host memory, and then transfers to GPU) will consume significant hardware resources (CPU and memory) on the host, which not only creates bottlenecks as we scale the number of GPUs and SSDs, but also increases the cost for LLM training.

TeraIO does not necessarily rely on NVIDIA’s GPUDirect solutions. It is compatible with any peer-to-peer (P2P) direct storage solutions, such as AMD's DirectGMA.

2. Offloading Target Choice in the Migration Algorithm

TeraIO uses both CPU memory and SSDs as tensor offloading targets, as described in our tensor migration algorithm in Section 3.2 as well as the evaluation in Section 4.

TeraIO supports simultaneous migrations. However, it prioritizes SSDs because of their unique advantages in terms of large capacity, low cost ($0.2/GB), and scalable I/O bandwidth with GPUDirect Storage support. When generating optimized tensor migration plans, TeraIO predicts the bandwidth utilization of SSDs based on previously decided tensor migrations (see Line 17 in Algorithm 1). If TeraIO estimates that the SSD bandwidth will be saturated, it will offload the candidate tensors to the host memory.

When simultaneous migrations to both host memory and SSDs are needed, the GPU-CPU transfers will utilize the remaining PCIe bandwidth. TeraIO's I/O-aware tensor migration planning algorithm has taken this into consideration by tracking bandwidth usage across both targets (see Algorithm 1).

3. DRAM Scalability Challenge

The DRAM scalability challenge refers to the fundamental limitations in scaling DRAM capacity. They include (1) physical limitation in which the number of pins on a processor package and DIMM modules is limited by the printed circuit board trace routing, which cannot easily scale; (2) power limitation in which higher-capacity DRAM modules consume significantly more power, resulting in thermal management challenges; and (3) scaling limitation in which DRAM cells cannot be made smaller without causing DRAM cell disturbances and retention time violations.

4. Paper Formatting Concerns

We will take an editorial pass at the paper and fix the formatting issues.

1. Comparison to Smart-Infinity

(1) The research approach of TeraIO is fundamentally different from Smart-Infinity. Smart-Infinity focuses on using specialized computational storage devices to accelerate the optimizer state updates. It can be viewed as “ZeRO-Infinity + computational storage”. TeraIO proposes a low-cost tensor offloading framework that focuses on optimizing tensor migrations with commodity SSDs.

(2) The hardware settings of TeraIO vs. Smart-Infinity are different. Smart-Infinity requires specialized computational storage devices. TeraIO uses commodity SSDs and does not require hardware changes or specialized hardware accelerators.

(3) Smart-Infinity's tensor offloading mechanism remains identical to ZeRO-Infinity. With lifetime-aware tensor migration schemes, TeraIO significantly outperforms the offloading approach proposed in ZeRO-Infinity, as examined in our evaluation. For the above reasons, our current submission did not include the comparison with Smart-Infinity. We will add the discussion in the paper.

1. Ablation Studies

We conducted the ablation study as follows.

We conducted experiments when disabling GPUDirect Storage (GDS) and using conventional host-based I/O approach (CPU reads data from SSDs to host memory, then transfers to GPU). We obtained similar end-to-end performance (see table below), showing that the performance speedup is mostly attributed to the optimized tensor migration planning.

Table: Training throughput w/ and w/o GDS (tokens/s)

GDS provides significant scalability advantages by eliminating host resource contention. Without GDS, tensor offloading frameworks suffer from scalability bottlenecks. To validate this, we scale the number of GPUs and SSDs and measure the host resource usage. As shown in the table below, host resource usage increases linearly with the number of GPUs and SSDs. With 2 GPUs and 4 SSDs per GPU, over 100 GB/s of host memory bandwidth is consumed, and more than 3 cores are fully used, as host memory is used as a bounce buffer for data transfers between GPUs and SSDs.

Table: Host resource utilization without GDS

*Usage relative to a single core.

We also show the breakdown of end-to-end execution time in the table below. It includes three components: (a) the time when tensor migrations stall GPU computation, (b) the time when tensor migrations overlap with GPU computation, and (c) the remaining GPU computation time.

Table: Breakdown of training iteration time (kernel stall%,overlap%,computation only%)

Compared to ZeRO-Infinity, TeraIO has less stall time, as it maximizes the overlap of tensor migrations with computation. TeraIO outperforms ZeRO-Infinity for the design differences in two aspects:

Offloading Granularity: ZeRO-Infinity plans the offloading of tensors at the layer granularity of ML models. This introduces burst I/O patterns and can underutilize migration bandwidth. In contrast, TeraIO performs tensor offloading at the GPU kernel granularity based on our precise tensor lifetime analysis. This approach ensures high and consistent migration bandwidth utilization (see Fig. 9 (2)).

Migration Strategy: ZeRO-Infinity uses a heuristic-based policy to decide which tensors should be offloaded and when the offloading should start, regardless of the storage bandwidth usage. It lacks global optimization and produces unpredictable I/O patterns. In contrast, TeraIO quantifies the benefit and cost of tensor offloading via I/O-aware planning (Section 3.2), which generates globally optimized migration plans that maximize I/O bandwidth utilization.

2. Benefits of GPUDirect Storage (GDS)

We choose GDS since it provides significant scalability advantages (see Question #1).

An important goal of TeraIO is to achieve low-cost LLM training, thus it is critical to minimize the use of host CPU and memory resources. Compared to ZeRO-Infinity, TeraIO significantly reduces the host resource requirements while offering better scalability.

3. LLM Workloads Evaluated

We add GPT2-20B in our evaluation and the results as follows:

Table: Training throughput (tokens/s) for GPT-20B

The results are consistent with our evaluation in Fig. 7.

4. Fig. 4

We will fix it in the paper.

5. ZeRO Settings

We show the performance-critical parameters in the table below. With these settings, we ensure ZeRO-Infinity achieves reasonable performance.

Table: Performance-critical parameters in ZeRO-Infinity

We enabled the pipeline_read/write parameters to optimize computation and data I/O overlap during optimizer state updates. We tuned parameters pin_memory, buffer_count, and buffer_size to optimize tensor offloading throughput. For param_persistence_threshold and model_persistence_threshold, we use default values.

6. Compilation Assumptions

Similar to state-of-the-art tensor offloading studies like ZeRO-Infinity, TeraIO uses ML compilers to extract rich semantic information from models for guiding the tensor offloading. This approach is a standard practice as current LLM training systems already use compilers like torch.compile() for performance optimizations. Extracting semantics from model execution comes at no additional cost.

For non-compiled pipelines, TeraIO can still track tensor operators and measure kernel execution times at runtime by instrumenting the PyTorch framework, and collecting kernel execution time, input/output tensors for each operator upon its execution at runtime (Section 3.1).

TeraIO is also extensible for supporting LLMs that involve dynamic computations. Take MoE for example, although expert routing is dynamic and is determined at runtime, most parts of the model are still static. TeraIO can identify the static parts during the profiling process and perform lifetime-aware migrations for their tensors. For the dynamic parts, such as expert routing, we can treat them as single operators and skip the fine-grained migrations for their tensors.

7. Tradeoffs Between Gradient Checkpointing and Activation Offloading

We choose activation offloading over recomputation for two reasons: (a) our roofline analysis shows that with the migration bandwidth of 32GB/s, TeraIO can achieve near-ideal performance. This bandwidth can be easily achieved by aggregating multiple SSDs. In contrast, recomputation will consume precious GPU resources and introduce extra computation overhead. (b) It is much easier to scale SSD resources at a lower cost by adding more SSD devices, compared to the approach of scaling expensive GPU resources.

Activation recomputation is orthogonal to offloading. As having both activation recomputation and offloading significantly increases the complexity of the framework design and implementation, we wish to explore this as future work.

8. Batch Size Boost Gain

Table: Performance of training LIama3-8B using TeraIO with various batch sizes

The table above shows the throughput of training Llama3-8B (seq_len=2K) with different batch sizes. The hardware setting of two H100 GPUs has sufficient memory for training with batch size 32 and micro-batch size 2. However, with larger (micro)batch sizes, out-of-memory errors will happen.

TeraIO enables training with batch size of 128, delivering 9% higher throughput compared to training without offloading (but using the maximum batch size). This shows that, for the models that can fit in the GPU memory, TeraIO can still achieve better performance by allowing larger batch sizes.

9. Clarification of the ZeRO-Infinity Data Point

As ZeRO-Infinity performs optimizer state updates on CPU, it requires CPU memory to store the optimizer states.

In Fig. 8, we evaluate the performance of ZeRO-Infinity with 4 SSDs. We obtained the performance with the minimum CPU memory it requires. For Llama3-8B, TeraIO outperforms ZeRO-Infinity in both performance and CPU memory usage. For Llama3-70B, we added an experiment and ran ZeRO-Infinity with 1TB CPU memory. TeraIO still outperforms ZeRO-Infinity. The result is in the table below.

Table: Throughput (tokens/s) of training Llama3-70B with different CPU memory capacities

10. Novelty of Tensor Lifecycle Management

TeraIO differs fundamentally from these existing works in both research goal and approach.

(1) SwapAdvisor, HOME, and PatrickStar focused on data migrations between GPU and CPU memory; Checkmate focused on alleviating the GPU memory bottleneck with tensor recomputation. These studies did not consider the cost-efficiency of LLM training as the first citizen in their designs. In contrast, TeraIO focuses on the tensor migrations between GPU memory and low-cost SSDs, aiming to provide a cost-effective approach for training LLMs.

(2) TeraIO introduces unique advantages in tensor lifetime analysis and migration planning. Unlike existing works that rely on model dataflow graphs, TeraIO uses precise tensor lifetime information for generating optimized and fine-grained tensor migration plans. Instead of relying on heuristic-based approaches, the I/O-aware migration algorithm in TeraIO tracks estimated storage bandwidth usage and quantifies the benefit and cost of each planned tensor migration. This enables TeraIO to generate globally optimized migration plans that maximize I/O bandwidth utilization.