Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage
Authors: Ziqi Yuan
Haoyang Zhang
Yirui Eric Zhou
Apoorve Mohan
I-Hsin Chung
Seetharami Seelam
Jian Huang
Venue: NeurIPS 2025 poster ()
Keywords: GPU Memory
GPUDirect Storage
Large Language Model
Offloading
Cost-Efficient Training
Solid State Drives
Primary area: infrastructure

Abstract: We present the design and implementation of a new lifetime-aware tensor offloading
framework for GPU memory expansion using low-cost PCIe-based solid-state
drives (SSDs). Our framework, TERAIO, is developed explicitly for large language
model (LLM) training with multiple GPUs and multiple SSDs. Its design is driven
by our observation that the active tensors take only a small fraction (1.7% on
average) of allocated GPU memory in each LLM training iteration, the inactive
tensors are usually large and will not be used for a long period of time, creating
ample opportunities for offloading/prefetching tensors to/from slow SSDs without
stalling the GPU training process. TERAIO accurately estimates the lifetime (active
period of time in GPU memory) of each tensor with the profiling of the first few
iterations in the training process. With the tensor lifetime analysis, TERAIO will
generate an optimized tensor offloading/prefetching plan and integrate it into the
compiled LLM program via PyTorch. TERAIO has a runtime tensor migration
engine to execute the offloading/prefetching plan via GPUDirect storage, which
allows direct tensor migration between GPUs and SSDs for alleviating the CPU
bottleneck and maximizing the SSD bandwidth utilization. In comparison with
state-of-the-art studies such as ZeRO-Offload and ZeRO-Infinity, we show that
TERAIO improves the training performance of various LLMs by 1.47× on average,
and achieves 80.7% of the ideal performance assuming unlimited GPU memory.

Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage
Ziqi Yuan', Haoyang Zhang', Yirui Eric Zhou', Apoorve Mohan I-Hsin Chung", Seetharami Scelam", Jian Huang' University of Illinois Urbana-Champaign, 2IBM Research {ziqiy6, zhang402, yiruiz2, jianh)@illinoii edu apoorve www.eee..mmmiiiim com, (ihchung, sseelam) 0us. ibm. com
Abstract
We present the design and implementation of new lifetime-aware tensor offloading framework for GPU memory expansion using low-cost PCIe-based solid state drives (SSDs). Our framework, TERAIO, developed explicitly for large language model (LLM) training with multiple GPUs and multiple SSDs. Its design is driven by our observation that the active tensors take only a small fraction (1.7% on average) of allocated GPU memory in each LLM training iteration, the inactive tensors are usually large and will not be used for a long period time, creating ample opportunities for fifloading/prffecching tensors to/from slow SSDs without stalling the GPU training process. TERAIO accurately estimates the lifetime (active period of time in GPU memory) of each tensor with the profiling of the first few iterations in the training process. With the tensor lifetime analysis, TERAIO will generate optimized tensor offfetciing/refeeccinng plan and integrate it into the compiled LLM program via PyTorch. TERAIO has a runtime tensor migration engine to execute the offloaddig//refetchingg plan via GPUDirect storage, which allows direct tensor migration between GPUs and SSDs for alleviating the CPU bottleneck and maximizing the SSD bandwidth utilization. In comparison with sate-offthe-aa studies such as ZeRO-Offloa and ZeRO-Infinity, we show that TERAIO improves the training performance of various LLMs by 1.47x on average, and achieves 80.7% of the ideal performance assuming unlimited GPU memory.
1 Introduction
I Large language models (LL.Ms) have been widely employed in various application domains 31 42 However, training LLMs is still a well-known challenging problem, as its memory demand increasing at a much faster speed than the scaling speed of GPU memory [8 25 Scaling the GPU memory capacity is both technically challenging and prohibitively expensive. Due to space constraints and DRAM scaling issues [19 [5]. it is hard scale up the GPU memory capacity on each server machine. Scaling out LLM training with a cluster of GPU servers can increase the aggregated memory capacity, however, it inevitably increases operational cost, placing major barrier for researchers and developers to entry into cutting-edge LLM development.
46
To overcome the GPU memory wall, prior studies proposed expanding GPU memory with external memory devices. One common approach to allow tensor offloading from GPU memory to host memory 17 39
20 However, due to the fundamental DRAM scalability challenge, such an approach is still limited by the host memory. Recent studies have extended tensor
32
41
43
offloading to PCle-based SSDs that offer larger capacity much lower cost M 37 16 23 47
51 29 [50]. But due to the incapability of efficiently utilizing the SSD bandwidth and hiding the slow SSD accesses, these offloading solutions still deliver suboptimal performance. For instance, ZeRO-infinity 337 enabled the offloading of tensors to SSDs at the granularity of deep neural network
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(DNN) layers, its coarse- -raiined fffloadinng/refetchiin scheme wastes not only the limited SSD bandwidth but also the precious GPU memory space, leaving the solution less attractive in practice.
Ideally, we wish to expand GPU memory with low-cost SSDs while achieving similar training performance as the ideal case assuming GPUs have unlimited on-board memory. To this end, we present a new tensor offloading framework TERAIO, which enables fine-grained offloading of tensors in an accurate fashion based on their activity patterns in GPUs, for best utilizing both SSD bandwidth and GPU memory. Our study of tensor activity patterns ( in LLM training shows (1) the active tensors, which are used in the current kernel during LLM training, consume only a small portion (1.7% on average) of requested GPU memory in total; (2) many inactive tensors are large and occupy a substantial GPU memory space in each training iteration; but (3) these inactive tensors are not used in the training for a long period of time, depending on the computation intensity of LLMs.
With these insights, TERATO develops lifetime-aware tensor offloading mechanism following three design principles: (1) offloading large inactive tensors to SSDs can save precious GPU memory and maximize SSD bandwidth utilization during the training process: (2) the distribution of their inactive periods of time will help TERAIO decide which inactive tensor should be offloaded at what time, and similarly, which tensor should be prefetched at what time; (3) precisely scheduling tensor offloading and prefetching in consideration of the available SSD bandwidth will help TERAIC effectively overlap the tensor movement with GPU computation. Our roofline model analysis (:2.2) shows that, given each GPU connected to multiple commodity SSDs today, the aggregated storage I/O bandwidth is sufficient to meet the tensor migration requirement without hurting the GPU training process.
To fulfill the design principles discussed above, we develop TERAIO with three major components: (1) tensor lifetime profiler that can extract tensor activity patterns (e.g., tensor size and lifetime) in advance with the assistance of deep learning compilers such as PyTorch, (2) lifetime-aware tensor migration algorithm that can generate optimal tensor offloading/prefethhing plans based on the learned tensor activity patterns, and (3) a tensor migration engine that will execute the generated herifaaiing/reeeccinnng plans with efficient direct data transfer between GPUs, host memory, and SSDs. We present each of these components as follows.
An open-source tensor lifetime profiler. TERAIO conducts the profiling of the tensor size and lifetime distributions by running the first few iterations of L.MM training on the target GPU setting. As the computation and dataflow patterns of each iteration almost the same, the profiling results can accurately represent the generic patterns of the entire LLM training process. To track the metadata information of each tensor, instrument the automatic operator generator in PyTorch rather than intrusively instrument the source code of each generated operator. Therefore, the proifler requires minimal code modifications to PyTorch. As the execution of LLM on GPUs has highly predictable dataflow patterns, TERAIO uses the execution time of GPU kernels to accurately estimate the tensor lifetime (i.e., the length of the active and inactive period of time). With the knowledge of tensor activity patterns, TERAIO will create the tensor offloxding/prefetching plans in advance.
Lifetime-aware tensor migration algorithm. TERAIO prioritizes offloading large tensors with long inactive period of time to SSDs, for fully utilizing the available storage I/O bandwidth. For tensors that have short inactive periods of time, TERAIO will make the best effort to retain them in GPU memory, for avoiding unnecessary migration overhead. As host memory and SSD offer different capacities, bandwidths, and costs, TERAIO prefers to offload tensors SSDs for taking advantage of their large capacity and low cost. However, when the SSD bandwidth is saturated at runtime, TERAIO will use the host memory the offloading destination. Given an execution plan for each LLM training iteration, TERAIO will iteratively search for the best offloading candidate based on the tensor size and lifetime. until the required GPU memory is below the capacity limit. After that, TERAIO will generate an optimized tensor migration plan by adding corresponding offloading and prefetching instructions in the compiled L.M training program.
Tensor migration engine using GPUDirect storage. Following the tensor migration plan integrated into the compiled LLM training program, the tenor migration engine of TERAIO will offload/prefetch tensors to/from SSDs or host memory at runtime. For the tensor migration between GPU memory and SSD, TERAIO uses GPUDirect storage enable direct data transfer between GPUs and SSDs, therefore, it can bypass the host CPU to alleviate the scalability bottleneck and maximize SSD bandwidth utilization. When the available SSD bandwidth is insufficient to support tensor offloading and prefetching, TERAIO will migrate tensors to the host memory. To track the latest locations of

---

tensors (GPU memory. host memory, or SSDs), TERAIO indexes tensors with their identification numbers using hash maps.
We implement the core components of TERAIO based on Py TyToch Therefore, TERATO does not require any code modifications LLM training programs. To evaluate the efficiency of TERATO, we train set of Llama and Granite models with different batch sizes and sequence lengths using TitehTitan [22 on a GPU server that has two NVIDIA HI00 GPUs and eight PCle-based SSDs. In comparison with state-of-the-art offloading solutions ZeRO-Offoad 40] and ZeRO-Infinity [37], TERAIO improves the training performance by 1.47x on average. achieves 80.7% of the ideal performance assuming unlimited GPU on on-boar memory, and delivers 1.45 improvement on cost efficiency for LLM training. In summary, we make the following contributions.
We conduct a quantitative characterization study of tensor memory usage when training different LLMs on multiple GPUs, and show that the high compute intensity of modern LLMs provide rich opportunities for tensor offloading.
We develop lightweight tensor lifetime profiler based on PyTorch, which can learn tensor activity patterns for multi-GPU LLM training.
We design lifetime-aware tensor migration planning algorithm that optimizes offload- ing/prefetching decisions based on tensor activity patterns, GPU memory capacity, and the available migration bandwidth.
We implement transparent tensor migration engine that enables direct data transfer between GPU and SSDs, alleviating the scalability bottleneck the host.
We conduct a thorough evaluation of TERAIO with the training of various LLMs, demon- strating significant improvement on training performance and cost efficiency, compared to oftte-of-theaan offloading solutions.
2 Characterization Study of Tensor Activity Patterns in LLM Training
In this section, we present our characterization study of tensor activity patterns in LLM training. To facilitate our study, we utilize our tensor lifetime profiler, which will be discussed in |3.1 to analyze the distributions of tensor sizes and lifetimes during the LLM training. We use two NVIDIA Hi00 GPUs and 2-stage 1f1b pipeline parallelism (8 in our experiments. Our study covers a variety of LLMs with different architectures, including decoder-only models such as Llama3-8B (batch size of 128, sequence length of 4096), Llama3- (batch size of 256, sequence length of 2048) [4, and GPT2-40B (batch size of 16, sequence length of 1024) 35 well encoder-decoder models like T5-11B (batch size of 64, sequence length of 512) 77] For models that require more memory than
GPU memory capacity, we offload tensors not needed by the current kernel SSDs. We summarize our study results as follows.
100%
100%
Memory Cocsumgtion
9.1%
10%
10000
23000
20008
40000
$3000
5000
100CS
15000
20006
Livred orUn thatch soe 125, tecuarson leegth
180 OPT -tB GPEO Botch coe 178, dequneee leegth 1804)
100%
0.1%
35900
100033
GPI Keravi
130000 200000
200000
30539
23900
33300
Shmmal
40000
50000
00000
Llmred 230-0700 thotch cae 64, requenoo inagth 2010)
00 TO-LLD-RUU (hatch cou 256, eaguanzo Imngth
(a) Memory consumption of different LLMs on each GPU.
Consumption
all
active
active
Monsry
100%
100%
10%
GPU Kernel Index (x) CIPT2-400--CPP (Slnge-0)
GPU Kernel Eadox Chi GWFT3--0----PUU (Stuge11)
(b) Memory consumption of different stages in the pipeline parallelism.
Figure 1: Memory consumption of all and active tensors (w.rt. the GPU memory capacity) in one iteration of the parallel training program. Logarithmic scale is used in the presentation.
3

---

100
100
100
100
00
TIEEER ILANE
80
80
TLAMN
80
CDP (16)
60
CDD (56)
60
CDF
40
40
4.0
40
20
20
20
20
10a
102
10
10a
10
10a
10a
10
10f
107
10a
102
10
10 10"
10"
16a
10f
107
10
Inactive Time (15)
Inactive Time (146)
Enactive Time (145)
Iractive Time (18)
Llama3-BB
(b)
GPT2-408
[c) Elama3-70B
TS-11B
Figure 2: The distribution of inactive periods of tensors.
2.1 Rich Opportunities for Tensor Migration
Small memory requirement of active tensors. We first study the memory demand and usage of tensors in each training iteration. We define the tensors that are currently used by a running GPU kernel as active tensors. We present the memory consumption of active tensors used by each GPU kernel in Figure | Figure (a) shows the memory consumption of tensors in different models, Figure (b) shows the memory consumption of tensors across different pipeline stages. For all LLMs examined in our study, the active tensors account for only less than 14% (1.7% on average) of the total GPU memory capacity, although their total memory usage greatly exceeds the GPU memory capacity. Most tensors in GPU memory are inactive and can be offloaded to low-cost SSDs, thus, we can best utilize the GPU memory for tensors that will be used by kernels in the near future.
Long inactive periods of inactive tensors. To understand how long the inactive tensors remain inactive and how much GPU memory they consume, we study the distribution of their inactive periods, as shown in Figure 2 For all the models we study, most tensors have sizes that range from 1OMB to IGB, We observe that more than 40% of these tensors remain inactive for more than 10a microseconds. These inactive periods are longer than the time needed to migrate these tensors to SSDs at a bandwidth of 6.5 GB/s. With this insight, we can ensure that these tensors can be migrated efficiently without introducing negative impact on the training performance.
Normelized Theoughpat
0.8
D Liama3-4a 9B
Liama3 700
0.6
TS-118
GPT2-400 40B
0.4
0.2
0.0
1216
24 32
48 64
Available Migration Bandwidth (CBS)
Figure 3: Roofline analysis with different migration bandwidths. The training performance nor- malized to the ideal case assum- ing GPU memory is infinite.
The long inactive periods are the cause of the sparse tensor access pattern and high compute intensity of LLMs. From spatial perspective, although LLMs have tens or hundreds of layers, many tensors are used within only a single layer. From a temporal perspective, the compute- intensive kernels (e.g., attention) in each layer take a considerable amount of time, providing rich opportunities for TERAIO to overlap the computation with the migration of inactive tensors.
We also observe that, compared with traditional DNN models, the higher compute intensity and larger model sizes of LLMs lead to substantially longer inactive periods. For example, in BERT-Large (3). 48% of tensors are larger than 100MB, and the inactive periods of more than 60% of these large tensors are two orders of magnitude shorter than those in LLMs.
2.2 Bandwidth Requirement for Tensor Migration
We now study how much migration bandwidth is needed for offloading to achieve near-ideal training performance. We quantify the roofline performance of different LLM models under different migra- tion bandwidths available to each GPU. To facilitate this study, we build performance model to estimate the training time. In the model, we assume that each kernel's execution time is the same as the value collected in our characterization study. We simulate tensor migration at the designated bandwidth, and check whether the tensors needed by the kernel are already GPU memory or not. If they are still being migrated due to limited SSD bandwidth, the waiting time for the migration is added to the total training time. Figure shows the normalized roofline training throughput of LLMs under different migration bandwidths. We observe that bandwidth of 32 to 48 GB/s is sufficient to achieve near-ideal performance for L.M.s. Such a bandwidth requirement can be easily achieved by aggregating multiple commodity SSDs (c.g., an SSD array), demonstrating the feasibility of TERAIO.
4

---

Lifetime-Aware Tensor Migration Algorithm (89.2
Per-PPU Tersor Migation Plats
Oye
start bamel dcadlarkrrat
GIFU Mamrrry
Office
Esccution Groph Trasing
Tensor Lifeiins information
GPU
vxage
IB
dise
kaads
rattino
645
Location Traoking Thble
nigunfans
12400s
GPU Karnel Phoflling
berrel AAK
Insurtt
G/U Man
CPU May
Tensor Lifetime Profiler ($3.1)
PyTorch Framework
GDS-Based Tensor Migration Engine ($3.3)
bain Mamao
Figure 4: System overview of TERAIO.
3 TERAIO Design and Implementation
We show the system overview of TERAIO in Figure 4 Given an LLM, TERAIO's tensor lifetime profiler works with PyTorch to track tensor sizes and lifetimes ($3.T) In the first few training iterations, TERAIO traces the execution graph and collects the execution time of each kernel. Since the training follows the same execution graph in subsequent iterations, the tensor activity patterns remain the same. The tensor migration algorithm ($3.2) creates a tensor migration plan that (1) maximally overlaps computation and migration, and (2) minimizes migration traffic. The algorithm iteratively selects the best offloading candidates until the required GPU memory fits within the actual GPU memory capacity. For the migration destination, it prefers to migrate tensors to SSDs. Once the SSD bandwidth is saturated, it also uses available CPU memory. During LLM training, TERAIO's tensor migration engine transparently executes the migration plan ($3.3).
3.1 Tensor Lifetime Profiler
Tracking tensors. TERAIO instruments PyTorch framework to track tensors and measure kernel exccution time at runtime. A tensor is considered active in one of the following three scenarios.
First, a tensor is active when it is the input or output of PyTorch CUDA operator. However, instrumenting Py PyTorch to track every operator is challenging, as there are thousands of operators. Instead, we leverage PyTorch's 's automatic operator generator, which produces source code for each operator, to insert profiling code that will mark all input and output tensors as active when the operator is executed at runtime. Second, for tensors that are involved in inter-GPU communication, they should be active in GPU memory. Third, a tensor is considered active when Py explicitly checks whether it resides in GPU memory. This happens when updating optimizer states, For the second and third scenarios, since there are only a few communication operators and PyTorch checks in total, we directly set the corresponding tensors as active in the source code.
Operator
IPU Communiciioo
PyTorch Creck
Perkod
Active
Inoctive Parlod
Dealiccated
Runtime Tensor Lifetime Tracking
Operator
Comm
Openstor
Check o
Openatar 2
Tensor Lifetime Table
Exacution
Flow
Tensor
Tensor
Terwor
Tensor
Tensor 0.1.2
Update
Tensor
active periods
global
Tonsor
(po1) op2)
False
12BMB
Tansor 2
64MB
True
Figure TERAIO tracks and analyzes tensor activity patterns for the tensor migration algorithm.
Analyzing tensor activity patterns. Figure 5 shows how tensor information is collected at runtime. To understand when a tensor consumes GPU memory and when it must reside in GPU memory, we need collect its tensor size and active time. When an operator is executed, the instrumented code records the tensor size and active time for the corresponding tensor. The profiler will calculate the inactive time period based the duration between its active states. Specifically, for intermediate tensors such as gradients and activations (Tensor 0 and Tensor in F)g that will be deallocated immediately after its computation completes, quantify its inactive time period the time interval between the two active periods. For global tensors such as model weights and optimizer states (Tensor 2 in Figure 5) that are used across multiple training iterations, they are allocated before training starts and never deallocated during training. Therefore, for some cases, the profiler may need to calculate its inactive period based on the active states across two iterations.
5

---

IO bar dwidth salureted by prev ouslyy planned lansor migralions
Texa arsount of required GPU mamory
Toral mey GPU memor afte offood
Io bandwict that be used offoading tensor
Denel can get from otfoading the barsor
Flaqy aned CPU Menurrry
Tansor Inactive Paried
Feed
PO Bendwitth Usage
ofoaaing
Memory Capardly
Prefetching
Tytert
fe/foot
foffowent
sme
Eaf/foaded
Untter
(a). 1/O bandwidth usage aflects actual migratiio time
(b). The benelit of otlloading one inactive lensor
Figure 6: Illustration examples to explain the key insights of lifetime-aware migration algorithm.
3.2 Lifetime-aware Tensor Migration Algorithm
The lifetime-aware tensor migration algorithm iteratively finds the hest offloading candidates in each LLM training iteration, until the required GPU memory is below the capacity limit. By tracking the amount of required GPU memory and the storage 1/O bandwidth utilization, the algorithm is able to evaluate the potential benefits of tensor offloading. We discuss our key ideas as follows.
Storage I/O bandwidth-aware planning. To offload an inactive tensor, we wish to keep it out of GPU memory long as possible. Therefore, ideally, would offload it as soon as it becomes inactive and prefetch instantly before it is needed by the subsequent kernel. However, in reality, the time when can offload and prefetch the inactive tensor is greatly affected by the storage bandwidth usage. For example, in Figure (a). the inactive period of the tensor is from Lstart to lend- However, since the 1/O bandwidth occupied by previously planned migrations, we have to delay the actual offloading until fo ftoad and start the actual prefetching by feref This means that we can only reduce the memory consumption from fo foaded to tyrefetch. When planning tensor migrations, our algorithm tracks the estimated storage I/O bandwidth usage and calculates the reduction in memory consumption in an 1/O-aware manner (Line 4-5 and 16-17 in Algorithm
Algorithm 1 LifetimeeAware Tensor Migration Planning
Require: tensor inactive periods = {(t,, dss start, end,)) } where 1, tensor ID. e; size, and [start, endd]] defines the kemel raago of inactivity: GPU memery capacity Maru; Estimation total amounn of required GPU memory M [mo. 113
across N kemnels; Kermel execution times T = [To: Th...., TN Tw-1]: 1]: 1OO bandwidth usage states {((,, trigger_time, deadline, target))
Ensure: Migration plan
1: Initialize migration plan list P = {} and estimation of required GPU memory M' = M
2: while peak memory consumption max (M') > McPu
3: for each inactive period (f, s, start, end) E do
4: Determine carliest feasible offload completion time floaded by analyzing available offloading bandwidth
5:
Determine latest required prefetch start time fprefetch by analyzing available prefetching handwidth
6:
fof fleuded < pre fetck
7:
Identify critical memory pressure rogions C (A|M'( > Maro, & [toffloaded. (prefetch)) Calculate the offloading benefit B Σxeo end
8:
9:
10:
end for
11:
Select inactive period best offleading cost-benefit (B/ smnce
12:
if viable offloading candidate exists then
13:
break end if
14:
15:
Add migration plans (t. start, flonded SSD/CPU) and (f, &prefetoh Sprefeteiked GPU) to
16:
Update A by subtracting tensor size a from affected kernels
17:
Update IOO bandwidth usage states
18:
Remave selected inactive period from consideration
19:
end while
20:
return P
Quantify the benefit and cost of tensor offloading. To fully utilize the available I/O bandwidth, we want to prioritize offloading large tensors with long inactive periods. Following this principle, our algorithm searches for the best offloading candidate by estimating its benefit and cost. To quantify the benefit, at a given time T, we define critical memory pressure as the part of GPU memory consumption that exceeds the capacity. The benefit of a tensor migration is defined as the integral of the reduction in critical memory pressure over time, as illustrated by the shaded area in Figure b). We quantify the cost as the sum of offloading and prefetching time of tensors. TERAIO's migration algorithm sorts all candidates by their benefit-to-cost ratio. It then selects the tensor with the highest ratio for migration in the current iteration. This procedure is shown in Line 6- 11 in Algorithm

---

Decide offloading destination. TERAIO prioritizes SSDs because of their large capacity and low cost. When the estimated SSD bandwidth is saturated, TERAIO will make the best effort migrate tensors to the host memory. TERATO allows users to define the maximum amount of host memory that can be used for tensor migration. Internally, the migration algorithm tracks the estimated host memory consumption. If it has already reached the user-defined limit, even if SSD bandwidth is saturated, TERAIO will not offload tensors to the host memory.
Minimize kernel stalls. If the required GPU memory exceeds the available GPU memory capacity, and TERAIO cannot offload tensors to SSDs or host memory, TERAIO will stall kernel execution to wait for more inactive tensors to be offloaded. It will also wait for the tensors needed by the kernel to be migrated back into GPU memory. This would cause storage I/O bandwidth contention. Since the next kernel will stall until the needed tensors are migrated to GPU memory, these migrations should complete as early as possible to avoid stalling the GPU training process. Therefore, TERATE marks these migrations that must finish before the next kernel as 'urgent'. At runtime, the tensor migration engine ($3.3) always prioritizes these urgent migrations over other pending migrations.
Compilation cost. The overhead of generating tensor migration plans based on the lifetime analysis is shown in Table Although this com- pilation step introduces additional latency, its cost is less significant in practice because enables efficient tensor offloading, which reduces overall model training time. After the plan generation, we do not need to recompile ML models, as the corresponding offloading and prefetching instructions have been integrated into the execution graph. The details are described in Section 3.3
Table 1: Compilation time for different models.
Model
Timeis)
Liana'
Cramle-code- base 30
37.9
Liunas 708
396.6
3.3 Tensor Migration Engine Using GPUDirect Storage
To execute the migration plan, we need to locate the addresses of the tensors to be migrated, based on the tensor identifier. TERAIO maintains a hashmap-based tensor location table in Py PyTorch map tensor identifiers to their current devices and addresses.
The migration engine migrates tensors between GPU memory and external memory. For tensor migrations between GPU memory and SSDs, we use GPUDirect storage to achieve direct data transfer between them. Our choice of GPUDirect storage is motivated by the potential scalability limitations of host CPU in multi-GPU systems. For example, on an 8-GPU system, to achieve near-ideal performance, each GPU requires 32 48 GB/s bidirectional bandwidth (see $2). With GPUDirect storage, we can directly connect SSDs to each GPU via PCIe Gen5 switches to meet the bandwidth demand. For the conventional approach that uses the host CPU to first read data from SSDs to the host memory, and then uses cudaMemcpy to move data to the GPU, the redundant data copy not only causes performance overhead but also wastes precious CPU cycles 33 But as discussed, TERAIO still supports migration between GPUs and CPU memory.
4 Experiments
We show that (1) TERAIO outperforms state-of-the-un offloading frameworks by 1.47x on average when training LLMS that greatly exceed GPU memory capacity ($4.2); (2) Compared to the case of training LLMS using only GPU memory, TERAIO reduces the cost by up to 5.41x ($4.3): (3) Compared to existing offloading frameworks, TERAIO improves the cost efficiency by 1.45x on average ($4.);; (4) TERAIO achieves better throughput than ZeRO-Infinity even with less CPU memory and fewer SSDs ($4.4).
4.1 Experimental Setup
Models. We evaluate TERAIO with Llama3-8B, Llama3-70B 4 and Granite-code-bar 8B 13 We use C4 (36 as our training dataset. To study how different memory demands impact the performance of ATRA we use batch sizes ranging from 16 to 128 and sequence lengths from 1,024 to 8,192.
Hardware configuration and MI framework. Table 2 shows the hardware configuration used in our experiments. Due to the limited PCIe slots on our machine, we can only install at most 8 PCIe SSDs. When evaluating TERAIO, we use H100 GPUs and 2 RAID-O arrays with 4 SSDs in each array. Each RAID-O array is logically assigned to one GPU, providing approximately 16 GB/s bandwidth for tensor migrations. We use PyTorch 2.5.0 30 and TorchTitan 22 train LLMs.
7

---

Baselines. We compare TERAIO with the Ideal case,
Table 2: Our GPU server configuration.
ZeRO- OffOOOff and ZeRO-Infnity. The Ideal case assumes
GPU
1x NVIDIA NVL.
that all GPUs in the system have infinite on-board mem-
GPU Momorr
saGB HBM per GPU
ory, which gives the theoretical best training performance.
CPU
2x AMD EPYC 93M
Memury
ZeRO-Offload [9) and ZeRO-Infinity [7] are popular offloading- based training systems. ZeRO-Offload offloads tensors from GPU memory only to CPU memory, while
Intervenect
PCle Gen'
8SDA
XX Samning" 1i0 PRO 2T8
SSD KentWoite Handwiddti
[6,76,3 G&S per SSD
ZeRO-Infinity leverages both SSDs and CPU memory to expand GPU memory. TERAIO-SSD and TERAIO-Mixed are two variations of TERAIO. TERAIO-SSN only migrates tensors to low-cost SSDs, while TERAIO-Mixe TERAIO-Mixed use the same amount of CPU memory for tensor migration ZeRO, we preserve most of its default parameters and fine-tune several parameters to maximize its throughput. We list its detailed configurations in Table 4 [Appendix] B).
uses both SSDs and CPU memory. To make a fair comparison, we let ZeRO-Infiniz For
In our evaluation, we aim to compare TERAIO with the best performance achievable by ZeRO. Therefore, for the parallelization strategy, use tensor parallelism for ZeRO series, it delivers optimal multi-GPU training performance. Moreover, we split each batch into multiple micro-batches in order to amortize the well-known performance bottleneck [16] of ZeRO's CPU-based optimizers. In addition, although activation checkpointing can reduce memory consumption, we disable it because it degrades ZeRO's training throughput.
In terms of training precision, we use cell-precision training in all experiments. Though mixed- precision training [26] is popular, the different memory requirements of the mixed-precision training strategies used by TorchTitan and ZeRO will lead an unfair comparison. Specifically, when we enable mixed-predd preccippe training, in the ZeRO series, all tensors in the GPU are represented by 16-bit floating points, while in TorchTitan, most tensors, including model weights, gradients, and optimizer states still remain in 32-bit floating point format. Such differences in numerical formats of tensors can lead to significant differences in GPU memory requirements for the same model, resulting in an unfair scenario where TERAIO has to migrate larger tensors than ZeRO
4.2 End-to-end Performance
We show the end-to-end average training throughput of Llama3-8B, Granite-code-basee and Llama3-70B with different batch sizes and sequence lengths in Figure 7 On average, TERAIO outperforms ZeRO-Offload and ZeRO-Infinity by 1.47x. Compared to the ideal system assuming unlimited GPU memory, TERAIO achieves 80.7% of the ideal performance.
Training Throughput
ZeRO-Infiilty Llame3-83
ZeRO-Offload
TeralO-SSD Granite-code- beatse--dde
TeralO-Mixed
Llama3-708
1500
200
(token/s) (9
1000
300
bs32
9000043
bus4 Pig4045
2s128 a0q2046
1602 5006
Iui44 3604000
ba128 s00000g
baB 5002040
2516 SeN2048
1632 se22048
(1). Average training throughput with different batch sizes.
Training Throughput
ZeRO Infinity Llama' 8B
2cRO Offload
TeralO SSD Granite code bose BD
TeralO Mixed
Ideal
Llama33 70D
2028
(tokents)
bus4
bubt 8606006
1004 3MH22043
2010 ANNASSH 2010 HMNNU72 1618 AMQSTNUG
168.81%
381
M
350 80%
1242 75%
(2). Average training throughput with different sequence lengths.
Figure 7: Average training throughput of Llama3-8B. Cranite-code-base-- different batch sizes and sequence lengths. bs the batch size. seq is the sequence length. M the overall peak memory consumption of the LLM training on one GPU w.r.t. the GPU memory capacity. "x" " means the framework failed to train this model due to out-of-memory errors.
and Llama3-70B with
Training throughput. As shown in Figure 7 when training Llama3-8B and Granite-8B. ZeRO- Offload achieves 65.9% and 65.5% of the ideal performance, respectively. Though ZeRO-infinit takes extra time to migrate tensors further from CPU memory to SSDs, its training throughput
8

---

is slightly lower than ZeRO-Offloa.. Such a subtle performance difference comes from the high aggregate bandwidth of 8 SSDs and the relatively small sizes of optimizers and parameters offloaded to SSDs, For Llama-70B, ZeRO-Ofiii -ffo--fffodd fails because the host memory capacity is too small to store offloaded tensors of such a large model. Since SSDs offer larger capacity, ZeRO-Infinity can still train Llama3-70B. However, it only achieves 43.0% of the ideal performance, because its coarse-gen -ggaiind offloading scheme cannot efficiently utilize the limited SSD bandwidth to migrate larger tensors.
TERAIO outperforms the ZeRO series by up to x.59x. For Llama3-8B, both TERAIO-Mixed and TERAIO-SSD can achieve near-ideal performance, demonstrating the effectiveness of our lifetime- aware tensor migration algorithm in choosing the most beneficial tensor to migrate. Moreover, we find that even though 2 H100 GPUs can provide sufficient GPU memory to train Llama3-8B without memory expansion, TERAIO allows us to increase the (micro-Jbatch size, improving the throughput by 9%. For Granite, TERATO-SSD achieves similar performance to the ZeRO series, while TERAIO-Mixen can still deliver near-ideal performance by utilizing CPU memory. The result shows that as the model size increases, achieving near-ideal performance requires not only our lifetime-aware tensor migration algorithm, but also migration bandwidth higher than what 4 SSDs can provide. For Llama-70B, since its memory requirement significantly exceeds GPU memory capacity, higher migration bandwidth is needed. Even TERAIO-Mixed can only achieve 59.7% of the ideal performance. Nonetheless, it still outperforms SeRO--nninity by 1.33x.
Impact of varying batch size and sequence length. As batch size and sequence length vary, the training throughput of TERAIO-Mixee Mixed is always the closest to the ideal throughput among all offloading frameworks, while TER AIO-SSD delivers similar or better performance than the 7eRO series. For the same model, increasing the batch size doesn't raise memory requirements, since each batch is split into micro-batches that contain the same number of training samples. However, when the sequence length increases, the memory requirement also increases. Therefore, when the sequence length is longer than 3,072, ZeRO-Infinity fails to train the model since the GPU memory capacity is smaller than the activation tensors of the model. In contrast, TERAIO can train all model configurations because it can offload any inactive tensor to save memory.
Breakdown of training time. To better understand the source of TERAIO's performance gains, we also show the breakdown of the end-to-en training time in Figure g 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.
Kernel %
Overlap %
Compotation Only
Llama3-88 (bs32 seq2048)
Cranite-88 (bs32,ang400)
1aama3-70r (bs8. seg204))
ZoRO Infinit-
ZeRO Infinity
ZeRO Infinity
TeralO-SSD
TeraIO-SSD
TeraIO-SDD
TeraIO- Mixed
ToratO Mixed
TeralO-Miae
066
2016:
40%
60%
80%
100%
0%%
20%
40%
80%
100%
0%
20%
40%
10%
100%
Percentage (56)
Percentage (%)
Percestage (1)
Figure 8: Latency breakdown of training iterations.
As shown in Figure 8 compared to InROntifintt, TERAIO-Mixed has less stall time, as it maximizes the overlap of tensor migrations with computation. TERAIO-Mixed outperforms ZeRO- -nfiniiy for the design differences in offloading granularity and migration strategy. As for offloading granularity. ZeRO-Infinity selects the offloading of tensors at the layer granularity ML models. This introduces burst I/O patterns and can underutilize migration bandwidth. In contrast, TERATO selects tensor offloading at the GPU kernel granularity based on our precise tensor lifetime analysis. This approach ensures high and consistent migration bandwidth utilization, as shown in Figure I0 (Appendix C). As for 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 1/O-aware planning (Section 3.2). which generates globally optimized migration plans that maximize I/O bandwidth utilizzation.

---

4.3 Training Cost
To evaluate the cost of TERATO, we sum- marize the prices of different devices used in each base- line setup in Ta- ble[ TERAIO-SSD needs a server with
Table 3: Cost of each device used by TERAIO and other baselines for Llama3- 70B training. Prices are quoted from Exxact [5] The PureGPUs setup contains the minimum number of 8-GPU H100 servers capable for training L 70B without offloading tensors to host memory or SSDs.
Server (wit CPUS
Server HIOGGPUS)
HIDSOGUU)
Sammung
memery
mustrry
91,04299
Traatt--sn
Taaal Mixed
only 128GB of CPU memory and 8 SSDs
Z:RO-Ollea
Zel00-Inireiy
PantiKii
since it migrates ten-
sors only to SSDs, while TER MIO-Mixdd uses both SSDs and ITB CPU memory for more efficient tensor migration. Since 7ZeRO-Offlod and ZeRO-Infinity consume large amount of CPU memory to train LLMs, they both need a server with 1TB of memory. ZeRO-Infinity additionally uses SSDs in evaluation. We also compare with the PureGPU setup, in which all tensors are kept within GPU memory. To provide enough GPU memory to train Llama3-70B in this setup, we need to pay 1499,591.40 for two servers, each equipped with 8 H100 GPUs. In comparison, TERAIO-SSD and TER ATo-Mixdd save costs by 5.88x x and 5,41x, respectively. Compared to ZeRO-Offload and ZeRO ZeRO-Infinity, since we have similar machine setups, TERAIO's 's 1.47 training performance improvement translates into 1.45x and 1.47 x improved cost efficiency, respectively.
4.4 Impact of Varying Number of SSDs and CPU Memory Capacity
Figure 9 shows the training throughput of TERAIO as we vary the available CPU
Training Throughpat.
120
1800
100
1500
memory capacity and the number of SSDs
1200
80
used for each GPU. The performance of TERAIO scales favorably the number of SSDs or the CPU memory capacity. For Llama-8B, TERAIO achieves near ideal performance with only 2 SSDs and 64 GB of CPU memory. Even for Llama-70B, we observe that using only 2 SSDs and 512GB of CPU memory still achieve 73.1% 1% of the training throughput obtained with 4 SSDs and 1, 1,024GB of CPU memory. These results validate our LLM characterization study's 's key observation: with lifetime-aware tensor offloading, we can achieve good performance even with limited hardware resources.
900
2 530 Carato)
60
-81 2 5504 (Dima))
-0- 3 550x (analD)
40
-C- 3 552h (Rena)))
we increase
300
* ZettO-tatretty
4 55Ds (Canal))
20-
* ZellD-itteily
4 8530 (Tena)))
16 32
64
128
128
512
1024
Avaliable CPU Momory Capanity (GB)
Anadiable CPU Momory Copacity (C3)
(al Lluma3-88
(b) Llama3-706
Figure 9: Training throughput as we vary the CPU memory capacity and the number of SSDs per GPU.
Compared to ZeRO-Infinity, TERAIO significantly reduces the CPU memory capacity requirement while achieving better performance. ZeRO-Infinity requires 170GB and 770GB of CPU memory to train Llama-8B and Llama-70B: respectively, as it has to offload gradients and optimizers into CPU. In comparison, TERAIO achieves better performance even with more limited hardware resources, as it efficiently utilizes both the large capacity of SSDs and extra bandwidth of CPU memory.
5 Conclusion
We present TERAIO, a lifetime-aware tensor offloading framework that can accurately plan and execute fine-grained tensor offloading and prefetching instructions for LLM training. With predictable tensor activity patterns, TERAIO best utilizes precious GPU memory for accelerating GPU training process, while leveraging large-capacin SSDs for lowering training cost. Compared to existing tensor offloading work, TERAIO provides a more practical and cost-efficient solution for LLM training.
Acknowledgment
We thank the anonymous reviewers for their insightful feedback. We thank the members in the Systems Platform Research Group (Illinois PlatformX) at UTUC for constructive discussions. This work was partially supported by the Hybrid Cloud and AI program at the IBM-Illinois Discovery Accelerator Institute (IIDAD), and NSF under the grants CAREER CNS-2144799 and CCF-2107470.
10

---

References
[I] Jonghyun Bae, Jongsung Lee, Yunho Jin, Sam Son, Shine Kim, Hakbeom Jang, Tae Jun Ham, and Jae W Lee. (FlashNeuron): (SSD-Enabed] [Large-Batch] training of very deep neural networks. In 19th USENIX Conference on File and Storage Technologies (FAST 21), pages 387-401, 2021.
[2] Olivier Beaumont, Lionel Eyraud-Dubois, and Alena Shilova. Efficient combination of remateri- alization and offloading for training dnns. Advances Neural Information Processing Systems, 84:23844-23857, 2021.
[3] Jacob Devlin, Ming-War Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[4] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv preprint arXi::240.221783 2024.
[5] EXXACT. Exxact.
2025.
[6] Talia Gershon, Seetharami Seelam, Brian Belgodere, Milton Bonilla, Lan Hoang, Danny Barnett, I-Hsin Chung, Apoorve Mohan, Ming-Hung Chen, Lixiang Luo, Robert Walkup, Constantinos Evangelinos, Shweta Salaria, Marc Dombrowa, Yoonho Park, Apo Kayi, Liran Schour, Alim Alim, Ali Sydney, Pavios Maniotis, Laurent Schares, Bernard Metzler, Bengi Karacali-Akyamal Sophia Wen, Tatsuhiro Chiba, Sunyanan Choochotkaew, Takeshi Yoshimura, Claudia Misale, Tonia Elengikal, Kevin O Connor, Zhuoran Liu, Richard Molina, Lars Schneidenbach, James Caden, Christopher Laibinis, Carlos Fonseca, Vasily Tarasov, Swaminathan Sundararaman, Frank Schmuck, Scott Guthridge, Jeremy Cohn, Mare Eshel, Paul Muench, Runyu Liu, William Pointer, Drew Wyskida, Bob Krull, Ray Rose, Brent Wolfe, William Cornejo, John Walter, Colm Malone, Clifford Perucci, Frank Franco, Nigel Hinds, Bob Calio, Pavel Druyan, Robert Kilduff, John Kienle, Connor McStay, Andrew Figueroa, Matthew Connolly, Edie Fost, Gina Roma, Jake Fonseca, Ido Levy, Michele Payne, Ryan Schenkel, Amir Malki, Lion Schneider, Aniruddha Narkhede, Shekeba Moshref, Alexandra Kisin, Olga Dodin, Bill Rippon, Henry Wrieth, John Ganci, Johnny Colino, Donna Habeger- Rose, Rakesh Pandey, Aditya Gidh, Aditya Gaur, Dennis Patterson, Samsuddin Salmani, Rambilas Varma, Rumana Rumana, Shubham Sharma, Aditya Gaur, Mayank Mishra, Rameswar Panda, Aditya Prasad, Mat Stallone, Gaoyuan Zhang, Yikang Shen, David Cox, Ruchir Puri, Dakshi Agrawal, Drew Thorstensen, Joel Belog, Brent Tang, Saurabh Kumar Gupta, Amitabha Biswas, Anup Maheshwari, Eran Gampel, Jason Van Patten, Matthew Runion, Sai Kaki, Yigal Bogin, Brian Reitz. Steve Pritko. Shahan Najam. Surya Nambala, Radhika Chirra, Rick Welp, Frank DiMitri, Felipe Telles, Amilear Arvelo, King Chu, Ed Seminaro, Andrew Schram, Felix Eickhoff, William Hanson, Eric Mckeever, Michael Light, Dinakaran Joseph, Piyush Chaudhary, Piyush Shivam, Puneet Chaudhary, Wesley Jones, Robert Guthrie. Chris Bostic, Rezaul Islam, Steve Duersch, Wayne Sawdon, John Lewars, Matthew Klos, Michael Spriggs, Bill McMillan, George Gao, Ashish Kamra, Gaurav Singh, Marc Curry, Tushar Katarki, Joe Talerico, Zenghui Shi, Sai Sindhur Malleni, and Erwan Gallen. The infrastructure powering ibm's gen ai model development, 2025.
[7] Google. T5 11b.
2025.
[8] GPUDirect Storage: A Direct Path Between Storage and GPU Memory. nvidia. wtt.:/ooo/oooooiiiiiiioooo http://lloooppuirreeeetttstrrgre
[9] Mark Hildebrand, Jawad Khan, Sanjeev Trika, Jason Lowe-Power, and Venkatesh Akella. Autotm: Automatic tensor movement heterogeneous memory systems using integer linear programming, In Proceedings of the Twenty- Fifth International Conference on Architectural Support for Programming Languages and Operating Systems, pages 875-890 890. 2020.
[10] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106. 09685, 2021.
11

---

[11] Chien-Chin Huang, Gu Jin, and Jinyang Li. Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. In Proceedings of the Twenty-Fifte International Conference on Architectural Support for Programming Languages and Operating Systems, pages 1341-1355, 2020.
[12] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
[13] IBM. Ibm granite. https://hugginggacee http://.........eoo///////// cwt..m-eociimme/ee granite 2025.
[14] Joseph Izraelevitz, Jian Yang, Lu Zhang, Juno Kim, Xiao Liu, Amirsaman Memaripour, Yun Joon Soh, Zixuan Wang, Yi Xu, Subramanya R Dulloor, et al. Basic performance mea- surements of the intel optane de persistent memory module. arXiv preprint arXiv:1903.05714 2019.
[15] 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.
[16] Hongsun Jang. Jaeyong Song, Jaewon Jung. Jaeyoung Park, Youngsok Kim, and Jinho Lee. Smart-infinity: Fast large language model training using near-storage processing real system. In 2024 IEEE International Symposium on Perf--Peroormncc Computer Architecture (HPCA), pages 345-360. IEEE, 2024.
[17] Jachoon Jung, Jinpyo Kim, and Jaejin Lee. Deepum: Tensor migration and prefetching in unified memory. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pages 207-221, 2023.
[18] 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 preprint TTt::2001.08361 2020.
[19] Youngeun Kwon and Minsoo Rhu. Beyond the memory wall: A case for memory-centric hpc system for deep learning. In Proceedings of the 51st Annual IEEE/ACM International Symposium on Microarchitecture (MICRO'18) 18), Fukuoka, Japan, 2018
[20] Tung D Le, Haruki Imai, Yasushi Negishi, and Kiyokuni Kawachiya. Tilms: Large model support in tensorflow by graph rewriting. arXiv preprint
2018.
[21] Shaobo Li, Yirui Eric Zhou, Yuqi Xue, Yuan Xu, and Jian Huang. Managing scalable direct storage accesses for gpus with gofs. In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles. SOSP '25, page 979-995, New York, NY, USA, 2025. Association for Computing Machinery.
[22] Wanchao Liang, Tianyu Liu, Less Wright, Will Constable, Andrew Gu, Chien-Chin Huang, Iris Zhang, Wei Feng, Howard Huang, Junjie Wang, et al. Torchtitan: One-stop pytorch native solution for production ready I1m pre-training. arXiv preprint arXiv:2410.06511, 2024.
[23] Changyue Liao, Mo Sun, Zihan Yang, Kaiqi Chen, Binhang Yuan, Fei Wu, and Zeke Wang. Adding nvme ssds to enable and accelerate 100b model fine-tuning on a single gpu. arXiv preprint arXiv:2403.06500 06504 2024.
[24] Jonas Markussen, Lars Bjørlykke Kristiansen, Päl Halvorsen, Halvor Kielland-Ge Gyruu, Häkon Kvale Stensland, and Carsten Griwodz. Smartio: Zero-overhead device sharing through pcie networking. ACM Transactions on Computer Systems, 38(1-2), jul 2021.
[25] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. arXiv preprint arXiv: 18/2.0666 2018.
[26] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, al. Mixed precision training. arXiv preprint arXiv: 2017.
12

---

[27] Apoorve Mohan, Robert Walkup, Bengi Karacali, Ming-hung Chen, Abdullah Kayi, Liran Schour, Shweta Salaria, Sophia Wen, I-hsin Chung, Abdul Alim, Constantinos Evangelinos, Lixiang Luo, Marc Dombrowa, Laurent Schares, Ali Sydney, Pavlos Maniotis, Sandhya Kotesh- wara, Brent Tang, Joel Belog, Rei Odaira, Vasily Tarasov, Eran Gampel, Drew Thorstensen, Talia Gershon, and Seetharami Seelam. Vela: A Virtualized LLM Training System with GPU Direct RoCE, page 1348-1364. Association for Computing Machinery, New York, NY, USA, 2025.
[28] 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.
[29] Xiaonan Nie, Yi Liu, Fangcheng Fu, Jinbao Xue, Dian Jiao, Xupeng Miao, Yangyu Tao, and Bin Cui. Angel-ptm: A scalable and economical large-scale pre-training system in tencent. arXiv preprint arXiv:2303.08 02806888,, 2023.
[30] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS-W W, 2017.
[31] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4200 2023.
[32] Bharadwaj Pudipeddi, Maral Mesmakhosroshahi, Jinwen Xi, and Sujeeth Bharadwaj. Training large neural networks with constant memory using new execution algorithm. arXiv preprint arXiv:2002.0556 2020.
[33] Zaid Qureshi, Vikram Sharma Mailthody, Isaac Gelado, Seungwon Min, Amna Masood, Jeong- min Park, Jinjun Xiong, C. J. Newburn, Dmitri Vainbrand, I-Hsin Chung, Michael Garland, William Dally, and Wen-mei Hwu. Gpu-initiated on-demand high-throughput storage access in the bam system architecture. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, ASPLOS 2023, page 325-339, New York, NY, USA, 2023. Association for Computing Machinery.
[34] Zaid Qureshi, Vikram Sharma Mailthody, Isaac Gelado, Seungwon Min, Amna Masood, Jeong- min Park, Jinjun Xiong, Chris Newburn, Dmitri Vainbrand, I-Hsin Chung, et al. Gpu-initiated on-demand throuhhruuhhput storage access in the bam system architecture. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pages 325-339, 2023.
[35] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 18):9, 2019.
[36] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21((40):1-67 67, 2020.
[37] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero- infinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, St. Louis, Missouri, 2021.
[38] Jie Ren, Jiaolin Luo, Kai Wu, Minjia Zhang, Hyeran Jeon, and Dong Li. Sentinel: Efficient tensor migration and allocation on heterogeneous memory systems for deep learning. In 2021 IEEE International Symposium on High-Performant Computer Architecture (HPCA), pages 598 6.1. IEBE, 2021.
[39] 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. In 202/ USENIX Annual Technical Conference (USENIX ATC 21), pages 551-564, 2021.
13

---

[40] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang. Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scai model training. CoRR, 1066400 2021.
[41] Minsoo Rhu, Natalia Gimelshein, Jason Clemons, Arslan Zulfiqar, and Stephen W Keckler. vdnn: Virtualized deep neural networks for scalable, memory-efficiet neural network design. In 2016 49th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), pages 1-13. IEEE, 2016.
[42] Manish Shetty, Yinfang Chen, Gagan Somashekar, Minghua Ma. Yogesh Simmhan, Xuchao Zhang, Jonathan Mace, Dax Vandevoorde, Pedro Las-Casas, Shachee Mishra Gupta, al. Building ai agents for autonomous clouds: Challenges and design principles. In Proceedings of the 2024 ACM Symposium on Cloud Computing, pages 99-110, 2024.
[43] 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.
[44] Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song, Zenglin Xu, and Tim Kraska. Supereuuons: 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.
[45] Bingyang Wu, Ruidong Zhu, Zili Zhang, Peng Sun, Xuanzhe Liu, and Xin Jin. {dLoRA}: Dynamically orchestrating requests and adapters for [LoRA) [LLM} serving. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 91 1-979 2024.
[46] Duo WU, Xianda Wang, Yaqi Qiao, Zhi Wang, Junchen Jiang, Shuguang Cui, and Fangxin Wang. Netllm: Adapting large language models for networking. In Proceedings of the ACM SIGCOMM 2024 Conference, pages 661-678, 2024.
[47] Kun Wu, Jeongmin Brian Park, Xiaofan Zhang, Mert Hidayctoğlu, Vikram Sharma Mailthody, Sitao Huang. Steven Sam Lumetta, and Wen-mei Hwu. Tba: Faster large language model training using ssd-based activation offloading. arXiv preprint arXiv:2408. 10013, 2024.
[48] Tailing Yuan, Yuliang Liu, Xucheng Ye, Shenglong Zhang, Jianchao Tan, Bin Chen, Chengru Song, and Di Zhang. Accelerating the training of large language models using efficient activation rematerialization and optimal hybrid parallelism. In 2024 USENIX Annual Technical Conference (USENIX ATC 24), pages 545-561, 2024.
[49] Haoyang Zhang, Yirui Zhou, Yuqi Xue, Yigi Liu, and Jian Huang. G10: Enabling an efficient unitied gpu memory and storage architecture with smart tensor migrations. In Proceedings of the 56th Annual IEEE/ACM International Symposium on Microarchitecture, pages 395-410 410, 2023.
[50] Jie Zhang and Myoungsoo Jung. Zng: Architecting gpu multi-processors with new flash for scalable data analysis. In 2020 ACM/IEEE 47th Annual International Symposium on Computer Architecture (ISCA), pages 1064-1075. IEEE, 2020.
[51] Jie Zhang, Miryeong Kwon, Hyojong Kim, Hyesoon Kim, and Myoungsoo Jung. Flashgpu: Placing new flash next gpu cores. In Proceedings of the 56th Annual Design Automation Conference 2019, pages 1-6, 2019.
[52] Haoran Zhou, Wei Rang, Hongyang Chen, Xiaobo Zhou, and Dazhao Cheng. Deeptm: Efficient tensor management in heterogeneous memory for dnn training. IEEE Transactions on Parallel and Distributed Systems, 2024.
14

---

We organize our appendix as follows:
1. Appendix
provides an extended discussion of related work.
2. Appendix provides an detailed settings of the ZeRO baseline used in our evaluation.
3. Appendix provides a set of additional results.
4. Appendix D discusses the limitation of our work.
5. Appendix lists the licenses code and datasets we used in our work.
A Related Work
A.1 Expanding GPU Memory
To overcome the GPU memory wall issue, researchers have expanded GPU memory with CPU memory, persistent memory, and SSDs. Previously, when models were smaller, expanding GPU memory solely with CPU memory [77 39
43 20 reasonable. However, expanding GPU memory with host CPU memory only is expensive. As recent LLMs become increasingly large, CPU memory cannot provide enough capacity. Therefore, researchers have proposed to use persistent Alternatively, to obtain even larger capacity, SSDs 29 50 are also used expand GPÚ memory.
memory [4] to expand GPU memory 38 52
49
37
Compared to existing studies, TERAIO utilizes both SSDs and CPU memory for tensor migration. It allows users to flexibly define the maximum CPU memory used for LLM training, our lifetime-awand tensor migration algorithm helps achieve better training performance by generating optimized tensor migration plans.
A.2 Memory Efficient LLM Training
Besides tensor offloading discussed in this paper, many memory-efficient training techniques widely used in pre-training and fine-tuning, such as activation checkpointing [15
12
48
2]. Low-
Rank Adaptation (LoRA) [10 45). Activation checkpointing balances memory consumption and computation time by partially recomputing the forward propagation during backward propagation, which slows down training. In contrast, TERAIO enables LLM training without recomputation. LoRA reduces the number of trainable parameters during LLM fine-tuning. TERAIO can be combined with LoRA further reduce memory consumption and is also suitable for pre-training.
A.3 LLM Training Infranstructure
Recent studies have developed large-scale training infrastructures for LLMs. For example, Vela 27 introduces a virtualized training system that employs GPUDirect over RoCE networks to reduce communication overhead and improve scalability across multi-node clusters. IBM's Gen Al training platform (6] provides production-read stack with distributed scheduling, optimized I/O pipelines, and efficient resource management for foundation model training.
These systems focus primarily on distributed orchestration and communication efficiency. TERAIO orthogonal these works and be seamlessly integrated with such infrastructures alleviate GPU memory limitations. By enabling smart tensor migration, TERAIO enhances memory scalability without requiring changes to communication layers or cluster management frameworks.
B ZeRO Settings
We show the performance-critica parameters in the table below. With these settings, we ensure ZeRO- BeRO-Infinity achieves reasonable performance with hardware setup.
We enabled the
parameters to optimize computation and data I/O over- lap during optimizer state updates. We tuned parameters sin_memory, buffer_sii size to optimize tensor offloading throughput. For aaram_perristence. threshold and model_persistence. hewpsssttennetteehhhlleeee threshold, we use their default values.
count, and
15

---

Table 4: Detailed ZeRO settings used in evaluation.
Parameter
Value
Description
stage
3
Uses Zero-3
Overlaps read/write of next/previous tile with
aapeline_read/weeee
True
computation current tile
pin_memory
True
Uses page-locked CPU memory for faster transfers
buffer_count
4
Number of async I/O buffers for optimizer states
buffer_coune
18
Number of asyne I/O buffers for parameters
buffer_sizs
300/540M
Size of each parameter buffer
threahold
100K
Do not partition parameters smaller than this
model_ pereeppperssstnccchheeee
ays. maxsize
Upper bound of unpartitioned parameters
C Additional Evaluation Results
Menory Consunyttion
GPU
in
in
400
800
500
800
10000
20000
30000
00000
50000
G0000
0000
10000
15000
20000
25000
30000
25000
GPU Kernel Index one training iteration
GPU Kernel Index one training steration
(i) Gramite-code-buee code-bise-88
(b) 1lama3-70B
(1) Memory usage breakdown of the program during training.
Migration Bandwedth
Eviction from GPU
Fatching GPU
Eviction from GPU
Fetching to GPU
Usage (CTR0))
10
10000
20008
30000
40000
50000
G0000
10000
20000
30000
40000
50000
G0000
GPU Kernel Index in une training iteration
GPU Kernel Index one training iteration (b)
(a) Granite-coe codee--oooo-----
(2) Average migration bandwidth usage during training.
Figure 10: The memory usage breakdown and average migration bandwidth utilization when training the Granite-code-bai base-aaaeBBB model and the Llama3-70B model in TERAIO-Mixed.
Migration bandwidth utilization. Figure 10 shows that TERATO maintains high utilization of bidirectional migration bandwidth for most of the training time, thanks our 1/O-aware migration algorithm. This is particularly true when we train large models that require more memory and higher tensor migration traffic. Specifically, Granite and Llama3-70B utilize more than 16 GB/s migration bandwidth for more than 64% and 82.3% of the training time, respectively.
Host resource utilization. GDS provides significant scalability advantages by eliminating host resource contention. Without GDS, tensor offloading frameworks suffer from scalability bottlenecks. We validate this by measuring host resource usage as we scale the number of GPUs and SSDs.
CPU (GDS Disabled)
CPU Usage (GDS Enabled)
CPU Usage (GDS Disabled)
CPU Usage (GDS Boeblec)
GPU
GPUs
Bamdwitth (CRA)
OPU Usege 88)
CPU Lauge (N)
300
150
60
Nonory
300
150
100
200
100
50
100
Number of SSDs per GPU
Number of SSDs per GPU
Figure 11: Host resource utilization. The measured CPU usage relative to a single core.
As shown in Figure II host CPU usage increases linearly with the increasing number of GPUs and SSDs. With GPUs and SSDs per GPU, over 100 GB/s host memory bandwidth is consumed, and more than 3 cores are fully used. This is because host memory is used as a bounce buffer for data transfers between GPUs and SSDs. When GDS is enabled, the host CPU consumption significantly reduced. The host CPU utilization is reduced by 12.3% and host memory bandwidth
16

---

usage is decreased by 97.4% on average. As TERAIO aims to achieve low-cost LLM training, it is critical minimize the use of host CPU and memory resources.
D Limitations
Our current implementation of GDS uses NVIDIA's cuFile library which still relies the host file system to manage SSDs (c.g., filesystem metadata operations). Recent GDS studies such as GPU-initialized storage [44 34 and GoFS [2]] allow the GPU to fully bypass the host CPU by moving both the control path and data path for interacting with SSDs to the GPU. We wish to employ these studies to further reduce the overheads the host side.
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 the cost is not concern, our lifetime-aware tensor offloading approach technologies such as CXL memory. By expanding the GPU/host memory with external memory devices via CXL, new performance tradeoff's need to be considered, as the bandwidth and latency of accessing CXL memory vs. SSDs are different. We wish to extend TERAIO with new memory technologies as future work.
Although NVLink-C2C
be applied to new and emerging memory
E
Code and Datasets License
Codebase. Our model implementation is based on the Huggingface transformer repo (Apache-2.0 License) and the Torchtitan repo (BSD 3-Clause License).
Our framework implementation is based on Pytorch (BSD 3-Clause License).
The ZeRO series baselines we evaluated are from Beggatron-DeepSpee repo (Apache 2.0 License).
Datasets. We train the models using the C4 dataset (ODC-BY License).
17

---

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 an 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 (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 to use the checklist as one of the factors their evaluation. While "[Yes] " is generally preferable to "[No]" it is perfectly acceptable to answer " "[No] provided 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 used"). In general, answering "[No] or "[NA]" is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often
nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims in the abstract and introduction accurately summarize the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discussed the limitations of our work in Appendix D
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
18

---

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specing well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms
and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
No theory is included in the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Thcorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We introduced detailed models, hardware configurations, MI. frameworks, and baselines in $4.1
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
19

---

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g.. registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
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:
We submit the code in the supplemental materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7/
gwwss httpg//dddee/odeebbiiiii CCodeSoodieeeeeeeeee ssi mollicyy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
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?
20

---

Answer: [Yes]
Justification: We presented all the training hyperparameters in $4.1 The hardware configu- ration for training all baselines are listed in Table For each batch size and sequence length, we select the best hyperparameter for each baseline. We use AdamW as the optimizer and C4 as the dataset for all training tasks.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Training LLMs incurs sigmificant costs, Although we did not train for hundreds of iterations, we conducted experiments across multiple LLMs, batch sizes, sequence lengths, and memory and storage configurations. The consistent results across these settings serve as form of repeated evaluation and support the robustness of our findings.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
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?
Question:
Answer: [Yes]
Justification:
The paper provides information the compute resources in 44.l
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
21

---

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https::/neurrii httpp///bi......ccccs.eeeeeeeee.
Answer:
[Yes]
Justification:
We checked that our paper conform with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: For organizations and companies with ample GPU resources, our work enables training fine-tuning LLMs with larger batch sizes, leading to better GPU utilization. For smaller organizations and educational institutions that previously faced barriers due to limited resources, our work lowers the cost of access to LLMs.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe 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: No model dataset is released in this paper.
Guidelines:
22

---

The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification:
We included the license for the used code and datasets in Appendix E
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- ce-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
23

---

Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our L.M policy http:://eerrrppeecooooooaaerrree for what should or should be described.
24


Reviews:
[{"review_id": "wWpWMbc5zd", "reviewer": "Reviewer_JcMn", "summary": "The authors proposed TERAIO, aiming at efficient transparent tensor placement under offloaded scenario with persistent storage devices. Given the fact that offload with SSDs has long been a bottleneck for large models, the well orchestrated tensor management is critical to the feasibility of this technique. Via transparently tracking life-cycle of all tensors, TERAIO can plan ahead and make optimal transfer scheme with high generalizability. Additionally, the usage of GDS is also beneficial.", "questions": "- Please clarify the ZeRO settings (stages, DeepSpeed json parameters). They are performance critical. Some of the DeepSpeed mechanisms like persistent params are very similar to TERAIO functionalities.\n- TERAIO seems to have strong assumption on compiliation. Will it also work for non-compiled pipelines?\n- Activations in training are accumulating during training without gradient checkpointing. As they will be later fetched, they are likely be captured by TERAIO and offloaded, which will inevitably take precious downward bandwidth. Please kindly check the tradeoff between gradient checkpointing and activation offload.\n- You have mentioned the batch size boost gain (l313), but not shown in figures and tables. Please clarify.\n- For ZeRO-Infinity data point in Fig. 8, you only tested one size of CPU-side buffer. Please justify.\n- There other tensor lifecycle management works like SwapAdvisor, HOME, Checkmate, PatrickStar. Is there any comparison? If not working on the similar ground, please explain.", "limitations": "Yes", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- TERAIO is built on the basis of PyTorch, indicating it is model- and workload-agnostic.\n- GPUDirect Storage further cut down the potential latency in disk I/O.\n- Optimized tensor migration planning is promising.\n\n## Weakness\n- The paper lacks ablation studies that highlighting the effectiveness of several key designs in TERAIO. (e.g. GPUDirect Storage, tensor migration planning)\n- GPUDirect storage can significantly cut down latency in I/O, but not too beneficial in bandwith improvements. The latency issue is somehow hinderable by planning in TERAIO.\n- The generalizability of TERAIO is not well demonstrated. The LLM workload in sec 4 is too homogenous.\n- Figure 4 embedded texts are incorrectly referred. Please double check.", "was_revised": false, "final_justification": "", "rebuttal": "# 1. Ablation Studies\n\nWe conducted the ablation study as follows.\n\nWe 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.\n\n**Table: Training throughput w/ and w/o GDS (tokens/s)**\n\n||Llama3-8B (bs32,seq2k)|Granite-8B (bs32,seq4k)|Llama3-70B (bs8,seq2k)|\n|-|-|-|-|\n|TeraIO-SSD w/ GDS|1719|861|57|\n|TeraIO-SSD w/o GDS|1668|830|56|\n|TeraIO-Mixed w/ GDS|1723|1319|91|\n|TeraIO-Mixed w/o GDS|1713|1286|90|\n\nGDS 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.\n\n**Table: Host resource utilization without GDS**\n\n|GPUs|Metric|1 SSD per GPU|2 SSDs per GPU|4 SSDs per GPU|\n|-|-|-|-|-|\n|1 GPU|CPU Usage*|51%|83%|189%|\n||Memory BW (GB/s)|14|27|52|\n|2 GPUs|CPU Usage*|93%|185%|353%|\n||Memory BW (GB/s)|28|58|102|\n\n*Usage relative to a single core.\n\nWe 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. \n\n**Table: Breakdown of training iteration time (kernel stall%,overlap%,computation only%)**\n\n||Llama3-8B (bs32,seq2k)|Granite-8B (bs32,seq4k)|Llama3-70B (bs8,seq2k)|\n|-|-|-|-|\n|TeraIO-Mixed|3%,75%,22%|5%,75%,20%|45%,53%,2%|\n|ZeRO-Infinity|31%,15%,54%|30%,16%,54%|49%,23%,28%|\n\nCompared 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:\n\nOffloading 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)).\n\nMigration 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.\n# 2. Benefits of GPUDirect Storage (GDS)\n\nWe choose GDS since it provides significant scalability advantages (see Question #1).\n\nAn 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.\n# 3. LLM Workloads Evaluated\n\nWe add GPT2-20B in our evaluation and the results as follows: \n\n**Table: Training throughput (tokens/s) for GPT-20B**\n\n|Config|bs8,seq2048|bs16,seq1024|bs16,seq2048|bs32,seq2048|\n|-|-|-|-|-|\n|Ideal|546|573|614|656|\n|TeraIO-SSD|277|294|374|414|\n|TeraIO-Mixed|438|325|463|469|\n|ZeRO-Offload|270|276|309|331|\n|ZeRO-Infinity|212|214|256|283|\n\nThe results are consistent with our evaluation in Fig. 7.  \n# 4. Fig. 4 \n\nWe will fix it in the paper.\n# 5. ZeRO Settings\n\nWe show the performance-critical parameters in the table below. With these settings, we ensure ZeRO-Infinity achieves reasonable performance.\n\n**Table: Performance-critical parameters in ZeRO-Infinity**\n\n|Parameter|Value|Description|\n|-|-|-|\n|`stage`|3|Uses ZeRO-3|\n|`pipeline_read/write`|true|Overlaps read/write of next/previous tile with computation of current tile|\n|`pin_memory`|true|Uses page-locked CPU memory for faster transfers|\n|`buffer_count` (optimizer)|4|Number of async I/O buffers for optimizer states|\n|`buffer_count` (param)|18|Number of async I/O buffers for parameters|\n|`buffer_size` (param)|300/540M|Size of each parameter buffer|\n|`param_persistence_threshold`|100K|Do not partition parameters smaller than this threshold|\n|`model_persistence_threshold`|sys.maxsize|Upper bound of unpartitioned parameters|\n\nWe 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. \n\n# 6. Compilation Assumptions\n\nSimilar 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.\n\nFor 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). \n\nTeraIO 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. \n# 7. Tradeoffs Between Gradient Checkpointing and Activation Offloading\n\nWe 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. \n\nActivation 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. \n# 8. Batch Size Boost Gain\n\n**Table: Performance of training LIama3-8B using TeraIO with various batch sizes**\n|Batch Size|Micro-batch Size|Can Run w/o Offloading|Throughput (tokens/s)|Shown in Fig. 7|\n|-|-|-|-|-|\n|32|2|✓|1749|✗|\n|32|4|✗|1723|✓|\n|128|4|✗|1911|✓|\n\nThe 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.\n\nTeraIO 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.\n# 9. Clarification of the ZeRO-Infinity Data Point\n\nAs ZeRO-Infinity performs optimizer state updates on CPU, it requires CPU memory to store the optimizer states.\n \nIn 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.\n\n**Table: Throughput (tokens/s) of training Llama3-70B with different CPU memory capacities**\n\n|Configuration|512GB CPU Mem|770GB CPU Mem|1TB CPU Mem|\n|-|-|-|-|\n|2 SSDs (TeraIO)|83|/|105|\n|4 SSDs (TeraIO)|101|/|113|\n|ZeRO-Infinity|/|78|81|\n# 10. Novelty of Tensor Lifecycle Management\n\n\nTeraIO differs fundamentally from these existing works in both research goal and approach.\n\n(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.\n\n(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."}, {"review_id": "pVnMqVcHEN", "reviewer": "Reviewer_PtrA", "summary": "The paper proposed a tensor offloading framework for PCIe SSDs to enable LLM model training by analyzing the lifetime of active tensors. Moreover, based on the lifetime profiler, the proposed tensor migration planning algorithm decides the offloading and prefetching period to the GPUDirect storage and CPU memory. The evaluation shows that TERAIO outperforms the training throughput of state-of-the-art studies by 1.47x on average.", "questions": "- Why does the tensor migration plan first target SSD but not CPU memory, given that there is 1.5TB of DDR5 memory? The analysis for offloading and prefetching should also apply to main memory as well. In particular, it would be interesting to see how main memory and SSDs interact in offloading. For instance, if both are used at the same time, the available bandwidth could be the sum, or main memory could be used as a L4 cache for SSD.\n\n- In Line 32, the \"DRAM scalability challenge\" is unclear.", "limitations": "(mentioned above)", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The paper proposed an improved offload strategy by introducing a runtime tensor profiler. The migration plan is general to different I/O devices. The weakness of the work is its dependence on the GPUDirect technology, which limits its generalizability to non-NVIDIA GPUs or older generation systems. I suggest that the authors discuss why the work is limited to GPUDirect and what alternative technologies are available.", "was_revised": true, "final_justification": "My primary concern was the practical applicability of the work, particularly regarding alternatives to GDS technology and the role of CPU memory in the tensor migration strategy. The authors clarified both points sufficiently in their response.\n\nHowever, I maintain a borderline accept decision. The paper, as well as the rebuttal, contains several vague or imprecise terms, such as \"no additional cost in compilation\" and \"TeraIO does not have a minimum CPU memory requirement.\" These statements, while perhaps comprehensible in context, reflect a lack of rigor in writing. Those casual phrasing suggests the authors have not thoroughly scrutinized their responses and the precision of their claims. While the contributions have potential, the overall paper falls on the borderline of an acceptable academic article.", "rebuttal": "# 1.  Generalizability Issue of NVIDIA GDS\n\nWe 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.\n\nTeraIO 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. \n\n# 2.  Offloading Target Choice in the Migration Algorithm\n\nTeraIO 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. \n\nTeraIO 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. \n\nWhen 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).\n\n\n\n# 3.  DRAM Scalability Challenge\n\n\nThe 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. \n\n\n# 4.  Paper Formatting Concerns\n\nWe will take an editorial pass at the paper and fix the formatting issues."}, {"review_id": "wP6BcJwlB8", "reviewer": "Reviewer_UjZn", "summary": "TERAIO introduces a lifetime-aware tensor offloading mechanism for cost-efficient LLM training. It profiles tensor activity at the first few iterations and leverages GPUDirect Storage to support DMA operations. TERAIO achieves fine-grained, bandwidth-aware offloading to SSDs and minimizes stalls when migrating inactive tensors back to the GPU storage. TERAIO significantly improves training throughput with an average of 1.47x and reduces hardware cost compared to state-of-the-art baselines.", "questions": "Smart-Infinity also claims to outperform ZeRO-Infinity while scaling up the number of SSDs, so why was it not included as a baseline for comparison in the experiments? If possible, I would like to know the current performance difference between TERAIO and Smart-Infinity. This could be addressed either through a explanation or, ideally, by providing direct comparative data to demonstrate the significance.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. It requires minimal code changes and easily integrates with PyTorch, making it practical for developers to adopt.\n2. Demonstrate the strong performance over ZeRO-Offload and ZeRO-Infinity.\n3. The experimental data are clear, and the experimental design is explained in detail.\n\nWeaknesses:\n1. While the paper provides thorough experimental comparisons against ZeRO-Offload and ZeRO-Infinity, it doesn’t include Smart-Infinity, a recent storage offloading framework.", "was_revised": true, "final_justification": "Regarding the concern about whether there was a comprehensive baseline comparison, I have now received a clear explanation from the authors and am convinced that the issue has been addressed. Therefore, I am willing to increase my score to 5.", "rebuttal": "# 1. Comparison to Smart-Infinity\n\n(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.\n\n(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. \n\n(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. \nFor the above reasons, our current submission did not include the comparison with Smart-Infinity. We will add the discussion in the paper."}, {"review_id": "vduXyUxwvR", "reviewer": "Reviewer_9SW7", "summary": "This paper proposes an active memory profiler and memory offloading algorithm for reducing the peak memory usage of LLM training workloads. The key insight is that the active memory required for LLM training is a small fraction of the total memory allocated, and therefore carefully scheduling memory offloads and corresponding prefetches from SSD memory can enable scaling beyond the fundamental device DRAM memory limit. The evaluation shows that the proposed offloading framework TERAIO achieves an average speedup of 1.47x compared to the ZeRO-Offload and ZeRO-Infinity techniques.", "questions": "- For the characterization study in Section 2, which sequence length / batch sizes were used for each model? Do the conclusions change with different parameters?\n-  Why is there still a ~20% gap between TERAIO and the ideal performance?\n- How often do kernel stalls happen in practice and what is the resulting performance impact? Can these kernel stalls be predicted in advance using theoretical analysis?\n- How does this technique change with emerging memory technologies like the high-bandwidth GPU-CPU interconnects available with NVIDIA Grace systems?", "limitations": "The authors have described the limitations of their work and there are no ethical concerns.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- Evaluation demonstrates significant performance improvement over baseline approaches (ZeRO-Offload / ZeRO-Infinity).\n- Implementation of tensor memory profiler instruments PyTorch to minimize the required developer effort.\n- Paper is well written and contributions are clear.\n\nWeaknesses:\n- Limited discussion of related work / analysis of baselines. The main comparison is to ZeRO-Offload and ZeRO-Infinity, but it’s unclear A) whether these are state-of-the-art approaches to memory offloading, and B) why TERAIO is able to outperform ZeRO-Infinity given that both ZeRO-Infinity and TERAIO are able to offload to both host memory as well as SSDs. In Section 4.2 the authors write:\n\t“Since SSDs offer larger capacity, ZeRO-Infinity can still train Llama-70B. However, it only achieves 43.0% of the ideal performance, because its coarse-grained offloading scheme cannot efficiently utilize the limited SSD bandwidth to migrate larger tensors.”\nWhat makes the ZeRO-Infinity scheme more coarse grained than TERAIO? Later in Section 4.4 the authors write:\n“ZeRO-Infinity requires 170GB and 770GB of CPU memory to train Llama-8B and Llama-70B, respectively, as it has to offload gradients and optimizers into CPU.”\nIs this fundamental to ZeRO-Infinity’s offloading design or just an implementation detail?\nIn general, is the speedup mostly attributed to a more intelligent offload / prefetch schedule, a more robust implementation (e.g. more effectively using GPUDirect), or something else?\n- Limited novelty of tensor profiler. Many papers have proposed granular profiling for DNNs, but TERAIO includes this profiler as a key contribution. Is there anything particularly different about the TERAIO profiler?\n- Results are mostly empirical, but it would be nice to see more theoretical analysis of how to compute the expected performance of this approach given key parameters such as the activation sizes, hardware bandwidth, and available memory capacity. This would help determine whether completely hiding the transfer latency is possible without having to actually run the system in practice.\n- No comparison to activation recomputation. The authors suggest that this will slow down training which is true, but how does this slowdown compare to potential kernel stalls caused by delays in the memory prefetching?", "was_revised": true, "final_justification": "Overall, the additional experiments provided by the authors help clarify the performance benefits of TeraIO. Including these experiments and expanding the results to include more in-depth study of the ZeRO-Infinity baseline would strengthen the paper. As a result, I increased my score.", "rebuttal": "# 1. Is ZeRO-Infinity the SOTA Approach\n\nZeRO-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.\n\nTo 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.\n\n# 2. Why TeraIO Outperforms ZeRO-Infinity\n\nTeraIO outperforms ZeRO-Infinity due to its design differences in offloading granularity and migration strategy.\n\n**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)).\n\n**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.\n\n# 3. Clarification of CPU Memory Requirements of ZeRO-Infinity\n\nSince 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. \n\n\n# 4. Clarification of Speedup Breakdown\n\nThe 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). \n\n**Table: Training throughput comparison with and without GDS (tokens/s)**\n\n||Llama3-8B (bs32, seq2k)|Granite-8B (bs32, seq4k)|Llama3-70B (bs8, seq2k)|\n|-|-|-|-|\n|TeraIO-SSD w/ GDS|1719|861|57|\n|TeraIO-SSD w/o GDS|1668|830|56|\n|TeraIO-Mixed w/ GDS|1723|1319|91|\n|TeraIO-Mixed w/o GDS|1713|1286|90|\n\nAlthough 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**.\n\nWithout 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.\n\n**Table: Host resource utilization without GDS**\n\n|GPUs|Metric|1 SSD per GPU|2 SSDs per GPU|4 SSDs per GPU|\n|-|-|-|-|-|\n|1 GPU|CPU Usage*|51%|83%|189%|\n||Memory BW|13.9 GB/s|27.3 GB/s|51.7 GB/s|\n|2 GPUs|CPU Usage*|93%|185%|353%|\n||Memory BW|27.7 GB/s|57.5 GB/s|102.2 GB/s|\n\n*CPU Usage: Usage relative to a single CPU core.\n\n**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.\n\n# 5. Novelty of Tensor Profiler\n\nCompared to existing DNN profilers, our tensor lifetime profiler has two unique features:\n\n\n(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.\n\n(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.\n\n# 6. Theoretical Analysis of Expected Performance\n\nTeraIO 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). \n\n# 7.  Activation Recomputation\n\nWe 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.\n\nMoreover, 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. \n\n# 8. Clarification of Settings in the Characterization Study\n\nWe 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.\n\n||Llama3-8B|T5-11B|GPT2-40B|Llama3-70B|\n|-|-|-|-|-|\n|Batch size|128|256|16|64|\n|Sequence length|4,096|512|1,024|2,048|\n\nAs 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.\n\n# 9. Clarification of Performance Gap\n\nThe 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).\n\n# 10. Kernel Stalls\n\nTo 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.\n\nThese kernel stalls can be predicted and minimized using our performance model as discussed in Section 2.2 and [Question 6](#6-theoretical-analysis-of-expected-performance)).\n\n**Table: Breakdown of training iteration time (kernel stall%, overlap%, computation only%)**\n\n||Llama3-8B (bs32,seq2k)|Granite-8B (bs32,seq4k)|Llama3-70B (bs8,seq2k)|\n|-|-|-|-|\n|TeraIO-SSD|4%,96%,~0%|38%,59%,3%| 66%,34%,~0%|\n|TeraIO-Mixed|3%,75%,22%|5%,75%,20%|45%,53%,2%|\n\n# 11. Emerging Memory Technologies\n\nAlthough 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. \n\nEven 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."}]

Author rebuttal:

Meta-review: This paper presents a new tensor offloading system, TERAIO,  to enhance the efficiency of LLM training over multiple GPUs and SSDs. The design of this system comes from the observation that only a small portion of tensors are active during LLM training. Based on this observation, it estimates the lifetime of each tensor by only profiling a small number of training iterations, and then generates the optimized tensor offloading strategy. TERAIO can be directly integrated with PyTorch to support compiled LLM training workloads, and experiments also demonstrate its superiority over baselines. 

The reviewers agree that this is a solid workload, well written, and the solution is interesting and promising. The authors’ responses with new experiments and explanations address the majority of reviewers’ comments. Given that each reviewer is positive about this paper, it is recommended for acceptance. The authors are still suggested to address the remaining points in the revision, including more clear explanation about the compilation cost, and performance benefits.
Decision: Accept (poster)