DeepNVMe

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

需求

确保您的环境已正确配置以使用DeepNVMe。首先，您需要安装DeepSpeed版本>= 0.15.0。接下来，确保DeepSpeed安装中提供了DeepNVMe操作符。任何DeepNVMe功能都需要async_io操作符，而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虚拟机（即将推出）的相应指南。

创建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的工作方式类似。

你可以通过在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 Accelerators提供的机制。以下示例展示了如何将固定的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在实际深度学习应用中的使用情况。

Parameter swapper 在 ZeRO-Inference 和 ZeRO-Infinity 中。
Optimizer swapper 在 ZeRO-Infinity 中。
Gradient swapper 在 ZeRO-Infinity 中。
简单的文件读写操作。

致谢

本教程通过Guanhua Wang、Masahiro Tanaka和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调整aio_handle在GPU内存和挂载在/local_nvme上的本地NVVMe SSD之间数据传输的示例。此示例使用了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的Lambda工作站上执行的，其峰值读写速度分别为5.6 GB/s和4.3 GB/s。调优大约需要四分钟半。根据结果，可以预期通过使用如下配置的aio_handle，分别实现3.69 GB/sec和3.18 GB/sec的读写传输速度。

>>> 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的列表和简要描述。

通用输入/输出 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