LazyTrain:大语言模型训练中面向零浪费产出优化的有限资源分配
摘要
LazyTrain 是一个研究系统,通过将检查点选择、激活放置、重计算以及 CPU-GPU-NVMe 通信重叠建模为混合整数调度问题,优化有限硬件上的大语言模型训练,相比匹配基线实现了约 1.24 倍的 TFLOPS 提升。
arXiv:2608.11919v1 公告类型:新
摘要:在有限硬件上训练大语言模型日益成为一个涉及 GPU 计算、主机内存、PCIe 传输和存储带宽的调度问题。现有卸载系统减少了 GPU 驻留,MegaTrain 表明 CPU 主控的逐层流式执行器可以在单块 GPU 上训练大模型,但固定的检查点和放置启发式仍使通信暴露在关键路径上。我们提出 LazyTrain,一个构建在逐层流式执行器之上的优化层。LazyTrain 将检查点选择、激活放置、重计算以及 CPU-GPU-NVMe 通信重叠建模为混合整数调度问题,然后在训练过程中执行求解得到的策略。它进一步将 8 位优化器状态与快速梯度裁剪耦合为单个混合 8 位算子:状态压缩减少了优化器状态内存,而快速裁剪抵消了额外的 CPU 侧更新开销。在从 Qwen2.5-3B 到 Qwen3.6-27B 的 H800 实验中,相比匹配基线运行,LazyTrain 将持续 TFLOPS 提升了约 1.24$\times$;RTX 3090 实验同样在每个模型规模下将最大可行批次大小增加 1。在主要的 Qwen3.6-27B H800 MetaMathQA 运行中,LazyTrain 在批次大小 72 时达到 219.95 TFLOPS 和 1361 tokens/s,GPU 内存峰值 68.84\,GB,并在完整评估分割上获得 95.42\% 的精确匹配准确率。源代码可在 https://github.com/DataArcTech/LazyTrain 获取。
查看缓存全文
缓存时间: 2026/08/13 15:29
# LazyTrain:面向大语言模型训练零浪费产出优化的有限资源分配
来源:https://arxiv.org/html/2608.11919
Cehao Yang 隶属:香港科技大学(广州)
Honghao Liu 隶属:香港科技大学(广州)
Xueyuan Lin
Xuhui Jiang 隶属:IDEA Research;DataArcTech Ltd.
Chengjin Xu 隶属:IDEA Research;DataArcTech Ltd.
Jia Li
Jian Guo
###### 摘要
在有限硬件上训练大语言模型越来越成为一个涉及 GPU 算力、主机内存、PCIe 传输和存储带宽的调度问题。现有的卸载系统降低了 GPU 驻留内存,MegaTrain 表明 CPU 主控的逐层流式执行器可以在单张 GPU 上训练大模型,但固定的检查点策略和放置启发式方法仍会让通信暴露在关键路径上。我们提出 LazyTrain,一个构建在逐层流式执行器之上的优化层。LazyTrain 将检查点选择、激活放置、重计算和 CPU-GPU-NVMe 通信重叠建模为一个混合整数调度问题,然后在训练过程中执行求解得到的策略。它进一步将 8 位优化器状态与快速梯度裁剪耦合为单一 Hybrid 8-bit 算子:状态压缩减少了优化器状态内存,而快速裁剪抵消了额外的 CPU 侧更新开销。在从 Qwen2.5-3B 到 Qwen3.6-27B 的 H800 实验中,LazyTrain 相较匹配的基线运行将持续 TFLOPS 提升了约 1.24 倍;RTX 3090 实验同样在每个模型规模下将最大可行批大小增加了一。在主要的 Qwen3.6-27B H800 MetaMathQA 运行中,LazyTrain 在批大小为 72 时达到 219.95 TFLOPS 和 1361 tokens/s,GPU 内存峰值 68.84 GB,并在完整评估拆分上获得 95.42% 的精确匹配准确率。源代码可在 https://github.com/DataArcTech/LazyTrain 获取。
11 同等贡献。22 通讯作者。
## 1 引言
大语言模型训练越来越依赖于系统如何充分利用完整的内存层级,而不仅仅是 GPU 内存。即使对于监督微调,模型参数、梯度、优化器状态和激活也可能超过单个加速器的内存,而 CPU DRAM 和本地存储提供了大得多的容量,但带宽明显更低。因此,高效训练需要一个调度,不仅要决定卸载什么,还要决定每次传输何时可以隐藏在 GPU 计算之下(5 (https://arxiv.org/html/2608.11919#bib.bib2))。
现有的内存高效训练系统解决了这个问题的关键部分。ZeRO 风格的分片和卸载降低了模型状态的设备驻留(15 (https://arxiv.org/html/2608.11919#bib.bib7);16 (https://arxiv.org/html/2608.11919#bib.bib8)),FSDP 和 Gemini 提供了广泛使用的异构内存运行时策略(23 (https://arxiv.org/html/2608.11919#bib.bib13);6 (https://arxiv.org/html/2608.11919#bib.bib3)),而存储感知系统表明 CPU 和 NVMe 可以充当有用的容量层(10 (https://arxiv.org/html/2608.11919#bib.bib5);21 (https://arxiv.org/html/2608.11919#bib.bib11))。MegaTrain 将 CPU 内存视为参数和优化器状态的权威存储,将层流式传输通过 GPU 内存,并使用块式重计算来限制激活驻留(20 (https://arxiv.org/html/2608.11919#bib.bib12))。然而,其激活调度仍然是一个固定的启发式方法。固定的检查点间隔可能是可行的,但它不会联合搜索检查点边界、GPU/CPU/NVMe 归属、重计算块和传输窗口。
图 1:优化引导训练调度的动机。固定的启发式方法可以满足内存上限,同时将 CPU-GPU 或 NVMe 传输暴露在关键路径上。LazyTrain 则搜索检查点放置和传输分配,尽可能在计算窗口内隐藏激活流量。
本文的核心观察是,有限资源训练应被优化为一个受约束的调度问题。如图 1 (https://arxiv.org/html/2608.11919#S1.F1) 所示,瓶颈不仅仅是张量无法放入 GPU 内存。瓶颈在于 GPU 计算窗口、CPU-GPU 流量、强制性的参数和梯度移动,以及可能的 NVMe 激活溢出竞争共享资源。通过将张量移动到较慢层级来节省内存的调度,如果这些传输在关键路径上可见,仍然会损害吞吐量。我们提出 LazyTrain,一个面向有限资源 LLM 训练的优化引导调度器。LazyTrain 保留了逐层流式执行器,但将固定激活调度替换为一个混合整数模型,其决策变量选择检查点路径、激活归属、重计算块和通信分配。目标是最小化重计算加上由激活放置引起的增量通信暴露,而不是最大化卸载量。求解得到的策略在训练前加载一次,然后由运行时执行。
我们的评估涵盖了从 3B 到 27B 的单 GPU H800 和 RTX 3090 实验。在匹配的模型、数据和批大小设置下,最终的 LazyTrain 配置达到 219.95 TFLOPS 和 1361 tokens/s,而 MegaTrain 基线为 176.90 TFLOPS 和 1075.8 tokens/s。它的 GPU 内存峰值为 68.84 GB,并在完整评估拆分上获得 95.42% 的精确匹配准确率。LazyTrain-MILP 变体移除了混合整数调度但保留了其他运行时改进,下降到 193.17 TFLOPS。这是最大的组件消融退化,确立了 MILP 调度作为核心贡献。移除 Hybrid 8-bit 算子(将混合 8 位优化器状态与快速梯度裁剪结合)但保留 MILP 调度,则得到 219.29 TFLOPS。更广泛的实验报告了跨模型规模和加速器预算的吞吐量、内存、最大可行批大小和训练质量。
本文的贡献是:
- • 我们将有限资源 LLM 训练形式化为一个混合整数规划(MIP)模型,涵盖检查点选择、激活放置、重计算和 CPU-GPU-NVMe 通信重叠。
- • 我们在逐层流式执行器之上实例化了这个 MIP 调度器,可以将 MegaTrain 基线视为同一调度空间内的固定启发式调度。
- • 我们将耦合的 Hybrid 8-bit 算子作为 LazyTrain 的正式组件:8 位优化器状态减少内存,而快速梯度裁剪抵消其额外的 CPU 侧更新开销。我们在 H800 和 RTX 3090 上跨模型规模评估了完整系统。
## 2 相关工作
#### 扩展 LLM 和智能体工作负载。
最近的 LLM 系列继续推动模型容量、上下文长度和多模态能力(13 (https://arxiv.org/html/2608.11919#bib.bib14);3 (https://arxiv.org/html/2608.11919#bib.bib15);8 (https://arxiv.org/html/2608.11919#bib.bib16);4 (https://arxiv.org/html/2608.11919#bib.bib17))。与此同时,面向智能体的模型和工具针对编码、视觉推理和更长视野的工具使用工作流(14 (https://arxiv.org/html/2608.11919#bib.bib18);18 (https://arxiv.org/html/2608.11919#bib.bib19);22 (https://arxiv.org/html/2608.11919#bib.bib20);2 (https://arxiv.org/html/2608.11919#bib.bib21);7 (https://arxiv.org/html/2608.11919#bib.bib22);12 (https://arxiv.org/html/2608.11919#bib.bib23))。训练和适配这些系统增加了内存、通信和调度压力:更大的模型状态和长上下文智能体轨迹提高了状态和激活足迹,使基础设施成为模型开发的核心部分。
#### 内存高效的 LLM 训练和卸载。
高效的 LLM 训练通过模型并行、混合精度、重计算、状态分片和卸载来降低加速器驻留(5 (https://arxiv.org/html/2608.11919#bib.bib2);17 (https://arxiv.org/html/2608.11919#bib.bib9);11 (https://arxiv.org/html/2608.11919#bib.bib6);9 (https://arxiv.org/html/2608.11919#bib.bib4))。ZeRO、ZeRO-Infinity、ZeRO++、PyTorch FSDP 和 Colossal-AI Gemini 在 GPU、CPU 和 NVMe 层级之间划分或迁移模型状态,使更大的模型可训练(15 (https://arxiv.org/html/2608.11919#bib.bib7);16 (https://arxiv.org/html/2608.11919#bib.bib8);19 (https://arxiv.org/html/2608.11919#bib.bib10);23 (https://arxiv.org/html/2608.11919#bib.bib13);6 (https://arxiv.org/html/2608.11919#bib.bib3))。存储感知系统扩展了这一层级:Ratel 针对消费级 GPU 上的大模型微调优化数据移动(10 (https://arxiv.org/html/2608.11919#bib.bib5));生命周期感知卸载将 GPUDirect Storage 视为训练层级(21 (https://arxiv.org/html/2608.11919#bib.bib11));异构张量缓存增加了资源感知的迁移策略(1 (https://arxiv.org/html/2608.11919#bib.bib1))。MegaTrain 将模型状态存储在主机内存中,将层流式传输通过 GPU 内存,并使用块式重计算来限制激活驻留(20 (https://arxiv.org/html/2608.11919#bib.bib12))。
#### 区别总结。
上述系统通过减少设备驻留、添加异构内存层级或将层流式传输通过 GPU 内存来使有限资源训练成为可能。其中,MegaTrain 通过固定启发式方法选择激活检查点。LazyTrain 解决了一个互补的调度问题:在相同的逐层流式执行空间内,它在一个混合整数模型中联合选择激活检查点、层级放置、重计算和暴露的通信。它还将 NVMe 建模为既是 PCIe 消费者又是独立的读/写端点,因此 SSD 流量仅在可以隐藏或能启用原本不可行的调度时使用。
## 3 方法
### 3.1 概述
LazyTrain 是 MegaTrain 风格逐层流式执行器之上的优化层。执行器定义了可行的运行时动作:将参数从 CPU 内存流式传输到 GPU 内存,计算一个层或重计算块,将梯度卸载回 CPU 内存,并从 GPU、CPU 或本地 NVMe 重新加载物化的激活检查点。LazyTrain 在训练开始前决定每个层边界应使用哪些可行动作。
图 2:LazyTrain 架构。资源配置文件、模型配置和训练目标定义了一个混合整数调度问题。求解器选择检查点边界、GPU/CPU/NVMe 归属和传输窗口;得到的策略由逐层流式运行时执行。
图 2 (https://arxiv.org/html/2608.11919#S3.F2) 总结了该方法。调度器接收模型常量、内存预算、带宽配置文件和训练目标。它求解一个混合整数线性规划(MILP),其中二元变量选择检查点路径和存储归属,连续变量将激活流量分配给计算窗口。求解得到的策略随后由运行时加载。这种设计保持执行简单,同时使调度具有资源感知性。
### 3.2 逐层边界调度空间
考虑一个有 \(L\) 层的 transformer,层边界激活为 \(\mathcal{B}=\{0,\ldots,L\}\)。边界 0 是块输入,边界 \(i\)(对于 \(1\le i\le L\))是第 \(i\) 层的输出。令 \(A_i\) 为边界 \(i\) 处的激活大小,令 \(\rho_j\) 为从边界 \(j\) 到边界 \(j+1\) 的前向转换成本。一个有向边 \((s,e)\),其中 \(0\le s<e\le L\),表示从边界 \(s\) 到边界 \(e\) 的重计算块:前向计算在边界 \(s\) 处开始,并在边界 \(e\) 处结束。边集合是 \(\mathcal{E}=\{(s,e):0\le s<e\le L\}\)。如果边界 \(s\) 和 \(e\) 都是物化的,则边 \((s,e)\) 上的重计算成本为:
\[
R_{s,e}=\begin{cases}0,&e\le s+1,\\\\ \sum_{j=s}^{e-2}\rho_j,&e>s+1.\end{cases}
\tag{1}
\]
该块的最终层输出在边界 \(e\) 处已经可用,因此成本排除该输出。这种边视图将激活检查点转换为路径选择问题。
### 3.3 决策变量和约束
LazyTrain 使用二元变量 \(x_{s,e}\) 选择重计算边,并使用 \(z_i\) 指示边界 \(i\) 是否被物化。如果边界被物化,则二元变量 \(g_i\)、\(c_i\) 和 \(n_i\) 将其分配给 GPU 高带宽内存(HBM)、CPU DRAM 或本地 NVMe:
\[
z_i=g_i+c_i+n_i,\qquad i\in\mathcal{B}.
\tag{2}
\]
执行器将输入和输出边界保留在 HBM 中,因此 \(g_0=g_L=1\) 且 \(c_0=c_L=n_0=n_L=0\)。选定的边必须形成从输入边界到输出边界的一条合法路径:
\[
\sum_{(0,e)\in\mathcal{E}}x_{0,e}=1,\qquad
\sum_{(s,L)\in\mathcal{E}}x_{s,L}=1,
\tag{3}
\]
\[
\sum_{(s,i)\in\mathcal{E}}x_{s,i}=z_i,\qquad
\sum_{(i,e)\in\mathcal{E}}x_{i,e}=z_i,
\]
\[
0
\]
(Note: The original text has "0" at the end, probably an incomplete equation or artifact. We'll preserve as "0" as in source? Actually the source includes "0" at end after the equations. We should translate the text but the "0" might be part of a formula. Let's examine the original:
```
\[
\sum_{(0,e)\in\mathcal{E}}x_{0,e}=1,\\
\displaystyle=1,
\sum_{(s,L)\in\mathcal{E}}x_{s,L}=1,\\
\displaystyle=1,
(3)
\sum_{(s,i)\in\mathcal{E}}x_{s,i}=z_i,\\
\displaystyle=z_i,
\sum_{(i,e)\in\mathcal{E}}x_{i,e}=z_i,0
```
It's garbled from HTML. We should preserve mathematically reasonable version. The original text includes "0" at the end. Since we are translating, we can keep the equations as clean as possible, but to stay faithful, maybe we should represent the math as best as possible. The instruction says preserve markdown content; we don't need to fix errors. But "0" is odd. Let's translate the surrounding text and keep the math as given in a cleaned form? The source is HTML with LaTeX. The final "0" might be part of a sentence? Actually "z_i,0" might be from a missing "?" but we should not invent. We can translate the text before the equations and then keep the equations as provided. For readability, we'll format the equations neatly.
Let's continue with the translation. The original text after the equations: "The selected edges must form one legal path..." That's before the equations. After equations, there is "0" then a period? Let's see: The text:
```
\sum_{(i,e)\in\mathcal{E}}x_{i,e}=z_i,0
```
It seems "0" is an artifact. Probably it's "z_i, \quad 0 \le ..." but truncated. To avoid introducing errors, we can keep the math as close as possible. In the translation, we'll include the equations as displayed, and if "0" remains, maybe keep it. But better to interpret: The original includes "0" at the end of the line. We can treat it as part of the mathematical expression: "z_i,0" might be a typo. Since we're translating, we can preserve the math environment as faithfully. I'll output the equations as in the source, including the final "0" maybe as a standalone? Let's look at the source again:
```
\[
\sum_{(0,e)\in\mathcal{E}}x_{0,e}=1,\\
\displaystyle=1,
\sum_{(s,L)\in\mathcal{E}}x_{s,L}=1,\\
\displaystyle=1,
(3)
\sum_{(s,i)\in\mathcal{E}}x_{s,i}=z_i,\\
\displaystyle=z_i,
\sum_{(i,e)\in\mathcal{E}}x_{i,e}=z_i,0
\]
Actually, in the source, there are LaTeX line breaks. The final "0" is clearly an artifact from HTML conversion. In the translation, we can either omit or keep? I think it's okay to keep the equation block as a representation. Since the surrounding text is what matters, we can format it as:
\[
\sum_{(0,e)\in\mathcal{E}}x_{0,e}=1,\quad
\sum_{(s,L)\in\mathcal{E}}x_{s,L}=1,
\tag{3}
\]
\[
\sum_{(s,i)\in\mathcal{E}}x_{s,i}=z_i,\quad
\sum_{(i,e)\in\mathcal{E}}x_{i,e}=z_i,
\]
And ignore the trailing "0". But that changes the source. As a translator, I should translate content, not fix math. However, the math is not intended to be translated. To maintain meaningfulness, I'll render the equations correctly as they likely appear in the paper. The original "0" probably indicates a constraint like "0 \le x_{s,e} \le 1" but truncated. To be safe, I'll include a cleaner version but preserve all parts. Since the instruction says "Preserve headings, lists, emphasis, inline code, and links when possible" and not necessarily the math, I think it's okay to format the math as standard.
We'll output the equations in a neat way: The constraints (3) are as above. Then the next line "0" might be a stray. I'll omit the stray "0" and just keep the constraints. But to remain faithful, maybe I'll put "0" as a new line? That would be weird. Let's see if the original text has any sentence after that? Actually after the equations, the section ends. So it's fine to close.
Let's proceed with the full translation.
We need to ensure the entire content is translated. The provided text ends with "0" after equations, so we'll translate up to there.
Now, let's write the final translated markdown. We'll use Simplified Chinese. We'll preserve the structure. For math, we'll use LaTeX as in original but cleaned.
Let's also handle the footnotes: "11footnotetext:Equal contribution.22footnotetext:Corresponding authors." In the abstract, the original has "11footnotetext:... 22footnotetext:..." We can translate as "11 同等贡献。22 通讯作者。" at the end of abstract. Or we can use superscript markers. Since the source has "11footnotetext" etc., we'll keep as text.
We'll now produce the output.# LazyTrain:面向大语言模型训练零浪费产出优化的有限资源分配
来源:https://arxiv.org/html/2608.11919
Cehao Yang 隶属:香港科技大学(广州)
Honghao Liu 隶属:香港科技大学(广州)
Xueyuan Lin
Xuhui Jiang 隶属:IDEA Research;DataArcTech Ltd.
Chengjin Xu 隶属:IDEA Research;DataArcTech Ltd.
Jia Li
Jian Guo
###### 摘要
在有限硬件上训练大语言模型越来越成为一个涉及 GPU 算力、主机内存、PCIe 传输和存储带宽的调度问题。现有的卸载系统降低了 GPU 驻留内存,MegaTrain 表明 CPU 主控的逐层流式执行器可以在单张 GPU 上训练大模型,但固定的检查点策略和放置启发式方法仍会让通信暴露在关键路径上。我们提出 LazyTrain,一个构建在逐层流式执行器之上的优化层。LazyTrain 将检查点选择、激活放置、重计算和 CPU-GPU-NVMe 通信重叠建模为一个混合整数调度问题,然后在训练过程中执行求解得到的策略。它进一步将 8 位优化器状态与快速梯度裁剪耦合为单一 Hybrid 8-bit 算子:状态压缩减少了优化器状态内存,而快速裁剪抵消了额外的 CPU 侧更新开销。在从 Qwen2.5-3B 到 Qwen3.6-27B 的 H800 实验中,LazyTrain 相较匹配的基线运行将持续 TFLOPS 提升了约 1.24 倍;RTX 3090 实验同样在每个模型规模下将最大可行批大小增加了一。在主要的 Qwen3.6-27B H800 MetaMathQA 运行中,LazyTrain 在批大小为 72 时达到 219.95 TFLOPS 和 1361 tokens/s,GPU 内存峰值 68.84 GB,并在完整评估拆分上获得 95.42% 的精确匹配准确率。源代码可在 https://github.com/DataArcTech/LazyTrain 获取。
11 同等贡献。22 通讯作者。
## 1 引言
大语言模型训练越来越依赖于系统如何充分利用完整的内存层级,而不仅仅是 GPU 内存。即使对于监督微调,模型参数、梯度、优化器状态和激活也可能超过单个加速器的内存,而 CPU DRAM 和本地存储提供了大得多的容量,但带宽明显更低。因此,高效训练需要一个调度,不仅要决定卸载什么,还要决定每次传输何时可以隐藏在 GPU 计算之下(5 (https://arxiv.org/html/2608.11919#bib.bib2))。
现有的内存高效训练系统解决了这个问题的关键部分。ZeRO 风格的分片和卸载降低了模型状态的设备驻留(15 (https://arxiv.org/html/2608.11919#bib.bib7);16 (https://arxiv.org/html/2608.11919#bib.bib8)),FSDP 和 Gemini 提供了广泛使用的异构内存运行时策略(23 (https://arxiv.org/html/2608.11919#bib.bib13);6 (https://arxiv.org/html/2608.11919#bib.bib3)),而存储感知系统表明 CPU 和 NVMe 可以充当有用的容量层(10 (https://arxiv.org/html/2608.11919#bib.bib5);21 (https://arxiv.org/html/2608.11919#bib.bib11))。MegaTrain 将 CPU 内存视为参数和优化器状态的权威存储,将层流式传输通过 GPU 内存,并使用块式重计算来限制激活驻留(20 (https://arxiv.org/html/2608.11919#bib.bib12))。然而,其激活调度仍然是一个固定的启发式方法。固定的检查点间隔可能是可行的,但它不会联合搜索检查点边界、GPU/CPU/NVMe 归属、重计算块和传输窗口。
图 1:优化引导训练调度的动机。固定的启发式方法可以满足内存上限,同时将 CPU-GPU 或 NVMe 传输暴露在关键路径上。LazyTrain 则搜索检查点放置和传输分配,尽可能在计算窗口内隐藏激活流量。
本文的核心观察是,有限资源训练应被优化为一个受约束的调度问题。如图 1 (https://arxiv.org/html/2608.11919#S1.F1) 所示,瓶颈不仅仅是张量无法放入 GPU 内存。瓶颈在于 GPU 计算窗口、CPU-GPU 流量、强制性的参数和梯度移动,以及可能的 NVMe 激活溢出竞争共享资源。通过将张量移动到较慢层级来节省内存的调度,如果这些传输在关键路径上可见,仍然会损害吞吐量。我们提出 LazyTrain,一个面向有限资源 LLM 训练的优化引导调度器。LazyTrain 保留了逐层流式执行器,但将固定激活调度替换为一个混合整数模型,其决策变量选择检查点路径、激活归属、重计算块和通信分配。目标是最小化重计算加上由激活放置引起的增量通信暴露,而不是最大化卸载量。求解得到的策略在训练前加载一次,然后由运行时执行。
我们的评估涵盖了从 3B 到 27B 的单 GPU H800 和 RTX 3090 实验。在匹配的模型、数据和批大小设置下,最终的 LazyTrain 配置达到 219.95 TFLOPS 和 1361 tokens/s,而 MegaTrain 基线为 176.90 TFLOPS 和 1075.8 tokens/s。它的 GPU 内存峰值为 68.84 GB,并在完整评估拆分上获得 95.42% 的精确匹配准确率。LazyTrain-MILP 变体移除了混合整数调度但保留了其他运行时改进,下降到 193.17 TFLOPS。这是最大的组件消融退化,确立了 MILP 调度作为核心贡献。移除 Hybrid 8-bit 算子(将混合 8 位优化器状态与快速梯度裁剪结合)但保留 MILP 调度,则得到 219.29 TFLOPS。更广泛的实验报告了跨模型规模和加速器预算的吞吐量、内存、最大可行批大小和训练质量。
本文的贡献是:
- • 我们将有限资源 LLM 训练形式化为一个混合整数规划(MIP)模型,涵盖检查点选择、激活放置、重计算和 CPU-GPU-NVMe 通信重叠。
- • 我们在逐层流式执行器之上实例化了这个 MIP 调度器,可以将 MegaTrain 基线视为同一调度空间内的固定启发式调度。
- • 我们将耦合的 Hybrid 8-bit 算子作为 LazyTrain 的正式组件:8 位优化器状态减少内存,而快速梯度裁剪抵消其额外的 CPU 侧更新开销。我们在 H800 和 RTX 3090 上跨模型规模评估了完整系统。
## 2 相关工作
#### 扩展 LLM 和智能体工作负载。
最近的 LLM 系列继续推动模型容量、上下文长度和多模态能力(13 (https://arxiv.org/html/2608.11919#bib.bib14);3 (https://arxiv.org/html/2608.11919#bib.bib15);8 (https://arxiv.org/html/2608.11919#bib.bib16);4 (https://arxiv.org/html/2608.11919#bib.bib17))。与此同时,面向智能体的模型和工具针对编码、视觉推理和更长视野的工具使用工作流(14 (https://arxiv.org/html/2608.11919#bib.bib18);18 (https://arxiv.org/html/2608.11919#bib.bib19);22 (https://arxiv.org/html/2608.11919#bib.bib20);2 (https://arxiv.org/html/2608.11919#bib.bib21);7 (https://arxiv.org/html/2608.11919#bib.bib22);12 (https://arxiv.org/html/2608.11919#bib.bib23))。训练和适配这些系统增加了内存、通信和调度压力:更大的模型状态和长上下文智能体轨迹提高了状态和激活足迹,使基础设施成为模型开发的核心部分。
#### 内存高效的 LLM 训练和卸载。
高效的 LLM 训练通过模型并行、混合精度、重计算、状态分片和卸载来降低加速器驻留(5 (https://arxiv.org/html/2608.11919#bib.bib2);17 (https://arxiv.org/html/2608.11919#bib.bib9);11 (https://arxiv.org/html/2608.11919#bib.bib6);9 (https://arxiv.org/html/2608.11919#bib.bib4))。ZeRO、ZeRO-Infinity、ZeRO++、PyTorch FSDP 和 Colossal-AI Gemini 在 GPU、CPU 和 NVMe 层级之间划分或迁移模型状态,使更大的模型可训练(15 (https://arxiv.org/html/2608.11919#bib.bib7);16 (https://arxiv.org/html/2608.11919#bib.bib8);19 (https://arxiv.org/html/2608.11919#bib.bib10);23 (https://arxiv.org/html/2608.11919#bib.bib13);6 (https://arxiv.org/html/2608.11919#bib.bib3))。存储感知系统扩展了这一层级:Ratel 针对消费级 GPU 上的大模型微调优化数据移动(10 (https://arxiv.org/html/2608.11919#bib.bib5));生命周期感知卸载将 GPUDirect Storage 视为训练层级(21 (https://arxiv.org/html/2608.11919#bib.bib11));异构张量缓存增加了资源感知的迁移策略(1 (https://arxiv.org/html/2608.11919#bib.bib1))。MegaTrain 将模型状态存储在主机内存中,将层流式传输通过 GPU 内存,并使用块式重计算来限制激活驻留(20 (https://arxiv.org/html/2608.11919#bib.bib12))。
#### 区别总结。
上述系统通过减少设备驻留、添加异构内存层级或将层流式传输通过 GPU 内存来使有限资源训练成为可能。其中,MegaTrain 通过固定启发式方法选择激活检查点。LazyTrain 解决了一个互补的调度问题:在相同的逐层流式执行空间内,它在一个混合整数模型中联合选择激活检查点、层级放置、重计算和暴露的通信。它还将 NVMe 建模为既是 PCIe 消费者又是独立的读/写端点,因此 SSD 流量仅在可以隐藏或能启用原本不可行的调度时使用。
## 3 方法
### 3.1 概述
LazyTrain 是 MegaTrain 风格逐层流式执行器之上的优化层。执行器定义了可行的运行时动作:将参数从 CPU 内存流式传输到 GPU 内存,计算一个层或重计算块,将梯度卸载回 CPU 内存,并从 GPU、CPU 或本地 NVMe 重新加载物化的激活检查点。LazyTrain 在训练开始前决定每个层边界应使用哪些可行动作。
图 2:LazyTrain 架构。资源配置文件、模型配置和训练目标定义了一个混合整数调度问题。求解器选择检查点边界、GPU/CPU/NVMe 归属和传输窗口;得到的策略由逐层流式运行时执行。
图 2 (https://arxiv.org/html/2608.11919#S3.F2) 总结了该方法。调度器接收模型常量、内存预算、带宽配置文件和训练目标。它求解一个混合整数线性规划(MILP),其中二元变量选择检查点路径和存储归属,连续变量将激活流量分配给计算窗口。求解得到的策略随后由运行时加载。这种设计保持执行简单,同时使调度具有资源感知性。
### 3.2 逐层边界调度空间
考虑一个有 \(L\) 层的 transformer,层边界激活为 \(\mathcal{B}=\{0,\ldots,L\}\)。边界 0 是块输入,边界 \(i\)(对于 \(1\le i\le L\))是第 \(i\) 层的输出。令 \(A_i\) 为边界 \(i\) 处的激活大小,令 \(\rho_j\) 为从边界 \(j\) 到边界 \(j+1\) 的前向转换成本。一个有向边 \((s,e)\),其中 \(0\le s<e\le L\),表示从边界 \(s\) 到边界 \(e\) 的重计算块:前向计算在边界 \(s\) 处开始,并在边界 \(e\) 处结束。边集合是 \(\mathcal{E}=\{(s,e):0\le s<e\le L\}\)。如果边界 \(s\) 和 \(e\) 都是物化的,则边 \((s,e)\) 上的重计算成本为:
\[
R_{s,e}=\begin{cases}0,&e\le s+1,\\ \sum_{j=s}^{e-2}\rho_j,&e>s+1.\end{cases}
\tag{1}
\]
该块的最终层输出在边界 \(e\) 处已经可用,因此成本排除该输出。这种边视图将激活检查点转换为路径选择问题。
### 3.3 决策变量和约束
LazyTrain 使用二元变量 \(x_{s,e}\) 选择重计算边,并使用 \(z_i\) 指示边界 \(i\) 是否被物化。如果边界被物化,则二元变量 \(g_i\)、\(c_i\) 和 \(n_i\) 将其分配给 GPU 高带宽内存(HBM)、CPU DRAM 或本地 NVMe:
\[
z_i=g_i+c_i+n_i,\qquad i\in\mathcal{B}.
\tag{2}
\]
执行器将输入和输出边界保留在 HBM 中,因此 \(g_0=g_L=1\) 且 \(c_0=c_L=n_0=n_L=0\)。选定的边必须形成从输入边界到输出边界的一条合法路径:
\[
\sum_{(0,e)\in\mathcal{E}}x_{0,e}=1,\qquad
\sum_{(s,L)\in\mathcal{E}}x_{s,L}=1,
\tag{3}
\]
\[
\sum_{(s,i)\in\mathcal{E}}x_{s,i}=z_i,\qquad
\sum_{(i,e)\in\mathcal{E}}x_{i,e}=z_i.
\]相似文章
训练利润最优LLM的理论
本文提出了一种经济模型,结合缩放定律与微观经济学理论,分析大语言模型在利润最优情况下的训练策略,权衡模型质量、训练成本与硬件效率等因素。
DynaTrain: 面向弹性大语言模型训练的快速在线并行度切换
DynaTrain 是一个分布式训练系统,能够在大语言模型上实现亚秒级在线并行度重配置,通过虚拟参数空间抽象,使转换速度比现有方法快多达三个数量级。
无资源,无基准,没问题?评估与改进针对无资源语言的代码生成LLMs
本文通过构建基准测试并提出一种方法,将进一步预训练与权重差异迁移相结合,以更低的成本创建专门的指令遵循模型,从而解决无资源编程语言的代码生成问题。
LiteCoder-Terminal:扩展用于学习语言智能体的长程终端环境
LiteCoder-Terminal-Gen 引入了一种零依赖的合成管道,可生成可执行的终端训练环境,并产出 SFT 和 RL 数据集,使语言智能体在 Terminal Bench 基准测试上取得显著的性能提升。
修剪不安全票:一种资源高效的框架,用于更安全、更鲁棒的大型语言模型
本文介绍了一种资源高效的修剪框架,该框架能够识别并移除大型语言模型中与不安全行为相关的参数,同时保持模型的实用性。该方法利用无梯度归因和彩票假说视角,在最小化性能损失的前提下,显著减少了不安全内容的生成,并增强了对越狱攻击的鲁棒性。