NCCL 2.23版本中的主机内存泄漏问题分析与修复

在NCCL 2.23版本中，用户报告了一个潜在的主机内存泄漏问题。这个问题在使用相同的通信器进行多次集合操作时会逐渐累积内存消耗，最终可能导致内存耗尽(OOM)的情况。

问题现象

当程序在单线程中无限循环执行ncclReduce操作时，可以观察到主机内存使用量持续增长。通过启用NCCL的调试日志(NCCL_DEBUG=INFO NCCL_DEBUG_SUBSYS=ALLOC)，可以看到类似如下的内存分配记录：

memory stack hunk malloc(26345472)
memory stack hunk malloc(26411008)

这些分配的大小会随着时间的推移不断增加，最终导致系统内存耗尽。

问题根源

经过深入分析，发现问题出在NCCL内部的内存管理机制上。具体来说，在enqueue.cc文件中，当执行集合操作时，会使用comm->memPermanent来分配任务结构(ncclTaskColl)。这种分配方式会导致内存不会被及时释放，从而造成内存泄漏。

解决方案

NCCL开发团队已经确认了这个问题，并提供了一个修复补丁。修复的核心思想是将内存分配从永久内存池(comm->memPermanent)改为使用作用域内存池(comm->memScoped)。具体修改如下：

// 原代码
struct ncclTaskColl* t = ncclMemoryPoolAlloc<struct ncclTaskColl>(&comm->memPool_ncclTaskColl, &comm->memPermanent);

// 修复后代码
struct ncclTaskColl* t = ncclMemoryPoolAlloc<struct ncclTaskColl>(&comm->memPool_ncclTaskColl, &comm->memScoped);

这种修改确保了在集合操作完成后，相关的内存资源能够被正确释放，从而避免了内存泄漏。

影响范围

这个问题主要影响以下情况：

1. 长时间运行的NCCL应用程序
2. 频繁使用相同通信器执行集合操作的程序
3. 使用NCCL 2.23版本的系统

临时解决方案

对于需要使用NCCL 2.23版本的用户，可以应用开发团队提供的补丁来临时修复这个问题。该补丁已经过验证，能够有效解决内存泄漏问题。

长期解决方案

NCCL开发团队表示这个问题已经在后续版本中修复，建议用户升级到包含该修复的NCCL新版本。

总结

内存管理是高性能计算库中的关键问题。NCCL团队对这类问题的快速响应和修复体现了其对软件质量的重视。对于开发者来说，理解这类问题的根源有助于更好地使用NCCL库，并在遇到类似问题时能够快速诊断和解决。