DeepNVMe

本教程将展示如何使用 DeepNVMe 在持久存储和驻留在主机或设备内存中的张量之间进行数据传输。DeepNVMe 通过基于非易失性内存快闪存储器 (NVMe) 固态硬盘 (SSD)、Linux 异步 I/O (libaio) 和 NVIDIA Magnum IOTM GPUDirect® Storage (GDS) 的强大优化，提高了深度学习应用中 I/O 操作的性能和效率。

需求

确保您的环境已正确配置以使用 DeepNVMe。首先，您需要安装 DeepSpeed 版本 >= 0.15.0。接下来，确保 DeepNVMe 运算符在 DeepSpeed 安装中可用。async_io 运算符对于任何 DeepNVMe 功能都是必需的，而 gds 运算符仅在使用 GDS 功能时才需要。您可以通过检查 ds_report 的输出结果来确认每个运算符的可用性，以检查兼容状态是否为 [OKAY]。以下是 ds_report 输出的片段，确认 async_io 和 gds 运算符都可用。

deepnvme_ops_report

如果 async_io 运算符不可用，则需要为您的 Linux 版本安装相应的 libaio 库二进制文件。例如，Ubuntu 用户需要运行 apt install libaio-dev。一般来说，您应该仔细检查 ds_report 的输出结果，以获取有用的提示，例如以下提示

[WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
[WARNING]  async_io: please install the libaio-dev package with apt
[WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.

要启用 gds 运算符，您需要安装 NVIDIA GDS，并查阅 裸机系统 或 Azure VM（即将推出）的相应指南。

创建 DeepNVMe 句柄

可以通过两种抽象访问 DeepNVMe 功能：aio_handle 和 gds_handle。aio_handle 可用于主机和设备张量，而 gds_handle 仅适用于 CUDA 张量，但效率更高。使用 DeepNVMe 的第一步是创建一个所需的句柄。aio_handle 需要 async_io 运算符，而 gds_handle 需要 async_io 和 gds 运算符。以下代码片段分别说明了 aio_handle 和 gds_handle 的创建。

### Create aio_handle
from deepspeed.ops.op_builder import AsyncIOBuilder
aio_handle = AsyncIOBuilder().load().aio_handle()

### Create gds_handle
from deepspeed.ops.op_builder import GDSBuilder
gds_handle = GDSBuilder().load().gds_handle()

为简单起见，以上示例说明了使用默认参数创建句柄。我们期望使用默认参数创建的句柄在大多数环境中都能提供良好的性能。但是，您可以查看 下面 以了解高级句柄创建。

使用 DeepNVMe 句柄

aio_handle 和 gds_handle 提供了相同的 API，用于将张量存储到文件或从文件加载张量。这些 API 的一个共同特征是它们将张量和文件路径作为所需 I/O 操作的参数。为了获得最佳性能，应将固定的设备或主机张量用于 I/O 操作（有关详细信息，请参阅 此处）。为简洁起见，本教程将使用 aio_handle 进行说明，但请记住 gds_handle 的工作方式类似。

您可以通过在 Python shell 中对 aio_handle 对象进行制表符补全来查看可用的 API。这将使用 h. 的制表符补全进行说明。

>python
Python 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> h = AsyncIOBuilder().load().aio_handle()
>>> h.
h.async_pread(             h.free_cpu_locked_tensor(  h.get_overlap_events(      h.get_single_submit(       h.new_cpu_locked_tensor(   h.pwrite(                  h.sync_pread(              h.wait(
h.async_pwrite(            h.get_block_size(          h.get_queue_depth(         h.get_intra_op_parallelism(        h.pread(                   h.read(                    h.sync_pwrite(             h.write(

用于执行 I/O 操作的感兴趣的 API 是那些名称包含 pread 和 pwrite 子字符串的 API。为简洁起见，我们将重点介绍文件写入 API，即 sync_pwrite、async_pwrite 和 pwrite。我们将在下面仅讨论 sync_pwrite 和 async_pwrite，因为它们是 pwrite 的特化。

阻塞文件写入

sync_pwrite 提供了 Python 文件写入的标准阻塞语义。以下示例说明了如何使用 sync_pwrite 将 1GB CUDA 张量存储到本地 NVMe 文件中。

>>> import os
>>> os.path.isfile('/local_nvme/test_1GB.pt')
False
>>> import torch
>>> t=torch.empty(1024**3, dtype=torch.uint8).cuda()
>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> h = AsyncIOBuilder().load().aio_handle()
>>> h.sync_pwrite(t,'/local_nvme/test_1GB.pt')
>>> os.path.isfile('/local_nvme/test_1GB.pt')
True
>>> os.path.getsize('/local_nvme/test_1GB.pt')
1073741824

非阻塞文件写入

一项重要的 DeepNVMe 优化是非阻塞 I/O 语义，它使 Python 线程能够将计算与 I/O 操作重叠。async_pwrite 为文件写入提供了非阻塞语义。Python 线程稍后可以使用 wait() 与 I/O 操作同步。async_write 也可用于提交多个连续的非阻塞 I/O 操作，然后可以使用单个 wait() 对其进行阻塞。以下示例说明了如何使用 async_pwrite 将 1GB CUDA 张量存储到本地 NVMe 文件中。

>>> import os
>>> os.path.isfile('/local_nvme/test_1GB.pt')
False
>>> import torch
>>> t=torch.empty(1024**3, dtype=torch.uint8).cuda()
>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> h = AsyncIOBuilder().load().aio_handle()
>>> h.async_pwrite(t,'/local_nvme/test_1GB.pt')
>>> h.wait()
1
>>> os.path.isfile('/local_nvme/test_1GB.pt')
True
>>> os.path.getsize('/local_nvme/test_1GB.pt')
1073741824

非阻塞 I/O 操作的警告： 为了避免数据竞争和损坏，必须谨慎使用 .wait() 来序列化源张量的写入和目标张量的读取。例如，在非阻塞文件写入期间更新 t 是不安全的，可能会损坏 /local_nvme/test_1GB.pt。

>>> t=torch.empty(1024**3, dtype=torch.uint8).cuda()
>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> h = AsyncIOBuilder().load().aio_handle()
>>> h.async_pwrite(t,'/local_nvme/test_1GB.pt')
>>> t += 1 # <--- Data race; avoid by preceding with `h.wait()`

类似的安全问题也适用于在没有 .wait() 同步的情况下读取非阻塞文件读取的目标张量。

并行文件写入

一项重要的 DeepNVMe 优化是能够并行化各个 I/O 操作。此优化通过在构造 DeepNVMe 句柄时指定所需的并行度来启用。随后使用该句柄的 I/O 操作会自动在请求的主机或设备线程数上并行化，具体取决于情况。I/O 并行性可以与阻塞或非阻塞 I/O API 组合使用。以下示例说明了使用 async_pwrite 进行 4 路文件写入并行化。请注意，使用 intra_op_parallelism 参数在句柄创建中指定所需的并行度。

>>> import os
>>> os.path.isfile('/local_nvme/test_1GB.pt')
False
>>> import torch
>>> t=torch.empty(1024**3, dtype=torch.uint8).cuda()
>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> h = AsyncIOBuilder().load().aio_handle(intra_op_parallelism=4)
>>> h.async_pwrite(t,'/local_nvme/test_1GB.pt')
>>> h.wait()
1
>>> os.path.isfile('/local_nvme/test_1GB.pt')
True
>>> os.path.getsize('/local_nvme/test_1GB.pt')
1073741824

固定张量

DeepNVMe 优化的关键部分是使用直接内存访问 (DMA) 进行 I/O 操作，这需要主机或设备张量被固定。要固定主机张量，您可以使用 Pytorch 或 DeepSpeed 加速器 提供的机制。以下示例说明了将固定的 CPU 张量写入本地 NVMe 文件。

>>> import os
>>> os.path.isfile('/local_nvme/test_1GB.pt')
False
>>> import torch
>>> t=torch.empty(1024**3, dtype=torch.uint8).pin_memory()
>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> h = AsyncIOBuilder().load().aio_handle()
>>> h.async_pwrite(t,'/local_nvme/test_1GB.pt')
>>> h.wait()
1
>>> os.path.isfile('/local_nvme/test_1GB.pt')
True
>>> os.path.getsize('/local_nvme/test_1GB.pt')
1073741824

另一方面，gds_handle 提供了 new_pinned_device_tensor() 和 pin_device_tensor() 函数来固定 CUDA 张量。以下示例说明了将固定的 CUDA 张量写入本地 NVMe 文件。

>>> import os
>>> os.path.isfile('/local_nvme/test_1GB.pt')
False
>>> import torch
>>> t=torch.empty(1024**3, dtype=torch.uint8).cuda()
>>> from deepspeed.ops.op_builder import GDSBuilder
>>> h = GDSBuilder().load().gds_handle()
>>> h.pin_device_tensor(t)
>>> h.async_pwrite(t,'/local_nvme/test_1GB.pt')
>>> h.wait()
1
>>> os.path.isfile('/local_nvme/test_1GB.pt')
True
>>> os.path.getsize('/local_nvme/test_1GB.pt')
1073741824
>>> h.unpin_device_tensor(t)

整合在一起

我们希望以上内容能帮助您开始使用 DeepNVMe。您还可以使用以下链接查看 DeepNVMe 在现实世界深度学习应用中的用法。

  1. 参数交换器 在 ZeRO-推理 和 ZeRO-Infinity 中。
  2. 优化器交换器 在 ZeRO-Infinity 中。
  3. 梯度交换器 在 ZeRO-Infinity 中。
  4. 简单的文件读取和写入 操作。

致谢

本教程已通过 王冠华、田中雅浩 和 Stas Bekman 的反馈进行了重大改进。

附录

高级句柄创建

使用 DeepNVMe 实现峰值 I/O 性能需要仔细配置句柄创建。特别是，aio_handle 和 gds_handle 构造函数的参数对于性能至关重要，因为它们决定了 DeepNVMe 如何有效地与底层存储子系统（即 libaio、GDS、PCIe 和 SSD）交互。为方便起见，我们使您可以使用默认参数值创建句柄，这将在大多数情况下提供不错的性能。但是，在您的环境中榨取所有可用性能可能需要调整构造函数参数，即 block_size、queue_depth、single_submit、overlap_events 和 intra_op_parallelism。aio_handle 构造函数参数和默认值如下所示

>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> help(AsyncIOBuilder().load().aio_handle())
Help on aio_handle in module async_io object:

class aio_handle(pybind11_builtins.pybind11_object)
 |  Method resolution order:
 |      aio_handle
 |      pybind11_builtins.pybind11_object
 |      builtins.object
 |
 |  Methods defined here:
 |
 |  __init__(...)
 |      __init__(self: async_io.aio_handle, block_size: int = 1048576, queue_depth: int = 128, single_submit: bool = False, overlap_events: bool = False, intra_op_parallelism: int = 1) -> None
 |
 |      AIO handle constructor

性能调优

如前所述前面，为了实现目标工作负载或环境下的 DeepNVMe 峰值性能，需要使用配置优化的aio_handle或gds_handle句柄。为了方便配置，我们提供了一个名为ds_nvme_tune的实用程序来自动发现最佳的 DeepNVMe 配置。ds_nvme_tune会自动探索用户指定或默认的配置空间，并推荐提供最佳读写性能的选项。以下是如何使用ds_nvme_tune来调整/local_nvme上安装的本地 NVVMe SSD 与 GPU 内存之间aio_handle数据传输的示例。此示例使用ds_nvme_tune的默认配置空间进行调整。

$ ds_nvme_tune --nvme_dir /local_nvme --gpu
Running DeepNVMe performance tuning on ['/local_nvme/']
Best performance (GB/sec): read =  3.69, write =  3.18
{
   "aio": {
      "single_submit": "false",
      "overlap_events": "true",
      "intra_op_parallelism": 8,
      "queue_depth": 32,
      "block_size": 1048576
   }
}

上述调整是在配备两块 NVIDIA A6000-48GB GPU、252GB DRAM 和一块CS3040 NVMe 2TB SDD（峰值读写速度分别为 5.6 GB/s 和 4.3 GB/s）的 Lambda 工作站上执行的。调整过程大约需要四分半钟。根据结果，通过使用如下配置的aio_handle，可以预期分别实现 3.69 GB/秒和 3.18 GB/秒的读写传输速度。

>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> h = AsyncIOBuilder().load().aio_handle(block_size=1048576,
                                           queue_depth=32,
                                           single_submit=False,
                                           overlap_events=True,
                                           intra_op_parallelism=8)

ds_nvme_tune的完整命令行选项可以通过正常的-h或--help获得。

usage: ds_nvme_tune [-h] --nvme_dir NVME_DIR [NVME_DIR ...] [--sweep_config SWEEP_CONFIG] [--no_read] [--no_write] [--io_size IO_SIZE] [--gpu] [--gds] [--flush_page_cache] [--log_dir LOG_DIR] [--loops LOOPS] [--verbose]

options:
  -h, --help            show this help message and exit
  --nvme_dir NVME_DIR [NVME_DIR ...]
                        Directory in which to perform I/O tests. A writeable directory on a NVMe device.
  --sweep_config SWEEP_CONFIG
                        Performance sweep configuration json file.
  --no_read             Disable read performance measurements.
  --no_write            Disable write performance measurements.
  --io_size IO_SIZE     Number of I/O bytes to read/write for performance measurements.
  --gpu                 Test tensor transfers between GPU device and NVME device.
  --gds                 Run the sweep over NVIDIA GPUDirectStorage operator
  --flush_page_cache    Page cache will not be flushed and reported read speeds may be higher than actual ***Requires sudo access***.
  --log_dir LOG_DIR     Output directory for performance log files. Default is ./_aio_bench_logs
  --loops LOOPS         Count of operation repetitions
  --verbose             Print debugging information.

DeepNVMe API

为了方便起见，我们提供了 DeepNVMe API 的列表和简要说明。

通用 I/O API

以下函数用于aio_handle和gds_handle的 I/O 操作。

函数 描述
async_pread 非阻塞文件读取到张量
sync_pread 阻塞文件读取到张量
pread 具有阻塞和非阻塞选项的文件读取
async_pwrite 非阻塞文件写入从张量
sync_pwrite 阻塞文件写入从张量
pwrite 具有阻塞和非阻塞选项的文件写入
wait 等待非阻塞 I/O 操作完成

GDS 特定 API

以下函数仅适用于gds_handle

函数 描述
new_pinned_device_tensor 分配并固定设备张量
free_pinned_device_tensor 取消固定并释放设备张量
pin_device_tensor 固定设备张量
unpin_device_tensor 取消固定设备张量

句柄设置 API

以下 API 可用于探测句柄配置。

函数 描述
get_queue_depth 返回队列深度设置
get_single_submit 返回是否启用 single_submit
get_intra_op_parallelism 返回 I/O 并行度
get_block_size 返回 I/O 块大小设置
get_overlap_events 返回是否启用 overlap_event

更新：