来源：DataFunSummit

导读本文将介绍 Oneflow 基于重计算的动态图显存优化工作——Coop。Coop 的核心创新是实现动态图重计算策略和显存分配机制的联合优化。

文章包括以下四大部分：

1. 动态图重计算（DTR）介绍

2. 现有方法的局限性

3. OneFlow Coop 的做法

4. 实验效果展示

分享嘉宾｜张建浩一流科技框架开发工程师

编辑整理｜孙彪北京航空航天大学

出品社区｜DataFun

01

动态图重计算（DTR）介绍

Coop，是对于动态图重计算策略，还有显存分配策略的联合优化。

动态图重计算，简称 DTR。首先介绍一下动态图重计算的背景。

在大部分场景下，神经网络训练过程中占用的显存量，大部分都来自于网络训练的中间特征，而不是网络的参数（即权重）。例如，ResNet-50 的权重总大小就只有不到 100MB，剩下的显存其实都是由训练过程中在前向时产生的那些中间特征占据的。这些中间特征之所以会留在显存里面而没有被及时释放掉，是因为在反向传播的过程中还会再用到它们。

可以看右边的示例图，这个是来自于陈天奇的一篇奠基性的论文 Sublinear 里的一张图。它有两列，左边这一列是前向的过程，从 input 到最后的 loss。右边这一列是反向的过程，从 label 一直到得到所有的梯度。可以看到，在右边反向这一列的反向过程中，会用到前向过程中的一些中间结果，这也就是为什么这些中间特征没有被释放，是因为其实后面它们还会被用到，所以暂时留在显存里面。

在动态图的场景下，对于这些中间特征，我们可以先释放掉一部分的中间变量。然后在后向过程中，等到要用到这些被释放掉的中间变量的时候，再根据前向记录下来的计算过程，把中间变量重新计算出来。这样会付出一定的额外的计算开销，但是减少了在训练过程中所需的显存占用。

动态图重计算是由这一篇 paper 开创的。其思路就类似于 cache。我们常见的 cache 会有一个预先设置好的容量，即 cache 容量。如果 cache 的容量已满，还要再往 cache 里面再插入一个新的元素，就会从这个 cache 已有的元素中挑出来一个，把它删除掉，然后再把新的元素插进去。至于是选 cache 里面的哪一个元素来释放，就可以用一些非常简单但是却很有效的策略，比如常用的 LRU 策略。它的选择策略很简单，就是根据最近一次被访问的时间，即距离现在最近的一次访问过了多久。如果过的时间越长，就认为它是最不应该被 cache 住的，所以会把它给释放掉。

神经网络框架的动态图重计算也是类似的思路，它也是用一个非常简单的判断标准：当显存已经满了的时候，会用一个非常简单的判断标准去找出一个我认为释放代价最低的一个 tensor，把它释放掉。释放掉之后，再去试一试，看空闲显存是不是已经够了，如果已经够了就结束了，如果还是不够，就会去循环地做上述这个过程，直到足够的空闲显存被释放出来，使得后续的显存申请能够成功。

DTR 的判断标准考虑了三个因素：第一个是重计算的开销越小越好，这个好指的是 tensor 的释放代价低，代价越低，它就越应该被释放；第二个是占用的显存越大越好；第三个是上一次访问距离现在的时间越长越好。这 3 个都是很符合直觉的考虑因素。DTR 会综合考虑这 3 个因素去决定哪一个 tensor 释放代价是最低的。其中还有一点是重计算的开销，DTR 为了防止出现很长的一条 tensor 链被连续地释放掉的情况，除了考虑 tensor 自己的重计算开销之外，还会递归地将不在显存内的父代 tensor 的计算开销考虑在内。可以看右边这张图，这个图的 T₀ 到 T₆ 就形成了一个计算图。现在想计算出 T₇，可以从 T₅ 和 T₆ 算出 T₇。这时候如果发现显存不够了，它就会从灰色的，也就是目前在显存里的那些 tensor 里面，去挑一部分释放掉，为 T₇ 腾出空间。这是一个在运行时去动态做的过程。例如，在这张示意图里面，它就挑了 T₂ 和 T₃ 来释放，为 T₇ 腾出了空间。所以算完 T₇ 之后，它的状态就变成了这些，灰色的这四个在显存里面，T₂ 和 T₃ 就已经被踢出去了。

02

现有方法的局限性

在我们的实践过程中，我们发现已有的工作即 DTR 是有一些局限性的。

首先，它是一种带有盲目性的贪心算法，因为每次只会考虑当前最该被释放即代价最低的那一个 tensor，再不断地去做循环直到显存申请能够成功。但是循环的每一步之间，每一步被释放的这些 tensor 之间的关系却没有被考虑进去，就会经常出现这样的场景：比如需要腾出一个 100MB 的空间，开始循环，先挑了一个 60 MB 的 tensor 释放掉，又挑了一个 40MB 的 tensor 释放掉，这时虽然看起来 40 加 60 等于 100，能够容纳 100MB 的 tensor，但是实际上却不是这么简单。因为在显存里面我们对显存的连续性也是有要求的。我们单一地分别释放一个 60MB 的和 40MB 的 tensor，很可能它们释放出来的空闲显存是两段分散的。这样分散的显存我们是用不了的，也就是常见的显存碎片，而不是能够容纳 100MB tensor 的一段连续的大显存。既然 DTR 用不了这两个碎片，那么这个算法还会继续去循环，去找其他的 tensor 再释放，直到正好释放了一整个 100MB 的 tensor，或者两个或多个小的 tensor 正好一起凑够了 100MB 的连续显存，这显然是一个不太好的性质。

第二个局限性就是它的重计算算法和 tensor 在显存池里面的排列方式是有一定冲突的。

一般来说，假设网络里面有一个直线型 OP 序列，直线型 OP 序列产生的 tensor a，b，c，d 被放在显存池里。对于显存池里的一个 tensor，它被释放的时候，它的显存不会直接归还给系统，还会在这个显存池里面，显存池会把它给 hold 住。后续如果再有人要新的显存，就直接从显存池里 hold 住的那些显存里面去直接分配，就不走操作系统了。也就是说，显著池会存在一个很大的空闲显存块。这样这些顺序的 tensor a，b，c，d 就会在很大的空闲显存空间里面连着排。所以如果 tensor 在网络里面是连续的，那么它在显存池里面即物理上的内存排列按照默认规则也会是连续的。这样会引起的问题可以参考下图。

（a）图是一个比较典型的网络结构，即 relu-conv，这里就省略了 BN，因为画起来要简单一些，但是不影响本质的分析。一般来说，它在显存池里面的排列也会是按照这个顺序 relu-conv-relu-conv…。此时如果新来一个 X₅，它的显存大小是 2 倍的小方块的长度，我就要挑两个 tensor 释放。如果按照经典的 DTR 的做法，它会把 X₀ 和 X₂ 这两个 tensor 先释放掉。这个背后有两个原因。一个是比较显而易见的，因为 X₀ 和 X₂ 它们都是 relu 产生的，relu 的重计算代价比 conv 的重计算代价要低很多。另一个原因是我先挑了 X₀ 来释放，前文讲到，计算重计算代价的时候，里面的 c，就是重计算的开销，它不仅是 tensor 本身的计算开销，也会包括不在显存内的那些父代 tensor 的重计算开销。也就是说当我释放了 X₀ 之后，它的子代 tensor X₁ 的重计算代价也会增加，使得 X₁ 更加难以被释放了。能看出来这其实是一件很矛盾的事情，我们想要的是 X₀ 和 X₁ 都被释放掉，这样才能够形成一段连续的空闲显存。

但是 DTR 对重计算代价的设计和 tensor 在显存池里的排列方式却一起作用，使得我们反而更加难以去形成连续的显存，这就是它的第二个局限性。它带来的结果就是，可能先释放掉 X₀，X₂，接下来还得再释放一个，才能够形成一段连续的显存容纳 X₅。如果运气好，可能下一次就选到了 X₁，但是也有可能又选到了 X₄ 或者后面别的 tensor，那就更糟糕了，因为可能要释放很多无辜的 tensor，才能够碰巧形成一段连续的空闲显存。

第三个局限性是动态图重计算的传统方法 DTR 里面，对 in-place op 也是有一定影响的。因为在重计算的时候，相当于是在重新计算历史上的 op。in-place op 会就地修改掉一个已有的 tensor 的值，而不是产生新的 tensor。这样一来，在重算的时候，拿到的数据可能就是一个已经被修改过的数据，而不是当时所用的数据了，即现在的数据可能已经被 in-place op 给修改过了。这样它们就会有一个相互的作用。

参见上图，现在有这样的一个计算轨迹，从 p 到 x 到 y，分别减 2 加 1。现在如果来了一个 in-place op，对 x 做 in-place relu。本来计算的轨迹是减 2 加 1，比如 y 被释放了，如果被释放了，要重计算 y，从 x 加 1 就能把 y 再重计算出来。但如果对 x 做了 in-place relu 之后，x 的值就变了，这个时候再对它加 1，得到的就不是我想要的 y 了，数据就错了。在 DTR 的做法里面，引入了一个 copy-on-write 的中间层，实质上把 in-place op 变成了一个非 in-place op，也就是 out-place op。它的具体做法是：申请一块新的空间去放置一个 x’。如果对 x 做 in-place 操作，它会申请新的空间放置 x’。在 x 和 x’ 之间去做一个普通的 op，也就不 in-place 了。这样一来，在计算图或者在计算轨迹上面，x 和 x’ 就变成是分离的两个没有关系的 tensor，这样就能保证它是可以被正确地重计算的。

这样，首先，它失去了 in-place op 的优点，比如节省显存，对 cache 更加友好，这些优点全都没有了；另外，在模型训练过程中，模型的权重是不能被释放的，释放的代价实在是太高了，而模型权重通常都会 in-place 地被更新，如果我们也用这个 copy-on-write，也就是不 in-place 更新了，而是 out-place 更新，就会生成很多新的 tensor，这些新的 tensor 就会分散地分布在显存池里面的各个地方，这样就相当于把显存池里面的显存给割裂了，因为这些 tensor 是不能被释放的。这样就更加难以在显存池里面形成一段很长的连续的空闲显存了。

以上就是 DTR 的局限性。在我们的工作里面，联合考虑两件事情，一个是显存的排列方式，即 tensor 的显存在显存池里的排列方式，以及重计算的问题，就能够把上面的三个局限性都解决掉。

03

OneFlow Coop 的做法

首先，我们考虑了显存的排列方式。不再是盲目地去循环，每次找出代价最低的一个 tensor 了，而是整体的考虑，当要释放一个 tensor 的时候，它的周围是怎样的，相当于是在找出一个 tensor 的集合，释放这个 tensor 集合之后，能够形成足够的连续显存，而且这个集合的释放代价是最低的，也就是上图中的公式。后面的 subject to 是一个限定条件，这个 M(S, L) 限定条件是指在显存排列方式 L 下，释放掉集合 S，集合里面的所有 tensor 能够形成期望的连续空闲显存的长度。在满足限定条件的情况下，想找出最优的 S 和 L，这个最优即集合 S 里面的所有 tensor 的释放代价加起来最小。

基于上述目标，我们提出了三个模块，分别是 recomputable in-place，op-guided tensor allocation 和 layout-aware eviction。

上图是显存池的状态，即内存的排列状态。显示池里面绿色的是释放代价低的 tensor，蓝色的是释放代价高的 tensor，红色的是不可被释放的 tensor，这些红色的 tensor 主要是 parameters、buffers 等。可以看到最开始的时候，这 3 类 tensor 是交叉随机排列的。经过 recomputable in-place 之后，这些不可释放的 tensor，会一直留在显存池最边缘的地方，这样就能够促使我们形成足够的连续显存。再接下来，经过 op-guided tensor allocation，就能够把 cheap 的 tensor 还有 expensive 的 tensor 给分散到两边。这样就很容易找出释放代价低的 tensor 也就是 cheap 的这部分 tensor 了。优化 tensor 的排列方式，如果对应到前面的公式，就是调整公式里面的L，再通过 layout-aware eviction，也就是通过滑动窗口，就可以找出想要的一个连续 tensor 的集合。

对于 recomputable in-place，首先它是可重计算的，与 naïve 计算方式不同。并且，它又是 in-place 的，具有 in-place op 的那些优势，比如节省显存和对 cache 友好。

具体做法是受启发于一个 PL 领域的工作，一个叫 Koka 的编程语言，其中引入了一个 functional but in-place 的机制，类似于尾递归优化。这是一种通用的优化方式，具体来说，假如已经知道某一个变量接下来不会被用到了，就可以把新的变量直接分配到已经知道不会被用到的变量的内存地址上面，这就相当于对已经确定未来不会被用到的变量的内存做一个复用，这样就达到了类似于 in-place 的效果。但是如果单看变量，前后确实是不同的两个变量，但这两个变量共享了一个内存空间，因为他们的生命周期不重叠。

Coop 里面的 recomputable in-place 正是受此启发，不过这里不是变量，而是 tensor，但核心思想类似，让两个 tensor 暂时共享同一个内存空间。现在在这个 recomputable in -place 里面，虽然也是有两个不同的 tensor X 和 X’，在这两个 tensor 之间，表面上是在做一个非 in-place 的操作，去形成一个正确的计算轨迹。但是实际上，因为这两个 tensor 共享了同一个内存空间，所以效果是和 in-place OP 完全一样的。等这个 in-place OP 结束了之后，再和原来 DTR 的 copy-on-write 一样，把输入 tensor 设置为已释放的状态，最终达到的状态和 DTR 的 copy-on-write 完全一样。但是我们仍然保留了它 in-place OP 的良好的性质，也就避免了前面所提到的 copy-on-write 的 2 个缺点。

接下来是 op-guided Allocation。我们根据产生 tensor 的 op 的性质去指导我们把 tensor 放在显存池里面的什么地方。我们发现如果计算 cost density，即代价的密度，也就相当于是单位显存下重计算的开销。可以看到 Conv 和 MatMul 跟其他的 op 相比，代价密度有一个数量级上的差别。这就引导我们把 Conv 或 MatMul 生成的 tensor 和其他 op 生成的 tensor 划分成两类，分别放在不同的位置。在我们的实现里，我们利用了重计算的一个特性，即最大显存阈值是预先给定的，这样我们就可以从左右两端去放 tensor，而不是只能一直放在左边。具体来说，比如 x₀ 它是 ReLU 产生的，我们把它放在左边，而 x₁ 是 Conv 产生的，放在右边，以此类推。这样代价密度小的 tensor 在一边，代价密度大的 tensor 在另一边，我们就可以更加容易地找到代价之和更低的 tensor 集合。

同时它也避免了前文提到的自相矛盾的问题，即 x₀ 的释放反而增加了 x₁ 的代价，导致 x₀ 和 x₁ 难以同时释放形成连续的空闲显存，我们的放在两边的方式就可以避免这个问题，因为现在 x₀ 和 x₁ 在显存池里面就已经不再连续了。虽然 x₀ 释放以后，x₁ 的代价还会上升，但是已经不影响连续空闲显存的形成了。

最后一个模块是 layout-aware eviction，滑动窗口。为了能够做滑动窗口，我们做了一些预处理，把问题转化了一下。具体来说，在显存池里面，我们把那些已有的空闲的显存块看作一个代价为 0 的特殊 tensor，这时候 tensor 就包括那些空闲的显存了。把所有的 tensor 按照在显存里面的地址排列成一个 list，如果两个 tensor 在 list 里面是连续的，那么在物理显存上面也是相邻的。问题就变成了如何在这个 list 里面找一个连续的子序列，子序列要满足显存之和大于预先给定的阈值 M_R，还希望它的代价之和最小。这其实是一个很经典的算法题，有两个指针，每一次都判断显存之和有没有大于 M_R，如果小于，就把尾指针往右移一格，如果大于，就把头指针往右移一格。在移动的过程中，边移动边记录当前的两个指针之间的 tensor 的 ht 即显存之和，同时记录 ht 之和最小的一组 tensor。于是一次遍历就能够得到最优的 tensor 的集合，而不像传统的 DTR 方法要去不断地循环遍历。从理论上可以计算出来，如果期望的显存大小长度符合一个均匀分布，DTR 方法的时间复杂度会是 O(n²)，这个 n 指的是显存池里面有多少个 tensor，但是我们的 Coop 方法的时间复杂度则只有 O(n)，即只需要线性的时间就能够找到想要的集合。

04

实验效果展示

接下来看一下我们的实验效果。蓝色的是 Coop，绿色的是 DTR，黄色的是 DTE，DTE 是 MegEngine 团队对 DTR 做的一个改进。横轴是显存的比例，假设不重计算的时候显存占用量是多少，现在如果只想用不重计算时的80%的显存，那么横轴就是 0.8。纵轴是计算时间的比例，以某一个显存阈值重计算之后，计算时间扩大了多少倍，比如扩大 1.1 倍、1.2 倍，线越往左下角越好。可以看到 Coop 在 8 个网络里面都比较明显地超过了另外两条线。

另外比较有趣的一点是，我们统计了显存碎片率，即在显存池里面空闲显存占的比例。可以看到 Coop 有一个数量级的提升。BiLSTM 是对数坐标，其他两种方法显存碎片率已经快要爆表了，但是 Coop 仍然是一个很低的状态，只有 10% 左右。另外三个网络也是类似的效果，Coop 的显存碎片率都是很低的。

前文提到，去找 tensor 集合的过程也是有时间开销的。传统方法的时间复杂度是 O(n²)，而 Coop 方法是 O(n)，在实验中也能够得到证实。在大部分情况下，Coop 去找 tensor 集合的时间即搜索时间都会比其它两种方法要好。

目前我们刚刚投了一篇 paper，对应的代码还没有合并到 OneFlow 的 master 分支里面去，不过预计很快。大家如果感兴趣可以扫码加入我们 OneFlow 的微信群或者 QQ 群。

05

问答环节

Q1：OneFlow 现在有使用 DTR 的场景吗？印象里 OneFlow 是静态的图。

A1：曾经 OneFlow 是静态图，现在我们 OneFlow 是动态图静态图都有。个人觉得未来的趋势动态图仍然会是主导地位。就像 PyTorch 一直是 eager-first 的策略，也一直是统治性的框架。而我们目前的 API 也是追求和 PyTorch 百分之百完全兼容，你可以直接 import oneflow as torch。相当于可以认为我们是 PyTorch 很好的一个增强版。

Q2：OneFlow 支持动静态转换吗？

A2：支持的，用我们的 nn.Graph 包一层就可以。因为默认的时候我们和 PyTorch 一样是静态图，我们有自己的 nn.Graph，用 nn.Graph 把 nn.module 包一层，就可以把动态图变成静态图。

Q3：Coop 是用计算换显存资源吗？主要的应用场景是什么？大模型训练吗？

A3：对，可以这么说。大模型训练是一个典型的场景。也有一些场景比如：可能有些个人用户，他的显卡不是很好，可能显存很少。但是那种常见的模型，比如检测模型，可能需要很大的显存量才能够训练起来。这种场景下虽然不是大模型，但是对于这种卡的显存很少的一些普通用户来说，也是会有很大的帮助。

Oneflow 基于重计算的动态图显存优化实践