针对大语言模型训练中普遍存在的显存不足核心瓶颈，我院王智彬老师课题组联合华为技术有限公司，提出了名为 SmartSwap 的端到端交换优化系统。作为首个能够适配 PyTorch Eager Mode 下算子序列动态变化的交换优化方案，实验数据显示，SmartSwap 可降低 84.25% 的算子信息采集开销，最高支持训练规模达到硬件内存 4 倍的模型，相较传统方案性能提升最高可达 38.94%。目前，该研究的相关论文已被计算机体系结构领域顶级会议 DAC 正式接收。

论文标题： SmartSwap: Swap-Based Memory Optimization for LLM Training under Varying Operator Sequences

01 研究背景：大模型时代的显存压力

近年来，以 Llama、Qwen 系列等为代表的大语言模型展现出了惊人的能力，但其参数量的激增也带来了巨大的显存压力。单个 AI 加速器（如 GPU/NPU）的显存往往难以容纳千亿参数模型的训练需求。

虽然模型并行、梯度累积等技术可以缓解压力，但它们引入了复杂的通信开销。交换（Swapping） 是一种有效的优化手段——它将不常用的激活值临时迁移到主机（Host）内存中，需要时再换回。这种方法既不损失精度，也不增加计算量，理论上是完美的内存“扩容”方案。

然而在实际应用中，PyTorch 框架通常运行在 Eager Mode下。与静态图模式不同，Eager Mode 为了追求灵活性，允许训练过程中的算子序列动态变化。这种动态性直接打破了传统交换技术“一次采集、永久适用”的静态假设，导致换入换出时机错乱，甚至引发训练崩溃。

02 核心挑战：动态算子序列下的三重困境

为了在 Eager Mode 中实现高效的 Swap，团队面临三大技术难题：

·如何低成本地持续“监控”？

传统的 Profiler（性能分析器）开销巨大。PyTorch 内置的Profiler会导致训练速度下降 219%，无法在训练过程中长期开启。如果无法持续监控，就无法感知算子序列的变化。

·如何在信息缺失下制定策略？

为了降低监控开销，不能收集每个算子的精确执行时间。在缺乏详细时序信息的情况下，依然生成高效的交换策略具有挑战。

·如何精准地执行策略？

Eager Mode 下，算子没有唯一的 ID。如果上一轮的策略直接搬到下一轮，由于序列变化，可能会引发内存错误。

这些问题在现有的工作中亟待解决。团队对相关问题广泛调研、分析、思考后，给出了SmartSwap解决方案。

03解决方案：SmartSwap 设计

为了解决上述问题，团队提出了 SmartSwap，这是一个包含轻量级在线 Profiler、策略生成器（Policy Generator）和执行器（Executor） 的完整系统。

1. 轻量级在线 Profiler

SmartSwap 设计了两种模式的 Profiler：

轻量模式（Low Overhead Mode）： 仅记录算子序列的整数哈希值，用于实时检测序列是否发生变化。其开销仅为 0.9%，几乎可以忽略不计。

详细模式（Detailed Mode）： 当检测到序列变化或需要新策略时，短暂开启，收集必要的内存和算子信息。

系统通过一个 阶段调整模块（Stage Adjusting），在 WarmUp（预热）、GenPolicy（生成策略）和 Stable（稳定运行）之间自动切换，实现了“按需 Profiling”。

2. 策略生成器：基于“逻辑层”的估算艺术

既然无法获取每个算子的精确时间，SmartSwap 采用了分组估算的策略。将算子序列均匀分组，形成“逻辑层（Logical Layers）”。利用统计学观察：只要分组数量合理，每组的执行时间方差极小。因此，可以用迭代总时间的均值来估算每组的时间。

基于此，系统构建了内存缩减列表（MRL） 和候选列表（CL），通过模拟器计算出最佳的换入（Swap-in）和换出（Swap-out）时机。

3. 执行器：定制化的“交通指挥”

多特征模糊匹配： 为了解决算子 ID 缺失问题，Executor 使用算子名、调用栈、数据类型等多维特征进行模糊匹配，即使序列微调也能精准定位张量。

定制化RecordStream： PyTorch 原生的 recordStream 机制依赖频繁的主机查询，容易导致Host瓶颈。SmartSwap 利用模拟器预知的完成时间，设计了定制化的流同步机制，消除了不必要的轮询，大幅降低了 CPU 开销。

04 实验评估：模型规模扩展与性能飞跃

SmartSwap 基于 PyTorch-NPU 实现了超过 8700 行代码，并已在生产环境部署一年。实验基于 Ascend 910B NPU 进行，取得了以下结果：

1. 极致的扩展能力

SmartSwap 展现了高线性扩展能力。在 Batch Size、序列长度和 Hidden Size 三个维度上，SmartSwap 均能支持训练远超硬件内存限制的模型：

Batch Size： 支持高达硬件内存 4倍 的模型。

序列长度： 支持高达硬件内存 4倍 的模型。

Hidden Size： 支持高达硬件内存 1.24倍 的模型。

2. 卓越的性能表现

对比全量重计算（Full Recomputation）： SmartSwap 平均提升了 16-19% 的性能。

对比高并行度方案： SmartSwap 最高可带来 38.94% 的性能提升（因为减少了昂贵的跨设备通信，转而利用 PCIe 带宽）。

算子信息采集开销极低： 相比 PyTorch 内置 Profiler 高达 219.7% 的开销，SmartSwap 的详细模式开销仅为 34.6%，轻量模式低至 0.9%。

3. 内存复用效率飞跃和消除性能抖动

相比原生机制，SmartSwap 将内存块的复用间隔缩短了 2-3 个数量级。实验表明，原生方案在模型增大时会出现显著的性能波动，而 SmartSwap 实现了近乎完美的线性扩展，彻底消除了因流同步带来的性能抖动。

4. 长期稳定性

在连续 5000 步的训练中，SmartSwap 的 Loss 曲线与全量重计算完全重合，证明了其在动态环境下不会引入数值误差或内存泄漏。

05 结语与展望

SmartSwap 成功打破了 Eager Mode 动态训练与静态内存优化之间的隔阂。它证明了在不改变模型代码、不损失精度的前提下，通过精细化的系统设计，我们可以利用廉价的主机内存来弥补昂贵的显存缺口。

这项技术不仅适用于华为 Ascend NPU，其核心设计理念（逻辑层分组、定制流同步）对于解决 NVIDIA GPU 等通用硬件上的大模型显存溢出问题同样具有极高的参考价值。未来，团队计划开源该系统，助力大模型训练的普惠化与高效化。

对我们工作感兴趣的老师和同学，请联系:wzbwangzhibin@gmail.com