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 与计算重叠。

我们的目标是使用单进程微基准测试（可在此处获取）来展示 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 个 SSD 扩展到 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() 在 CPU 版本的 torch 中无法工作，如下所示。

>>> 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 性能扩展和易用性等多方面的改进。

本博客介绍了 Microsoft DeepSpeed 团队的 Joe Mayer、Logan Adams 和 Olatunji Ruwase 所做的工作。