DeepNVMe：为深度学习应用提供经济实惠的 I/O 扩展

我们使用 DeepNVMe 开发了 FastPersist 和 ZeRO-Inference，分别针对 DL 训练和推理中的 I/O 瓶颈。我们的实验在单个虚拟机上进行，其中我们将可用的 NVMe SSD 组合成一个 RAID-0（即磁盘条带化）卷，以利用聚合的读写带宽。由于 DeepNVMe 可以使用 CPU 弹出缓冲区（又称 AIO）或 NVIDIA GPUDirect Storage（又称 GDS）来卸载张量，我们报告了这两种模式的结果。

虽然将模型检查点保存到持久存储在模型训练中至关重要，但由于现有方法的低效率，它也是一个主要瓶颈。我们开发了 FastPersist 来应对检查点设置的性能挑战。FastPersist 通过三种关键技术使检查点开销在训练期间可以忽略不计：(i) DeepNVMe，(ii) 数据并行，以及 (iii) I/O 和计算重叠。

我们在此的目标是使用单进程微基准测试（可在此处获取 here）来展示 DeepNVMe 在 FastPersist 中的影响，该测试将模型检查点状态从 HBM 序列化到本地 NVMe。我们在实验中使用流行的 PyTorch torch.save() 作为基线，并将 FastPersist 集成到 torch.save() 中，以简化采用和性能比较。

我们测量了将 Phi-3-Mini 检查点状态从 HBM 序列化到本地 NVMe 存储的吞吐量。结果总结在下图中。我们观察到，与基线相比，使用 FastPersist 的检查点速度显著加快。在 8xGen5 NVMe 设置中，我们看到了超过 20 倍的速度提升。我们还观察到，与 4xGen5 相比，FastPersist 随着 8xGen5 NVMe 带宽的增加而扩展。

我们使用我们的 ds_io 基准测试工具来证明 DeepNVMe 能够与可用的 NVMe 带宽成比例地扩展 I/O 性能。这使用户能够通过使用更多或更快的 NVMe SSD，以适度的成本加速受 I/O 限制的 DL 应用。在我们的实验中，我们测量了 HBM 和 NVMe 之间 1GB 数据传输所达到的读写带宽。我们评估了将 NVMe 从 PCIe Gen4 扩展到 Gen5，以及从 4 个扩展到 8 个 SSD。这些 SSD 被组合成一个单一的 RAID-0（磁盘条带化）卷。我们在下图中总结了结果，显示 DeepNVMe 在两个维度上都扩展了 I/O 性能。从 4xGen4 SSD 扩展到 4xGen5 SSD，读取速度从 10GB/秒提高到 27GB/秒，写入速度从 5GB/秒提高到 11GB/秒。从 4xGen5 扩展到 8xGen5，读取速度进一步提高到 48GB/秒，写入速度提高到 26GB/秒。

我们通过取消对硬件环境和 I/O 操作的限制，增加了 DeepNVMe 的使用场景，具体如下所述。

尽管 GPU（及类似加速器）在 DL 中占主导地位，但 CPU 仍在重要的机器学习 (ML) 工作负载中使用，例如推荐系统。然而，DeepNVMe 以前在纯 CPU 环境中无法使用。这是因为 DeepNVMe 依赖 torch.pin_memory() 来处理页锁定的 CPU 张量，而 torch.pin_memory() 在 torch 的 CPU 版本中不起作用，如下所示。

>>> import torch
>>> torch.__version__
'2.6.0+cpu'
>>> x = torch.empty(1024).pin_memory()
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: Cannot access accelerator device when none is available.
>>>

我们通过添加用于分配 (new_cpu_locked_tensor()) 和释放 (free_cpu_locked_tensor()) 页锁定 CPU 张量的机制，使 DeepNVMe 在 CPU 环境中可用。下面的代码片段演示了分配一个固定的 CPU 张量 (x)。

>> import torch
>>> torch.__version__
'2.6.0+cpu'
>>> from deepspeed.ops.op_builder import AsyncIOBuilder
>>> h = AsyncIOBuilder().load().aio_handle()
>>> x = h.new_cpu_locked_tensor(1024, torch.Tensor())
>>> x.shape
torch.Size([1024])
>>> x.dtype
torch.float32

以前，DeepNVMe 的功能仅限于读取或写入文件的全部内容。我们现在改进了 DeepNVMe，使其能够从给定的偏移量读取或写入用户指定的文件内容部分。具体来说，我们扩展了现有的读/写 API，以接受用户指定的 file offset 参数（默认值为 0），如下所示。

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

pread(...) method of async_io.aio_handle instance
pread(self: async_io.aio_handle, buffer: torch.Tensor, filename: str, validate: bool, async: bool, file_offset: int = 0) -> int

在开发 FastPersist 时，我们需要以字节格式操作模型张量（通常是浮点数据类型），这既是为了性能，也是为了 I/O 操作的便利。然而，我们没有找到一个零拷贝的机制来将张量从任意数据类型转换为字节数据类型（即 torch.uint8），所以我们决定创建一个。此功能可通过 UtilsBuilder 操作获得，如下例所示。在该示例中，我们将一个 torch.bfloat16 张量转换为 torch.uint8。请注意，由于该功能的零拷贝特性，bf16_tensor 和 byte_tensor 是别名。

>>> import torch
>>> from deepspeed.ops.op_builder import UtilsBuilder
>>> util_ops = UtilsBuilder().load()
>>> bf16_tensor = torch.zeros(1024, dtype=torch.bfloat16, device='cuda')
>>> bf16_tensor
tensor([0., 0., 0., ..., 0., 0., 0.], device='cuda:0', dtype=torch.bfloat16)
>>> byte_tensor = util_ops.cast_to_byte_tensor(bf16_tensor)
>>> byte_tensor
tensor([0, 0, 0, ..., 0, 0, 0], device='cuda:0', dtype=torch.uint8)
>>> bf16_tensor += 1.0
>>> bf16_tensor
tensor([1., 1., 1., ..., 1., 1., 1.], device='cuda:0', dtype=torch.bfloat16)
>>> byte_tensor
tensor([128, 63, 128, ..., 63, 128, 63], device='cuda:0',
dtype=torch.uint8)

本博文更新了我们对 DeepNVMe 的持续开发情况，这是一种用于加速 DL 应用的 I/O 优化技术。我们宣布了 DeepNVMe 在多个方面的改进，包括应用覆盖范围、I/O 性能扩展和可用性。

本博文描述了微软 DeepSpeed 团队的 Joe Mayer、Logan Adams 和 Olatunji Ruwase 所做的工作。