Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage

A1 主要贡献

本文介绍了一个名为 Teraio 的新型张量卸载框架，旨在通过使用低成本的基于 PCIe 的固态硬盘（SSD）来扩展 GPU 内存，从而实现大语言模型（LLM）的高效训练。

核心问题与研究目标：
LLM 的训练面临着巨大的内存需求挑战，其增长速度远超 GPU 内存的扩展速度。传统的解决方案，如使用主机内存或现有的 SSD 卸载方案（如 ZeRO-Infinity），存在效率不高、无法充分利用 SSD 带宽和有效隐藏访问延迟等问题，导致性能不佳。理想的目标是在使用低成本 SSD 扩展 GPU 内存的同时，实现接近于拥有无限板载内存的 GPU 的训练性能。

核心观察与创新点：
研究发现，在 LLM 训练的每次迭代中：
1. 活动张量内存占用小：当前正在被 GPU 内核使用的“活动张量”仅占已分配 GPU 内存的很小一部分（平均为 1.7%）。
2. 非活动张量大且空闲期长：大多数非活动张量体积庞大，且在很长一段时间内不会被使用。
这些观察为在不拖慢 GPU 训练进程的情况下，向/从慢速 SSD 卸载/预取张量创造了充足的机会。

Teraio 框架的设计与贡献：
基于以上观察，Teraio 框架遵循三个设计原则：(1) 卸载大的非活动张量以节省 GPU 内存并最大化 SSD 带宽利用率；(2) 根据非活动周期分布来精确决定卸载和预取的时机；(3) 精确调度张量移动以有效重叠 I/O 与 GPU 计算。Teraio 主要由三个组件构成：

  1. 张量生命周期分析器 (Tensor Lifetime Profiler)：

    • 通过在训练开始的几次迭代中进行性能分析，精确提取张量的活动模式（如大小和生命周期）。
    • 该分析器通过对 PyTorch 中的自动算子生成器进行插桩来实现，而非侵入式地修改每个算子，从而实现了对 PyTorch 的最小化代码修改。

  2. 生命周期感知的张量迁移算法 (Lifetime-aware Tensor Migration Algorithm)：

    • 根据分析器获得的张量活动模式，生成最优的张量卸载/预取计划。
    • 该算法优先卸载具有长非活动周期的大张量，以充分利用可用的存储 I/O 带宽。对于非活动周期短的张量，则尽量保留在 GPU 内存中。
    • 当 SSD 带宽饱和时，算法会将主机内存作为备用卸载目标。

  3. 使用 GPUDirect Storage 的张量迁移引擎 (Tensor Migration Engine)：

    • 在运行时执行生成的迁移计划。
    • 利用 GPUDirect Storage 技术实现 GPU 和 SSD 之间的直接数据传输，绕过主机 CPU，从而缓解可扩展性瓶颈并最大化 SSD 带宽利用率。

主要贡献总结：

A3 关键观察与设计原则

本节介绍了在 LLM 训练中对张量活动模式的表征研究。该研究利用了将在 §3.1 中讨论的张量生命周期分析器，分析了 LLM 训练期间张量大小和生命周期的分布。实验使用了两块 NVIDIA H100 GPU 和两阶段 1f1b 流水线并行【索引21，PipeDream: Generalized pipeline parallelism for dnn training，2019，SOSP】。研究涵盖了多种 LLM 架构，包括仅解码器模型（如 Llama3-8B、Llama3-70B【索引4，The llama 3 herd of models，2024，arXiv】和 GPT2-40B【索引27，Language models are unsupervised multitask learners，2019，OpenAI blog】）以及编码器-解码器模型（如 T5-11B【索引6，T5 11b，2025，Hugging Face】）。对于内存需求超过 GPU 内存容量的模型，会将当前内核不需要的张量卸载到 SSD。


(b) 流水线并行中不同阶段的内存消耗。
图 1：并行训练程序一次迭代中所有张量和活动张量的内存消耗（相对于 GPU 内存容量）。图中使用了对数刻度。


图 2：张量的非活动周期分布。

2.1 丰富的张量迁移机会

活动张量的内存需求小。我们首先研究了每次训练迭代中张量的内存需求和使用情况。我们将当前正在被运行的 GPU 内核使用的张量定义为活动张量。图 1 展示了每个 GPU 内核使用的活动张量的内存消耗。图 1 (a) 显示了不同模型中张量的内存消耗，图 1 (b) 显示了不同流水线阶段张量的内存消耗。在我们研究的所有 LLM 中，尽管它们的总内存使用量远超 GPU 内存容量，但活动张量仅占总 GPU 内存容量的不到 14%（平均为 1.7%）。GPU 内存中的大多数张量是非活动的，可以被卸载到低成本的 SSD，因此，我们可以最好地利用 GPU 内存来存放那些即将被内核使用的张量。

非活动张量的非活动周期长。为了解非活动张量保持非活动状态的时长以及它们消耗多少 GPU 内存，我们研究了其非活动周期的分布，如图 2 所示。对于我们研究的所有模型，大多数张量的大小范围在 10MB 到 1GB 之间。我们观察到，超过 40% 的这些张量保持非活动状态超过 104 微秒。这些非活动周期比以 6.5 GB/s 的带宽将这些张量迁移到 SSD 所需的时间要长。基于这一洞察，我们可以确保这些张量能够被高效迁移，而不会对训练性能产生负面影响。

长非活动周期的成因。长非活动周期是由于 LLM 的稀疏张量访问模式和高计算强度造成的。从空间角度看，尽管 LLM 有数十或数百个层，但许多张量仅在单个层内使用。从时间角度看，每层中的计算密集型内核（如注意力机制）需要相当长的时间，这为 Teraio 提供了丰富的机会来将计算与非活动张量的迁移重叠起来。

与传统 DNN 模型的对比。我们还观察到，与传统 DNN 模型相比，LLM 更高的计算强度和更大的模型尺寸导致了显著更长的非活动周期。例如，在 BERT-Large【索引3，Bert: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv】中，48% 的张量大于 100MB，而这些大张量中超过 60% 的非活动周期比 LLM 中的短两个数量级。

2.2 张量迁移的带宽需求

不同迁移带宽下的性能分析。我们现在研究需要多少迁移带宽才能使卸载达到接近理想的训练性能。我们量化了不同 LLM 模型在每个 GPU 可用的不同迁移带宽下的屋顶线性能。为了方便这项研究，我们建立了一个性能模型来估计训练时间。在该模型中，我们假设每个内核的执行时间与我们在表征研究中收集到的值相同。我们模拟在指定带宽下的张量迁移，并检查内核所需的张量是否已在 GPU 内存中。如果由于 SSD 带宽有限，它们仍在迁移中，则迁移的等待时间将加到总训练时间中。


图 3：不同迁移带宽下的屋顶线分析。训练性能相对于假设 GPU 内存无限的理想情况进行了归一化。

带宽需求的结论。图 3 显示了 LLM 在不同迁移带宽下的归一化屋顶线训练吞吐量。我们观察到，32 到 48 GB/s 的带宽足以让 LLM 实现接近理想的性能。这样的带宽需求可以通过聚合多个商用 SSD（例如，一个 SSD 阵列）轻松实现，证明了 Teraio 的可行性。

A2 方法细节


图 4：Teraio 的系统概览。

Teraio 系统概览。我们在图 4 中展示了 Teraio 的系统概览。给定一个 LLM，Teraio 的张量生命周期分析器与 PyTorch 协同工作，以跟踪张量的大小和生命周期（§3.1）。在最初几次训练迭代中，Teraio 追踪执行图并收集每个内核的执行时间。由于后续迭代遵循相同的执行图，张量的活动模式保持不变。张量迁移算法（§3.2）创建一个张量迁移计划，该计划（1）最大化计算与迁移的重叠，以及（2）最小化迁移流量。该算法迭代地选择最佳的卸载候选者，直到所需的 GPU 内存适合实际的 GPU 内存容量。对于迁移目的地，它倾向于将张量迁移到 SSD。一旦 SSD 带宽饱和，它也会使用可用的 CPU 内存。在 LLM 训练期间，Teraio 的张量迁移引擎透明地执行迁移计划（§3.3）。

3.1 张量生命周期分析器

张量追踪。Teraio 对 PyTorch 框架进行插桩，以在运行时跟踪张量并测量内核执行时间。一个张量在以下三种情况之一被认为是活动的。首先，当一个张量是 PyTorch CUDA 算子的输入或输出时，它是活动的。然而，对 PyTorch 进行插桩以跟踪每个算子是具有挑战性的，因为有数千个算子。相反，我们利用 PyTorch 的自动算子生成器，它为每个算子生成源代码，来插入分析代码，该代码将在算子运行时将所有输入和输出张量标记为活动。其次，对于涉及 GPU 间通信的张量，它们应该在 GPU 内存中是活动的。第三，当 PyTorch 明确检查一个张量是否驻留在 GPU 内存中时，该张量被认为是活动的。这在更新优化器状态时发生。对于第二和第三种情况，由于总共只有少数几个通信算子和 PyTorch 检查，我们直接在源代码中将相应的张量设置为活动。


图 5：Teraio 跟踪并分析张量活动模式，以供张量迁移算法使用。

张量活动模式分析。图 5 展示了张量信息在运行时是如何被收集的。为了解一个张量何时消耗 GPU 内存以及何时必须驻留在 GPU 内存中，我们需要收集其张量大小和活动时间。当一个算子被执行时，插桩代码会记录相应张量的张量大小和活动时间。分析器将根据其活动状态之间的持续时间计算非活动时间段。具体来说，对于像梯度和激活（图 5 中的张量 0 和张量 1）这样的中间张量，它们在计算完成后会立即被释放，我们将其非活动时间段量化为两个活动周期之间的时间间隔。对于像模型权重和优化器状态（图 5 中的张量 2）这样的全局张量，它们在多个训练迭代中被使用，在训练开始前被分配，并且在训练期间从不被释放。因此，在某些情况下，分析器可能需要根据跨越两次迭代的活动状态来计算其非活动周期。

3.2 生命周期感知的张量迁移算法

算法核心思想。生命周期感知的张量迁移算法在每次 LLM 训练迭代中迭代地寻找最佳的卸载候选者，直到所需的 GPU 内存低于容量限制。通过跟踪所需的 GPU 内存量和存储 I/O 带宽利用率，该算法能够评估张量卸载的潜在好处。我们讨论我们的关键思想如下。


(a). I/O 带宽使用情况影响实际迁移时间


(b). 卸载一个非活动张量的好处

存储 I/O 带宽感知的规划。为了卸载一个非活动张量，我们希望它尽可能长时间地不占用 GPU 内存。因此，理想情况下，我们会在它一变为非活动状态时就将其卸载，并在后续内核需要它之前立即将其预取回来。然而，在现实中，我们能够卸载和预取非活动张量的时间受到存储带宽使用的极大影响。例如，在图 6(a) 中，张量的非活动周期 I 是从 $t_{start}$ 到 $t_{end}$。然而，由于 I/O 带宽被先前计划的迁移占用，我们必须将实际卸载推迟到 $t_{offload}$，并在 $t_{prefetch}$ 之前开始实际预取。这意味着我们只能在 $t_{offloaded}$ 到 $t_{prefetch}$ 的时间段内减少内存消耗。在规划张量迁移时，我们的算法会跟踪估计的存储 I/O 带宽使用情况，并以一种 I/O 感知的方式计算内存消耗的减少量（算法 1 中的第 4-5 行和第 16-17 行）。

算法 1 生命周期感知的张量迁移规划

Require: 张量非活动周期集合 I = {(ti, si, starti, endi)}，其中 ti 是张量 ID，si 是大小，[starti, endi] 定义了非活动的内核范围；GPU 内存容量 M_GPU；N 个内核所需的 GPU 总内存量估计 M = [m0, m1, ..., mN−1]；内核执行时间 T = [τ0, τ1, ..., τN−1]；I/O 带宽使用状态
Ensure: 迁移计划列表 P = {(ti, trigger_time, deadline, target)}
1: 初始化迁移计划列表 P = {} 和所需 GPU 内存估计 M' = M
2: while 峰值内存消耗 max(M) > M_GPU do
3:   for each inactive period (t, s, start, end) ∈ I do
4:     通过分析可用的卸载带宽，确定最早可行的卸载完成时间 t_offloaded
5:     通过分析可用的预取带宽，确定最晚需要的预取开始时间 t_prefetch
6:     if t_offloaded < t_prefetch then
7:       识别关键内存压力区域 C = {k|M'[k] > M_GPU, k ∈ [t_offloaded, t_prefetch]}
8:       计算卸载效益 B = s × Σ_{k∈C} τk
9:     end if
10:  end for
11:  选择具有最佳卸载成本效益 (B/s)_max 的非活动周期
12:  if 不存在可行的卸载候选者 then
13:    break
14:  end if
15:  将迁移计划 (t, start, t_offloaded, SSD/CPU) 和 (t, t_prefetch, t_prefetched, GPU) 添加到 P
16:  从受影响的内核中减去张量大小 s 来更新 M'
17:  更新 I/O 带宽使用状态
18:  从考虑范围中移除选定的非活动周期
19: end while
20: return P

量化张量卸载的效益与成本。为了充分利用可用的 I/O 带宽，我们希望优先卸载具有长非活动周期的大张量。遵循这一原则，我们的算法通过估算其效益和成本来寻找最佳的卸载候选者。为了量化效益，在给定时间 T，我们将关键内存压力定义为 GPU 内存消耗超过容量的部分。张量迁移的效益被定义为关键内存压力减少量随时间的积分，如图 6(b) 中的阴影区域所示。我们将成本量化为张量卸载和预取时间的总和。Teraio 的迁移算法按所有候选者的效益成本比进行排序。然后，它选择比率最高的张量在本轮迭代中进行迁移。此过程显示在算法 1 的第 6-11 行。

决定卸载目的地。Teraio 优先选择 SSD，因为它们容量大、成本低。当估计的 SSD 带宽饱和时，Teraio 会尽最大努力将张量迁移到主机内存。Teraio 允许用户定义可用于张量迁移的主机内存的最大量。在内部，迁移算法会跟踪估计的主机内存消耗。如果已达到用户定义的限制，即使 SSD 带宽饱和，Teraio 也不会将张量卸载到主机内存。

最小化内核停顿。如果所需的 GPU 内存超过了可用的 GPU 内存容量，并且 Teraio 无法将张量卸载到 SSD 或主机内存，Teraio 将会停顿内核执行，以等待更多的非活动张量被卸载。它还会等待内核所需的张量被迁回 GPU 内存。这会导致存储 I/O 带宽竞争。由于下一个内核会一直停顿，直到所需的张量迁移到 GPU 内存，因此这些迁移应尽早完成，以避免拖慢 GPU 训练过程。因此，Teraio 将这些必须在下一个内核开始前完成的迁移标记为“紧急”。在运行时，张量迁移引擎（§3.3）总是优先处理这些紧急迁移，而不是其他待处理的迁移。

3.3 使用 GPUDirect Storage 的张量迁移引擎

张量定位。为了执行迁移计划，我们需要根据张量标识符定位待迁移张量的地址。Teraio 在 PyTorch 中维护一个基于哈希表的张量位置表，用于将张量标识符映射到它们当前的设备和地址。

迁移机制。迁移引擎在 GPU 内存和外部内存之间迁移张量。对于 GPU 内存和 SSD 之间的张量迁移，我们使用 GPUDirect Storage 来实现它们之间的直接数据传输。我们选择 GPUDirect Storage 的动机是主机 CPU 在多 GPU 系统中可能存在的可扩展性限制。例如，在一个 8-GPU 系统上，为了实现接近理想的性能，每个 GPU 需要 32 到 48 GB/s 的双向带宽（见 §2）。通过 GPUDirect Storage，我们可以通过 PCIe Gen5 交换机将 8 个 SSD 直接连接到每个 GPU，以满足带宽需求。对于传统方法，即使用主机 CPU 先将数据从 SSD 读取到主机内存，然后再使用 cudaMemcpy 将数据移动到 GPU，这种冗余的数据复制不仅会造成性能开销，还会浪费宝贵的 CPU 周期【索引7，GPUDirect Storage: A Direct Path Between Storage and GPU Memory，NVIDIA Developer Blog】、【索引26，GPU-initiated on-demand high-throughput storage access in the bam system architecture，2023，ASPLOS】。但是，如前所述，Teraio 仍然支持 GPU 和 CPU 内存之间的迁移。

A4 实验环境

实验目标：
我们展示了 (1) 在训练远超 GPU 内存容量的 LLM 时，Teraio 的性能平均比最先进的卸载框架高 1.47 倍（§4.2）；(2) 与仅使用 GPU 内存训练 LLM 的情况相比，Teraio 将成本降低了高达 5.41 倍（§4.3）；(3) 与现有的卸载框架相比，Teraio 的成本效益平均提高了 1.45 倍（§4.3）；(4) 即使使用更少的 CPU 内存和 SSD，Teraio 也能实现比 ZeRO-Infinity 更高的吞吐量（§4.4）。

4.1 实验设置

A4 实验结果

4.2 端到端性能


(1). 不同批次大小下的平均训练吞吐量。 (2). 不同序列长度下的平均训练吞t量。
图 7：Llama3-8B、Granite-code-base-8B 和 Llama3-70B 在不同批次大小和序列长度下的平均训练吞吐量。bs 是批次大小。seq 是序列长度。M 是 LLM 训练在一个 GPU 上的总体峰值内存消耗与 GPU 内存容量的比值。“×” 表示该框架因内存不足错误而无法训练此模型。

训练吞吐量。如图 7 所示，在训练 Llama3-8B 和 Granite-8B 时，ZeRO-Offload 分别达到了理想性能的 65.9% 和 65.5%。尽管 ZeRO-Infinity 需要额外的时间将张量从 CPU 内存进一步迁移到 SSD，但其训练吞吐量略低于 ZeRO-Offload。这种微小的性能差异源于 8 个 SSD 的高聚合带宽以及卸载到 SSD 的优化器和参数尺寸相对较小。对于 Llama-70B，ZeRO-Offload 失败，因为主机内存容量太小，无法存储如此大模型的卸载张量。由于 SSD 提供了更大的容量，ZeRO-Infinity 仍然可以训练 Llama3-70B。然而，它仅达到了理想性能的 43.0%，因为其粗粒度的卸载方案无法有效利用有限的 SSD 带宽来迁移更大的张量。

Teraio 性能。Teraio 的性能比 ZeRO 系列高出多达 1.59 倍。对于 Llama3-8B，Teraio-Mixed 和 Teraio-SSD 都能达到接近理想的性能，证明了我们的生命周期感知张量迁移算法在选择最有益的张量进行迁移方面的有效性。此外，我们发现即使 2 个 H100 GPU 能为训练 Llama3-8B 提供足够的 GPU 内存而无需内存扩展，Teraio 仍允许我们增加（微）批次大小，从而将吞吐量提高了 9%。对于 Granite，Teraio-SSD 的性能与 ZeRO 系列相似，而 Teraio-Mixed 通过利用 CPU 内存仍能提供接近理想的性能。结果表明，随着模型尺寸的增加，实现接近理想的性能不仅需要我们的生命周期感知张量迁移算法，还需要比 4 个 SSD 所能提供的更高的迁移带宽。对于 Llama-70B，由于其内存需求显著超过 GPU 内存容量，需要更高的迁移带宽。即使是 Teraio-Mixed 也只能达到理想性能的 59.7%。尽管如此，它仍然比 ZeRO-Infinity 快 1.33 倍。

不同批次大小和序列长度的影响。随着批次大小和序列长度的变化，Teraio-Mixed 的训练吞吐量在所有卸载框架中始终最接近理想吞吐量，而 Teraio-SSD 的性能与 ZeRO 系列相当或更好。对于同一模型，增加批次大小不会增加内存需求，因为每个批次被分割成包含相同数量训练样本的微批次。然而，当序列长度增加时，内存需求也会增加。因此，当序列长度超过 3,072 时，ZeRO-Infinity 无法训练该模型，因为 GPU 内存容量小于模型的激活张量。相比之下，Teraio 可以训练所有模型配置，因为它能够卸载任何非活动张量以节省内存。


图 8：在 Teraio-Mixed 中训练 Granite-code-base-8B 模型和 Llama3-70B 模型时的平均迁移带宽利用率。

迁移带宽利用率。为了进一步理解 Teraio 为何能取得良好性能，我们研究了其迁移带宽利用率。图 8 显示，得益于我们的 I/O 感知迁移算法，Teraio 保持了双向迁移带宽的高利用率。

4.3 训练成本

成本对比。为了评估 Teraio 的成本，我们在表 2 中总结了每个基准设置中使用的不同设备的价格。Teraio-SSD 只需要一台配备 128GB CPU 内存和 8 个 SSD 的服务器，因为它只将张量迁移到 SSD，而 Teraio-Mixed 则同时使用 SSD 和 1TB CPU 内存以实现更高效的张量迁移。由于 ZeRO-Offload 和 ZeRO-Infinity 在训练 LLM 时消耗大量 CPU 内存，它们都需要一台配备 1TB 内存的服务器。在我们的评估中，ZeRO-Infinity 还额外使用了 8 个 SSD。

表 2：Teraio 及其他基准在 Llama3-70B 训练中使用的各设备成本。价格引自 Exxact [5]。PureGPUs 设置包含能够训练 Llama3-70B 而无需将张量卸载到主机内存或 SSD 的最少数量的 8-GPU H100 服务器。

成本效益分析。我们还与 PureGPU 设置进行了比较，在该设置中，所有张量都保存在 GPU 内存中。为了提供足够的 GPU 内存来训练 Llama3-70B，我们需要为两台各配备 8 个 H100 GPU 的服务器支付 499,591.40 美元。相比之下，Teraio-SSD 和 Teraio-Mixed 分别节省了 5.88 倍和 5.41 倍的成本。与 ZeRO-Offload 和 ZeRO-Infinity 相比，由于我们的机器设置相似，Teraio 1.47 倍的训练性能提升分别转化为 1.45 倍和 1.47 倍的成本效益提升。

4.4 不同数量 SSD 和 CPU 内存容量的影响


图 9：随着每个 GPU 的 CPU 内存容量和 SSD 数量变化时的训练吞吐量。

资源扩展性。图 9 显示了当我们改变每个 GPU 可用的 CPU 内存容量和使用的 SSD 数量时，Teraio 的训练吞吐量。随着我们增加 SSD 数量或 CPU 内存容量，Teraio 的性能呈有利的扩展趋势。对于 Llama-8B，Teraio 仅用 2 个 SSD 和 64 GB 的 CPU 内存即可达到接近理想的性能。即使对于 Llama-70B，我们观察到仅用 2 个 SSD 和 512GB 的 CPU 内存仍可达到使用 4 个 SSD 和 1,024GB CPU 内存所获训练吞吐量的 73.1%。这些结果验证了我们 LLM 表征研究的关键观察：通过生命周期感知的张量卸载，即使硬件资源有限，我们也能实现良好性能。

与 ZeRO-Infinity 的资源对比。与 ZeRO-Infinity 相比，Teraio 在实现更好性能的同时，显著降低了对 CPU 内存容量的需求。ZeRO-Infinity 分别需要 170GB 和 770GB 的 CPU 内存来训练 Llama-8B 和 Llama-70B，因为它必须将梯度和优化器卸载到 CPU。相比之下，Teraio 即使在更有限的硬件资源下也能实现更好的性能，因为它高效地利用了 SSD 的大容量和 CPU 内存的额外带宽。

A5 结论

我们提出了 Teraio，一个生命周期感知的张量卸载框架，它能够为 LLM 训练精确地规划和执行细粒度的张量卸载和预取指令。借助可预测的张量活动模式，Teraio 能够最好地利用宝贵的 GPU 内存来加速 GPU 训练过程，同时利用大容量 SSD 来降低训练成本。与现有的张量卸载工作相比，Teraio 为 LLM 训练提供了一个更实用、更具成本效益的解决方案。