现代机器学习框架支持在多种硬件设备上执行运行,而不同设备间的内存往往是独立管理的,即在某个时刻数据可能存储在不同的设备内存中。那么当网络的不同部分在不同的硬件设备上运行时,各个设备不一定总是拥有最新的数据,因而存在设备间同步数据的问题。本文分析各个框架的内存同步机制。

直观地,设备间的内存同步机制应当达到两个目标:

  1. 尽可能减少内存拷贝。即如无必要,则不进行任何同步。
  2. 尽量减少资源分配。即如无必要,用于描述同一数据的内存不在多个设备上分配。

运行在机器学习框架之上的神经网络是一张有向图,其中的节点是计算操作、边是操作间的张量。当这些节点在不同的设备上运行时,除了“边界”张量所描述的数据可以在不同的设备上有拷贝,其他的张量所描述的数据应当都只有一份拷贝。所谓“边界”张量是指连接两个分别在不同硬件设备上计算的张量。

Caffe 中的 SyncedMemory

SyncedMemory 是 Caffe 中专门负责内存同步的模块。SyncedMemory本质上是个状态机,包含四种状态:

  • UNINITIALIZED:未初始化。尚未分配任何资源
  • HEAD_AT_CPU:CPU 拥有最新的数据
  • HEAD_AT_GPU:GPU 拥有最新的数据
  • SYNCED:CPU 和 GPU 都拥有最新的数据

下面是状态转换图。

Caffe 原生内存同步状态转换

SyncedMemory 同时减少了资源分配和内存拷贝。当张量不是“边界”时,它们的状态转换只有一种。例如当张量在 CPU 端时,只有从 UNINITIALIZED 转换到 HEAD_AT_CPU。这样不需要在 GPU 端分配资源,也无需拷贝数据。数据拷贝只有在从 HEAD_AT_CPUHEAD_AT_GPU 转换到 SYNCED 时才会发生。

Caffe 的 SyncedMemory 看起来简单直观,但在集成更多的硬件设备时会遇到一些麻烦。例如,当我们需要在三种硬件设备间同步数据时,SYNCED 状态的语义就不够明确。一般而言,同步数据的过程是将数据从一个设备拷贝到另一个设备。在这一拷贝过程结束后,数据可能只在两个设备间是同步的,因此我们可能需要引入更多的表示同步的状态。例如,当引入 NPU 之后,我们需要引入三种额外的同步状态:SYNCED_CPU_NPUSYNCED_CPU_GPUSYNCED_GPU_NPU。状态转换图如下。

Caffe 三设备存同步状态转换

上图的转换只是理想中的情况,实际上 GPU 和 NPU 之间可能无法直接同步数据。因为 GPU 和 NPU 这类相关底层库一般只会提供与 CPU 交互的拷贝接口,那么在实践中可能所有的数据拷贝 CPU 都会参与。即从 GPU 到 NPU 的拷贝分两步,首先将数据从 GPU 拷贝到 CPU,然后再从 CPU 拷贝到 NPU。那么在GPU、CPU、NPU 的例子中,便不存在 SYNCED_GPU_NPU 状态,因为 CPU 端的数据也是同步的。这种情况的转换图如下所示。

Caffe 三设备经由CPU的内存同步状态转换

类似地,当增加更多的设备时,需要对 SyncedMemory 进行大量的改造,添加足够多的同步“中间”状态。

在这些状态转换中,只有当第一次到达 HEAD_AT_[CPU|GPU|NPU]时,才会在相应的设备上分配资源。这些资源一旦分配,一般会在 SyncedMemory 对象的生命周期中留存。