CUDA 有 unified memory 还需要内存优化吗？

CUDA 有 unified memory 还需要内存优化吗？

先说结论，不管有没有 unified memory，只要有性能要求的软件，都必需自己做优化。

因为 unified memory 的主要工作，并不是去掉了 Host Memory 与 Device Memory 之间的拷贝，而是将原有的需要人工手写的这部分代码，交给了运行时，使得这些工作对程序员透明，但是 Host2Device、Device2Host 的内存拷贝所造成的 overhead 依然存在。

其次，“内存优化”其实是一个非常宽泛的话题，他通常可以指内存生命周期的自动管理（如各种解释型语言使用的垃圾回收机制）；也包括减少动态内存重复申请、释放频率提高效率；还包括用更少的内存占用量完成一样的功能（当年阿波罗登月飞船上的计算机内存只有4KB）等。

以深度学习框架 OneFlow 为例，研发团队在不依赖 unified memory 的前提下，就在框架中自己实现了对开发者透明的 Host/Device 内存管理。

在 OneFlow 中，开发者只需要使用简单的表达式：user_op::HobDeviceTag() == DeviceType::kGPU 或者 user_op::HobDeviceTag() == DeviceType::kCPU 就可以告之框架相关功能是使用 GPU 内存还是 CPU 内存，编程过程中拿到的内存地址就是对应内存空间的地址。那些烦人的内存分配、设备之间的内存拷贝、动态内存的生命周期管理统统不用操心。

除了编程易用性外，怎样使用更少的内存完成更多的工作量也是内存优化重点。很多人会认为这块应该交给硬件厂商做，同时也常认为：如果硬件厂商做得足够好，就没有软件开发者什么事情了，遗憾的是，这是一个无法达到的理想境界。因为极致的内存优化与业务逻辑强相关，而硬件厂商要兼顾通用。

举个业务逻辑相关的例子，比如做深度学习框架，假设我们有 N 个参数，那么在反向过程中，需要为这 N 个参数分配空间存放梯度。一个新手在一次训练迭代中，可能会这样管理内存生命周期：

• 获取参数信息 N
• 根据参数信息 N ，分配空间 R
• 反向传播，更新梯度、更新参数
• 释放空间 R

当第二次训练迭代过程中，会重复以上过程。然而，只要稍加分析，就可以发现，在一次迭代结束后，不必要立即释放空间 R，而可以留给下一次迭代复用。这是一个很简单的例子，但是硬件厂商无法为此”擅作主张“，需要结合业务逻辑分析。

在实际的实现中，业务逻辑对代码复用的影响就更加复杂，而优化结果的回报自然更加丰厚。

可以有多丰厚呢？以深度学习 OneFlow 为例，在一流的算法工程师和框架工程师的配合下，同样的任务，OneFlow 和英伟达的亲儿子 HugeCTR 比较，OneFlow 内存节省惊人。

对于不熟悉 HugeCTR 的同学，简单介绍下，它是英伟达为解决大模型训练问题开发的框架，在 OneFlow 发布之前，HugeCTR 性能一直是世界第一。

以下是跑 Wide&Deep 不同 batch size 的的对比结果，可以看到，从 batch szie 一万到百万数量级，OneFlow 的内存占用率都只有一半不到，最少时只占 24%，四舍五入就是，用 OneFlow 买一张显卡可以完成的工作，换其他人得买4张显卡。当 batch size 达到 2097152 之后，就无法比较了，因为 OneFlow 还可以顶一下，HugeCTR 已经内存溢出（OOM Out of Memory）。

综上所述， unified memory 简化了编程的过程，但是我们不能期待它可以帮助我们一劳永逸地解决内存优化的问题。反过来说，恰恰是因为特定的领域（包括但不限于深度学习框架开发）还存在着许多颇具挑战、回报丰厚、悬而未决的问题，才让这些领域充满魅力。

最后，安利一下 OneFlow 的测试仓库 DLPerf。