【论文笔记】【存储】FlashNeuron: SSD-Enabled Large-Batch Training of Very Deep Neural Networks

[offloading data to SSD] FlashNeuron: SSD-Enabled Large-Batch Training of Very Deep Neural Networks. Jonghyun Bae, Seoul National University; Jongsung Lee, Seoul National University and Samsung Electronics; Yunho Jin and Sam Son, Seoul National University; Shine Kim, Seoul National University and Samsung Electronics; Hakbeom Jang, Samsung Electronics; Tae Jun Ham and Jae W. Lee, Seoul National University. FAST’21

论文地址：https://www.usenix.org/system/files/fast21-bae.pdf

---

总的来说

增加批次大小是提高硬件利用率的一种流行技术，这可能导致训练吞吐量的亚优化。
文章中提到了这一点，我觉得可以记一下。

本文也是解决GPU内存不足以训练大型模型的问题的。
本文是第一篇使用SSD作为交换内存来解决这个问题的。

最近的提案通过将一些中间数据（例如，特征映射）卸载到主内存来解决这个问题。然而，它们未能提供稳健的性能，因为GPU上的训练过程与CPU上运行的应用程序争用内存带宽和容量。

因此，我们提出了FlashNeuron，这是第一个使用NVMe SSD作为后备存储的DNN训练系统。为了充分利用有限的SSD写带宽，FlashNeuron引入了一个卸载调度器，该调度器选择性地将一组中间数据以压缩格式卸载到SSD上，而不增加DNN评估时间。FlashNeuron对CPU进程的干扰最小，因为GPU和SSD直接进行数据传输。

本文一个重要的观点就是，cpu-offload会干扰到cpu上运行的程序，所以要用SSD

我们对FlashNeuron和四个最先进的DNNs的评估表明，FlashNeuron可以将批次大小增加12.4倍至14.0倍，超过NVIDIA Tesla V100 GPU 16GB DRAM的最大允许批次大小。通过采用更大的批次大小，FlashNeuron还提高了训练吞吐量，最高可达37.8%（平均为30.3%），而且只会对CPU上运行的应用程序产生最小的干扰，这是仅使用GPU内存的基线。

一些现有工作的问题

1. 用多个GPU可以部分绕过内存容量壁垒，因为多GPU的精细使用可以实现接近线性的吞吐量提升。然而，这种吞吐量的提高伴随着GPU成本的线性增加，而这通常是整个系统成本的主要组成部分。因此，使用多个GPU通常会导致次优的成本效率（即，吞吐量/系统成本），因为每个GPU由于每个GPU的批次大小限制，没有在其完全容量下运行。

2. 最流行的方法是使用主机CPU内存作为后备存储，卸载一些不立即使用的张量。然而，这种基于内存的缓冲方法未能提供稳定的性能，因为GPU上的训练过程需要与CPU上运行的应用程序争用内存带宽和容量（例如，数据增强任务以提高训练精度）。此外，这些方案主要关注增加批次大小，而不是提高训练吞吐量。因此，由于CPU-GPU数据传输的成本超过了更大批次的好处，它们通常产生低训练吞吐量。

   • 常见的克服GPU内存容量瓶颈的方法是将数据缓存在主机CPU内存中。例如，vDNN [这篇值得一读，先记个TODO] 和 SuperNeurons [见前博客] 都选择性地将激活张量卸载到CPU内存。这些在内存上缓冲的解决方案可能会与CPU进程对内存带宽和容量的需求产生冲突，从而付出显著的机会成本。例如，每次DNN训练迭代时在CPU上进行数据增强是一种常见的做法，以防止对训练数据集过度拟合。典型的数据增强流程包括图像加载、解码和一系列的几何变换，这些都需要高内存带宽。
     （这一点之前没有注意到，是一个切入点）

3. 本文也没有提到优化器状态占用的情况，本文用的术语叫中间结果，我认为优化器状态应该就是在这里面的，可见下图。
   

本文的工作

FlashNeuron引入了一个卸载调度器，它明智地选择一组张量卸载到SSD。在主机侧，FlashNeuron是通过一个轻量级的用户级I/O堆栈实现的，由于GPU和SSD直接利用GPUDirect技术进行张量数据传输，因此它在CPU周期和内存使用上只占用最小资源。

具体来说，我们利用GPU和NVMe SSD设备之间的直接对等通信，使得数据缓冲不消耗主机CPU周期或内存带宽。这种在SSD上缓冲的方法补充了流行的在内存上缓冲的方法，从而在各种使用情况中提高了整体资源利用率。

对批量的见解：

• 通过增加批量大小，仍有大量的吞吐量提升空间。即使在只使用GPU内存的情况下，GPU资源在最大的每个GPU批量大小下仍然被低利用。在这种情况下，通过改善GPU内存容量瓶颈，可以提高GPU的吞吐量。**可能的一个担忧是，较大的批量大小有时可能会负面影响模型的准确性。**然而，对于极深的神经网络模型，基础批量大小相对较小，因此预计增加批量大小不会严重影响最终模型的准确性。
• 所以这个批量要大还是小还是一个很重要的超参的。在很大的模型里，batch size大点没关系。

FlashNeuron

FlashNeuron由三部分组成：卸载调度器、内存管理器和对等直接存储访问

具体来说，卸载调度器确定了一组需要卸载的张量（即，多维矩阵），并通过考虑多种因素生成一个卸载调度，如张量大小、张量传输时间以及前向/后向传播运行时间。一旦确定了调度，内存管理器就使用对等直接存储访问来协调GPU内存和SSD之间的数据传输，以最小化卸载带来的性能开销。

内存管理器

这部分涉及到张量的分配/释放，卸载和预取的管理，以及对张量进行压缩和精度降低的处理。

张量分配/释放

由于频繁的GPU内存分配和释放会带来显著的性能开销，FlashNeuron使用了一个定制的内存分配器。这个内存分配器首先预留了整个GPU内存空间，并自行管理内存的分配和释放。

1. 在FlashNeuron中，当

• i)在前向传播过程中第一次创建一个张量，或者
• ii)一个被卸载到SSD的张量在反向传播过程中被预取到内存时，就会分配张量。

2. 另一方面，当

• i)一个张量完全从内存卸载到SSD，或者
• ii)在迭代过程中一个张量不再被任何层使用时，就会释放张量。

为了追踪张量的生命周期，使用了一个引用计数机制（在PyTorch [50] 和 TensorFlow [1]中使用）。对于使用静态计算图的DNN框架，如Caffe，内存管理器会遍历计算图并通过追踪每一层所附带的张量的指针来追踪张量的生命周期。

3. 碎片化
   在张量的分配和释放中，一个关键问题是碎片化。为了避免这个问题，我们**从内存地址空间的最低端分配常驻内存（即，未被卸载）的张量，从最高端分配短暂存在（即，被卸载）的张量。**由于短暂数据在前向传播过程中的生命周期非常短，因此只有一小部分的内存地址空间被这些数据使用，从而使得碎片化的内存空间变得可以忽略。

管理卸载和预取：

内存管理器与对等直接存储访问（P2PDSA）进行交互，进行卸载和预取。

• 它在前向传播过程中，一旦一个张量被下一层使用后，就向P2P-DSA发起一个卸载请求。
• 在每一层执行结束后，内存管理器检查卸载请求是否已完成（即，张量是否已完全卸载到SSD）。
• 然后，此时张量被从GPU内存中释放。
• 在反向传播过程中，内存管理器向SSD发出预取请求。
• 在反向传播开始时，它首先为即将使用的一组张量分配内存并发起预取请求。
• 然后，每当这些张量被使用和释放时，内存管理器会预取额外的张量(这个思想在观察到DNN迭代具有一定的重复模式的论文的思路很像)，同时保留足够的内存空间来运行最大的层。

增强的压缩稀疏行（CSR）压缩和解压缩

一个关键的观察是：ReLU的输出是非常稀疏的，对于训练过程（43%~75%）稀疏率。

作者的方法是这样的：

• 他们将张量转换为一个二维矩阵，其列有128个条目。

• 然后，他们使用了一种增强的CSR格式。在原始的CSR格式中，我们需要存储每个非零元素的值和它在矩阵中的位置（行索引和列索引）。但在这种增强的CSR格式中，作者用一组位向量（每个位向量表示一行的非零元素集合）来代替存储每个元素的列索引的向量。这样，存储的大小就减小了8位（表示列索引）× 矩阵中非零元素的数量。

• 这种表示方式只有在超过八分之一的所有元素是非零元素时才是有益的。因为在这种情况下，你节省的存储空间（不必存储那么多的列索引）比你增加的存储空间（必须存储位向量）要多。

作者在GPU中实现了一个专门的程序来执行这个增强的CSR压缩/解压缩。根据他们的评估，这些压缩/解压缩操作的运行时开销可以忽略不计。

使用半精度浮点数（FP16）表示卸载的张量

为了进一步减少GPU和SSD之间的通信，内存管理器利用了神经网络可以在不显著降低最终模型精度的情况下容忍一定程度的精度损失的事实。具体来说，在前向路径中，内存管理器首先将卸载的张量转换为FP16格式（从FP32），然后存储在SSD中。然后，在反向传播过程中，预取的FP16张量被填充到FP32格式并重复使用。

对性能的损失是不大的，现在很多论文都在做混合精度训练，尤其是Deepspeed的ZeRO

卸载调度器

FlashNeuron 中的卸载调度器（Offloading Scheduler）的工作方式。这个调度器接收一个深度神经网络模型作为输入，并生成一个最优的张量卸载计划。

该计划的设计基于以下原则：

• 它应该卸载足够多的张量，使得 GPU 可以在目标批量大小下正确运行，而不会触发内存溢出错误。
• 它应该避免过多的从 GPU 到 SSD（反之亦然）的数据传输，从而最小化由张量卸载引起的迭代时间的增加。

卸载调度器的工作分为两个阶段：

我们先过一遍这个流程。

• initial state，溢出了16MB的内存
• 在阶段一的第一行，先线性搜索，把A B都offload了，这个时间是可以被计算重叠的，但是还是溢出8MB的内存
• 在阶段一的第二行，继续把CD offload掉，这时候已经不溢出了，但是通讯和计算不能完全重叠了，说明是卡在通信瓶颈了
• 故开始运行阶段二（不能简单的线性搜索了）
• 从后到前面去搜索，去掉压缩比不行的块，这里是D，然后压缩下一个；发现还是不行，继续向前搜索，去掉压缩比不行的块，这里是B；发现行了，计算和通讯完全重叠，算法结束。（这里的核心是选择压缩比比较好的张量去进行压缩然后offload，而不是简单的线性的搜索）
  

阶段 1：线性张量选择：
调度器首先迭代选择一定数量的张量，直到未选择的张量的总大小加上其他内存驻留对象（例如，权重和临时工作空间）的总大小小于 GPU 内存的总大小。然后，**调度器检查总数据传输时间（计算方法是将每个张量的卸载时间相加）是否小于前向传播中所有层的总执行时间。**如果满足这个条件，调度器就会采用这个计划并停止，因为它可以通过调度在层计算和张量卸载之间完全重叠。如果不满足，调度器进入第二阶段。
（这个时间检查是个灵魂，如果可以完全重叠，那就不用管了就没有第二步了）

阶段 2：压缩感知张量选择：
只有在第一阶段找不到满意的计划时，调度器才会运行第二阶段。这表明当前的计划在卸载张量上花费了太多时间，传输时间已经成为新的瓶颈。为了解决这个问题，调度器会用预期具有高 CSR 和 FP16 转换压缩比的压缩友好张量替换已选择的张量。

具体来说，调度器会以迭代的方式执行以下步骤来优化现有的计划：

1. 首先，调度器排除第一阶段选择的最后一个不可压缩的张量。它用一个或多个在尚未选择的张量中预期压缩比最高的张量替换它，使得新选择的张量的大小超过排除的张量的大小。
2. 然后，它重新计算预期的总数据传输时间，假设压缩张量的卸载时间是原始卸载时间的一部分（与压缩比例成反比）。

在这个阶段，调度器首先排除（或者说不再考虑）在第一阶段选择的最后一个不可压缩的张量。所谓 “不可压缩的张量”，指的是那些即使经过压缩，其数据大小也不会有显著减小的张量。这些张量在卸载过程中，由于数据大小没有显著减小，所以卸载时间长，不利于优化数据传输时间。

然后，调度器会在剩下的、还未被选择的张量中，找出预期压缩比最高的张量（或者如果一个还不够，就找多个）来替换刚刚被排除的那个张量。所谓 “预期压缩比最高的张量”，指的是那些经过压缩后，其数据大小会显著减小的张量。这些张量在卸载过程中，由于数据大小大幅减小，所以卸载时间短，有利于优化数据传输时间。

在选择新的张量替换旧的张量时，调度器会尽量保证新选择的张量的总大小超过被排除的那个张量的大小。这是为了保证在优化数据传输时间的同时，也能满足 GPU 的内存需求。

总的来说，这个阶段的目的是通过选择那些可以通过压缩显著减少数据大小的张量，来优化数据传输时间，从而提升整体的系统性能。

3. 如果这个总传输时间不超过前向传播的总执行时间，调度器就会停止。否则，它会重复这个过程，直到满足条件或者不存在压缩友好的张量（即，压缩后大小不减小的张量）。
4. 如果找到了一个满意的计划，对应的批量大小可能不会增加迭代时间，并达到比基线更高的吞吐量。另一方面，如果调度器因为找不到更多的压缩友好张量而停止，生成的计划可能会因为张量传输而产生一些延迟。然而，这个计划仍然可以用来在更大的批量大小下运行 DNN 训练（但可能吞吐量较低）。

这个过程还是比较清晰的。

Peer-to-Peer Direct Storage Access

P2P-DSA 允许 GPU 和 NVMe SSDs 之间进行直接内存访问，而无需使用主机 DRAM 缓冲区，从而在 SSD 读/写过程中最大限度地减少主机的介入。P2P-DSA 是一个轻量级的层，利用 GDRCopy 和 SPDK 这两项技术，实现张量从 GPU 到 NVMe SSDs 的通信。

GDRCopy 是一个基于 NVIDIA GPUDirect 的快速 GPU 内存复制库，它使 GPU 内存可以被其他 PCIe 设备直接访问。

Intel 的 SPDK 则将块级 I/O 接口直接暴露给用户空间软件。

P2P-DSA 有一个元数据表，用于维护卸载到 SSD 的每个张量的元数据。元数据表包含一个长的逻辑块地址（LBA）值和一个用于检查 I/O 完成的布尔值。

具体的传输请求（从 GPU 到 SSD 的卸载）操作如下：

当调用 P2PDSA_issue 时，P2P-DSA 从传输请求中获取索引、缓冲区和方向（写）信息。
调用逻辑块地址（LBA）分配器，为单个 SSD 设备或多个 SSD 设备（当使用多个 SSD 以提高卸载/预取带宽时）分配一组连续的块。将从 LBA 分配器分配的第一个块的 LBA 更新到元数据表的适当位置。
P2P-DSA 为每个逻辑块创建一个命令，然后将其排入命令队列。这里，NVMe 命令包含：

• i) 源地址（GPU 内存地址由 GPUDirect 转换为 PCIe 总线地址），和
• ii) 设备地址（使用元数据表中的 LBA 计算）。

当调用 P2PDSA_update 时：

从软件命令队列中获取排队的命令，并将其发给 NVMe SSD，只要 NVMe 设备的提交队列有空间。
NVMe SSD 设备将执行这些请求，并在 SSD 设备和 GPU 之间进行直接数据传输。一段时间后，这些传输请求将完成，并在 NVMe 设备完成队列中更新其状态。

当再次调用 P2PDSA_update 时，P2PDSA_update 将清空完成队列，并通过设置相应的完成位来更新元数据表。

此时，如果应用程序调用 P2PDSA_is_done 来查询已经卸载的张量，它将返回 true。

以下是这个过程的七个主要步骤：

1. 发起传输请求：当 P2PDSA_issue 被调用时，P2P-DSA （Peer-to-Peer Direct Storage Access）会从传输请求中获取索引、缓冲区和方向（写入）信息。

2. 分配逻辑块地址：调用逻辑块地址（LBA）分配器在单个或多个 SSD 设备上分配一组连续的块（当使用多个 SSD 时，可以提高卸载/预取的带宽）。然后，将从 LBA 分配器分配的第一个块的 LBA 更新到元数据表的相应位置。

3. 创建并推入命令：P2P-DSA 为每个逻辑块创建一个命令，然后将其加入命令队列。此处，一个 NVMe 命令包含：i) 源地址（GPU 内存地址由 GPUDirect 转换为 PCIe 总线地址）；ii) 设备地址（使用元数据表中的 LBA 计算）。

4. 发出命令：当调用 P2PDSA_update 时，只要 NVMe 设备提交队列有空间，就从软件命令队列中获取排队的命令并发送到 NVMe SSD。

5. 执行请求并传输数据：NVMe SSD 设备将执行这些请求，并在 SSD 设备和 GPU 之间直接进行数据传输。

6. 清理完成队列并更新元数据表：稍后，这些传输请求将完成，其状态将在 NVMe 设备完成队列中更新。当再次调用 P2PDSA_update 时，它将清理完成队列并通过设置相应的完成位来更新元数据表。

7. 检查已卸载的张量是否完成：如果此时应用程序调用 P2PDSA_is_done 来询问已卸载的张量是否完成，它将返回 true。

反向路径（即从 SSD 预取数据到 GPU 内存）的执行过程类似，只是：

• i) LBAs 是从元数据表中读取的，而不是被分配的；
• ii) 发出的是读命令，而不是写命令。

在卸载和预取的过程中，大部分数据访问都是顺序访问，这比随机访问在吞吐量和耐用性上更有优势。

两个关键函数

P2PDSA_issue 和 P2PDSA_update 是在论文中提到的两个函数，这两个函数是 Peer-to-Peer Direct Storage Access (P2P-DSA) 技术中的关键组成部分。这两个函数的作用是处理和管理 GPU 和 NVMe SSDs 之间的直接内存访问。

• P2PDSA_issue：这个函数的作用是处理从 GPU 到 SSD 的数据传输请求。当这个函数被调用时，它会从传输请求中获取索引、缓冲区和方向（写）信息。然后，它会调用逻辑块地址（LBA）分配器，为单个 SSD 设备或多个 SSD 设备分配一组连续的块。最后，P2PDSA_issue 会为每个逻辑块创建一个命令，然后将其排入命令队列。

• P2PDSA_update：这个函数的作用是管理 NVMe SSD 设备的完成队列。当这个函数被调用时，它会从软件命令队列中获取排队的命令，并将其发给 NVMe SSD 设备。然后，NVMe SSD 设备会执行这些请求，并在 SSD 设备和 GPU 之间进行直接数据传输。一段时间后，这些传输请求将完成，并在 NVMe 设备完成队列中更新其状态。在再次调用 P2PDSA_update 时，它会清空完成队列，并通过设置相应的完成位来更新元数据表。

通过这两个函数，P2P-DSA 技术能够有效地管理 GPU 和 NVMe SSDs 之间的直接内存访问，从而提高数据处理的效率。

小故事
这个地方偏系统的实现，我们通过一个小故事来看一下他在干什么。

你可以把 GPU 想象成一个大型仓库，里面储存了大量的包裹（也就是数据）。现在，你需要将这些包裹送到 SSD，也就是你的目的地。你可以把SSD 想象成一个远在城市另一头的大型收货点。P2PDSA_issue 就像是负责分配包裹和计划路线的调度员。

当你调用 P2PDSA_issue 时，调度员开始工作。他会先检查包裹的详细信息，然后为每个包裹分配一个独特的追踪号（也就是逻辑块地址，LBA）。然后，他会根据追踪号和包裹的目的地，为每个包裹创建一份详细的送货单（也就是命令），然后将这些送货单排入命令队列。

接着，你调用 P2PDSA_update，这就像是一辆大型货车，负责将包裹从仓库运送到收货点。每当货车有空位时，它就会从命令队列中取出一个送货单，然后从仓库中取出相应的包裹，直接送到收货点。在送货的过程中，货车会直接从仓库出发，直接到达收货点，不需要途径任何其他地方。

当包裹成功送达收货点后，货车会在完成队列中更新包裹的状态。然后，当你再次调用 P2PDSA_update 时，货车会清空完成队列，并通过设置相应的完成位来更新包裹的状态。

这就是整个过程的一个比喻。通过这个比喻，你可以看到，P2PDSA_issue 和 P2PDSA_update 的作用就像是一个高效的物流系统，它们能够有效地管理包裹（也就是数据）的传输过程，从而提高整个系统的效率。

---

这篇比较偏向系统，就写到这里~