
DeepSpeed DeepNVMe 技术深度解析用 NVMe 与 GPUDirect Storage 打通大模型 I/O 瓶颈【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文以 DeepSpeed 官方技术博客《DeepNVMe: Improving DL Applications through I/O Optimizations》blogs/deepnvme/08-2024/README.md为骨架系统讲解 DeepNVMe 如何通过 NVMe SSD、NVIDIA Magnum IO GPUDirect StorageGDS与 Linux 异步 I/Olibaio三大底层技术将数据加载、模型 checkpoint 与参数 offload 等 I/O 操作从大模型训练的瓶颈变为可扩展的性能引擎。读完本文你将理解 DeepNVMe 的设计原则、三类核心 I/O 优化、其开源实现与 API 入口以及它在 ZeRO-Inference 大模型生成场景中带来的实测收益。一、背景被忽视的存储子系统正在成为 DL 可扩展性的新瓶颈深度学习DL在语言、语音、视频和多模态等 AI 领域持续带来突破其关键推动力之一是在模型规模、序列长度与硬件并行度多个维度上的剧烈扩展。从系统视角看这种扩展把巨大压力施加在计算、内存、通信与存储等关键子系统上。然而现有的深度学习优化工作大多忽略了存储子系统——数据加载、模型 checkpoint 与 offload交换等 I/O 操作因此常常成为大规模 DL 的主要瓶颈。针对这一问题DeepSpeed 打造了一套统称为DeepNVMe的 I/O 优化方案通过加速 I/O 操作、降低硬件需求来提升 I/O 密集型 DL 应用的性能与效率。DeepNVMe 的加速建立在 NVMe SSD 与 NVIDIA GDS 等存储创新之上。在 Azure NC96ads_A100_v4 虚拟机上开展的实验中DeepNVMe 能够饱和可用的 NVMe 带宽用于与 GPU/CPU 内存的数据传输实测最高达到约 10 GB/s 读、5 GB/s 写。二、三项底层技术NVMe SSD、GDS 与 libaio对持久化存储的高性能访问是包括 DL 在内的众多计算领域面临的共同挑战业界也提出了大量软硬件方案。DeepNVMe 构建在以下三项技术之上NVMe SSD基于 Flash 的存储设备正取代慢得多的机械硬盘HDD成为现代服务器的主力持久存储。例如一台 Azure NC96ads_A100_v4 虚拟机配备 4 块 NVMe SSD单块读写带宽约 3.25 GB/s以 RAID-0 组合后理论上限可达约 13 GB/s 的聚合读带宽。DeepNVMe 通过并行化单个 I/O 操作、批量提交请求把多块 SSD 的聚合带宽真正吃满。NVIDIA Magnum IO GPUDirect StorageGDS支持在 NVMe 与 GPU 显存之间直接传输数据从而绕开传统方案中经由 CPU 内存中转bounce buffer带来的额外拷贝开销。GDS 在 CUDA 11.4 及以上版本中可用。Linux 异步 I/OlibaioLinux 引入的异步 I/O 栈相比传统 I/O 栈能更好地榨取 NVMe 这类快速存储设备的原始性能使 I/O 请求不必阻塞调用线程。三、DeepNVMe 模块的设计与实现四个优化 × POSIX 风格接口从仓库源码可以看到DeepNVMe 在 DeepSpeed 中并非单一模块而是一套算子 Python 工具的组合C/CUDA 算子与 Python 绑定位于 csrc/aiolibaio 后端与 csrc/gdsGDS 后端构建器 op_builder/async_io.py 与 op_builder/gds.py 负责编译链接async_io链接-laioGDS 额外依赖 CUDA 的libcufile并通过检查cuFileDriverOpen判定兼容性参见 op_builder/gds.pyI/O 基准与压测工具集中在 deepspeed/nvme如 ds_aio_handle.py、ds_io相关脚本面向应用的集成入口ZeRO-Inference / ZeRO-Infinity 的参数交换器位于 deepspeed/runtime/swap_tensor。该模块的设计遵循两个关键原则原则一用底层存储技术实现强有力的优化。具体包括非阻塞 I/Onon-blocking I/O提交后线程可立即返回后续用wait()同步批量提交bulk submission of I/O operations单次调用即可把大量请求交给内核单次 I/O 操作的并行化parallelization of an individual I/O operation将一个大 tensor 的读写切分到多个线程并行执行以扩展 NVMe 数量带来的带宽轻量级运行时lightweight runtime避免复杂框架带来的额外开销。原则二对外暴露简单、类 POSIX 的接口。上层 Python 代码只需要pread/pwrite这类直白 API即可获得上述底层优化而无需理解 NVMe、libaio 或 GDS 的复杂细节。这一设计使得它很容易嵌入既有 DL 应用。四、评测环境与基准工具实验在Azure NC96ads_A100_v4虚拟机上完成具体软硬件环境参见下表多设备实验将 SSD 以 RAID-0 组合。评估使用三套基准工具工具语言定位fioC广泛使用的通用 I/O 基准工具gdsioNVIDIA 提供专门评测 GDS 性能ds_ioPythonDeepSpeed 自研与 DeepNVMe 无缝集成更能代表以 Python 为主的 DL 应用ds_io及配套的aio_handle/gds_handle引擎、性能扫描工具等实现均保留在仓库中例如 deepspeed/nvme/ds_aio_handle.py 演示了如何按参数分别用GDSBuilder().load().gds_handle(...)或AsyncIOBuilder().load().aio_handle(...)创建引擎并支持以 bounce buffer 或 GPU pinned tensor 作为数据缓冲区见 deepspeed/nvme/ds_aio_handle.py#L85-L109。对应的单元测试位于 tests/unit/v1/nvme/test_aio.py。五、CPU 缓冲区近线性扩展至 10 GB/s 读、5 GB/s 写第一组微基准使用fio配置 libaio 后端与ds_io测量1GB 数据在 NVMe 与 CPU 内存之间的传输性能结果如图 1。从结果可以得出两点结论高性能尽管ds_io比fio更贴近真实 DL 应用Python 绑定、tensor 语义其吞吐仍大致与fio持平近线性扩展DeepNVMe 的 I/O 性能几乎随可用 NVMe 带宽线性增长达到10 GB/s 读、5 GB/s 写。这一带宽可扩展特性的根源正是第三节提到的单次 I/O 并行化 批量提交机制NVMe 数量越多可被并行切分与批量下发的请求越多聚合带宽利用越充分。六、GPU 缓冲区GDS 带来最高 37% 加速且匹配乃至超越 gdsio第二组微基准使用gdsio与ds_io测量1GB 数据在 NVMe 与 GPU 显存之间的传输。实验中ds_io分别配置了传统 bounce buffer 路径与更高效的 GDS 路径结果如图 2。观察结果有三点GDS 显著优于 bounce bufferDeepNVMe 启用 GDS 相比传统 bounce buffer 路径最高有37% 的加速——这正是 GDS 省去 CPU 中转拷贝的直接收益高性能即便更贴近 DL 应用DeepNVMe 也能匹配有时甚至超越gdsio带宽可扩展无论是否启用 GDSDeepNVMe 的 I/O 性能都随可用 NVMe 带宽增长。启用 GDS 时最高9.6 GB/s 读、5 GB/s 写不使用 GDS 时为7 GB/s 读、4 GB/s 写。七、端到端验证ZeRO-Inference 上的 LLAMA3-70B token 生成DeepNVMe 在真实 DL 应用中的代表性场景是ZeRO-Inference——一项通过 DeepNVMe 将模型权重 offload 到 CPU 或 NVMe 内存、从而大幅降低大模型推理硬件门槛的 AI 普惠技术。它非常适合离线推理等吞吐型负载以及硬件预算有限的场景。在仓库中ZeRO 的 NVMe offload 通路清晰可见零冗余配置模型在 deepspeed/runtime/zero/offload_config.py 中通过OffloadDeviceEnum.nvme与nvme_path指定deepspeed/runtime/zero/offload_config.py#L21-L30并强制要求devicenvme时nvme_path非空实际的参数/优化器/梯度交换器parameter/optimizer/gradient swapper位于 deepspeed/runtime/swap_tensor其中 partitioned_param_swapper.py 正是 ZeRO-Inference 权重 offload 的核心层间 NVMe 交换配置项aio的默认值定义在 deepspeed/runtime/swap_tensor/constants.py。为量化 NVMe offload 的性能博客以 token 生成负载做了端到端评测在单张 NVIDIA A100-80GB上推理LLaMA3-70Bprompt 长度 512、生成长度 32、batch size 96并将 NVMe SSD 数量从 1 扩展到 4。ZeRO-Inference 有/无 GDS 的结果如图 3。两点关键观察GDS 持续领先相比 bounce buffer 方案GDS 使 token 生成快 10%–18%随 NVMe 带宽扩展4 块 NVMe SSD 时DeepNVMe 达到GDS 路径约 7 tokens/s、bounce buffer 路径约 6 tokens/s的生成吞吐。Profiling 结果还表明DeepNVMe 会随更多 NVMe 带宽继续扩展因此它是提升生成式应用性能的一条经济路径。八、从博客到仓库DeepNVMe 的 API 与上手方式博客正文之外DeepSpeed 仓库的 docs/_tutorials/deepnvme.md 给出了完整的上手指南可用于验证上文所述能力在开源代码中的落点算子依赖所有 DeepNVMe 功能都依赖async_io算子仅 GDS 功能额外需要gds算子可通过ds_report输出检查两者的兼容状态如async_io不可用时通常需安装libaio-devUbuntu 下执行apt install libaio-dev。两个句柄抽象aio_handle适用于主机与设备 tensorgds_handle只支持 CUDA tensor 但效率更高。创建方式分别如下### aio_handle需要 async_io 算子 from deepspeed.ops.op_builder import AsyncIOBuilder aio_handle AsyncIOBuilder().load().aio_handle() ### gds_handle需要 async_io gds 算子 from deepspeed.ops.op_builder import GDSBuilder gds_handle GDSBuilder().load().gds_handle()I/O API 家族两个句柄提供一致的接口核心包括async_pread/sync_pread/pread读与async_pwrite/sync_pwrite/pwrite写以及非阻塞同步用的wait()。其中sync_*提供标准阻塞语义async_*提供非阻塞语义提交后线程可继续计算稍后以一次wait()阻塞等待一批已提交请求完成。构造参数即性能开关aio_handle/gds_handle构造函数中的block_size、queue_depth、single_submit、overlap_events、intra_op_parallelism是性能关键参数Python 侧默认值依次为1048576、128、False、False、1。这五个参数与 ZeRO 配置 JSON 中aio段相对应——后者默认定义在 deepspeed/runtime/swap_tensor/constants.py解析逻辑见 deepspeed/runtime/swap_tensor/aio_config.py。可结合 deepspeed/runtime/swap_tensor/async_swapper.py 等实现继续深入阅读。自动化调优工具仓库提供了ds_nvme_tune命令行工具可自动搜索并推荐最优句柄配置。例如在挂载于/local_nvme的 NVMe 上进行 GPU↔NVMe 传输调优ds_nvme_tune --nvme_dir /local_nvme --gpu其完整选项--sweep_config、--no_read、--io_size、--gds、--flush_page_cache、--loops等可通过ds_nvme_tune -h查看。九、总结与展望DeepNVMe 是 DeepSpeed 为应对 I/O 操作成为深度学习可扩展性关键瓶颈而打造的 I/O 优化技术它基于 NVMe SSD 与 NVIDIA GDS 等存储技术通过非阻塞 I/O、批量提交、单次 I/O 并行化与轻量运行时等优化实现持久化存储与 DL 应用内存之间的高速、高效数据传输。在 Azure NC96ads_A100_v4 上以 NVMe offload 推理 LLaMA3-70B单张 A100-80GB的实测中DeepNVMe 达到最高约7 tokens/s的生成吞吐。博客宣布 DeepNVMe 随DeepSpeed v0.15.0 及以上版本开源可用后续博客还将报告其在模型 checkpoint、数据加载等其他 I/O 密集型 DL 应用上的改进。本文所引实验数据与图表均出自 blogs/deepnvme/08-2024/README.md仓库中的算子构建器op_builder/async_io.py、op_builder/gds.py、句柄实现与 ZeRO 交换器代码可作为进一步验证与复现的参考。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考