#8531·DeepSpeed

[请求] 为 NPU 将本机固定主机内存注册到设备运行时(跟进 #8283 和 #8315)

作者: VenusTZZ创建于 2026年9月16日更新于 2026年9月16日
标签enhancement

**请描述您的功能请求是否与问题有关。**DeepSpeed 的原生固定内存后端(`DS_PIN_MEMORY_BACKEND=native`)使用 `posix_memalign` + `mlock` 锁定页面,自 #8283 起,它还可以将每个缓冲区注册到设备运行时中,以便异步复制可以使用 DMA 引擎。这种更改在发布时带来了很大的改变:在 H200 上,H2D 带宽从 ~10-17 GB/s(仅使用 mlock)提升到 ~53-55 GB/s 后注册。然而,注册功能是一个针对每个加速器的钩子,而今天只有 CUDA 加速器实现了它。其他所有平台都会进入一个空的默认值:在 Ascend NPU 上,设置为 `DS_PIN_MEMORY_REGISTER_DEVICE=1` 时,该钩子会默默地返回 `False`,缓冲区仍然只能使用 mlock,而原生后端的性能与可分页内存相同。这不是一个假设的差距 - 在 Ascend 910B4(CANN 8.5.0)上,注册相同的原生缓冲区将 4 MiB H2D 从 ~10 GB/s 提升到 ~23 GB/s,而 64 MiB 的注册副本保持 ~23.4 GB/s,而 mlock-only 则在 ~13 和 ~22.5 GB/s 之间波动。XPU 平台面临着相同的问题,已提交为 #8315。NPU 和 XPU 用户独立地触及了相同的盲点,这表明 #8283 的好处是根据平台进行限制的,而非 CUDA 平台运行越来越多的实际训练作业(尤其是 Huawei Ascend),在没有实现两个钩子的情况下,它们会默默地在每个 H2D 复制中丢失 2x+。此问题涵盖了 NPU 方面;#8315 则涵盖 XPU。

内容来源: deepspeedai/DeepSpeed