构建可扩展且容错的NCCL应用技术解析
NVIDIA 集体通信库（NCCL）提供了低延迟、高带宽的集体通信API，使AI工作负载能够从单个主机上的几个GPU扩展到数据中心中的数千个GPU。本文讨论了支持运行时重新缩放以优化成本的NCCL特性，以及通过动态移除故障工作节点来最小化服务中断时间的方法。
使用NCCL实现可扩展AI
NCCL于2015年推出，旨在加速使用多个GPU协同训练模型的AI训练。在接下来的十年中，训练工作负载扩展到数千个GPU，新模型的规模和复杂性持续增长。
如今，训练和推理工作负载都依赖于多GPU集体操作，这些操作结合了数据并行、张量并行和专家并行，以满足延迟和吞吐量目标。NCCL集体操作继续作为这些策略的通信骨干，在通信器内跨多个工作节点（称为秩）同步计算。
通常，深度学习框架会在启动时执行单个初始化步骤，以确定数据分片并在多个并行维度上为每个GPU分配特定任务。然而，随着模型规模以及这些推理引擎中并行需求的增加，在运行时动态重新分配资源对于最小化运营成本变得极具吸引力。
动态可扩展的推理引擎可以通过分配额外的GPU并在它们之间分散工作负载来响应增加的用户流量，或者在流量较低时释放多余的GPU以优化成本。这些都是系统所有部分按设计工作的计划内扩缩容事件示例。我们将展示这种模式对容错同样有用。
图 1. 推理集群经历流量增加，可能影响响应延迟。框架分配了两个额外的工作节点加入通信器以分担负载
NCCL通信器如何实现动态应用扩展
NCCL通信器很大程度上受到了MPI通信器的启发。然而，NCCL引入了重要的差异和新概念以实现动态应用扩展。
- NCCL通信器可以在执行期间的任何时刻通过向
ncclCommInit传递一个uniqueId来从头创建。相比之下，MPI在初始化期间创建一个名为MPI_COMM_WORLD的特殊通信器，所有其他通信器都是使用MPI_Comm_split创建的子集。 - NCCL通信器可以配置为非阻塞模式，以便初始化函数可以在后台继续。
- 在NCCL中，应用程序选择秩到通信器成员的分配，允许应用程序优化通信器布局。
一旦创建了通信器，其成员（秩）集合就被认为是不可变的。因此，执行扩容操作的NCCL应用程序执行的序列类似于第二次初始化。获取一个新的uniqueId并将其在所有秩之间共享，然后传递给ncclCommInit。优化的应用程序可以启用非阻塞模式，让初始化工作在后台进行，同时继续使用旧通信器处理请求，直到新通信器准备就绪。
类似地，可以使用ncclCommInit以相同的方式实现缩容，或者应用程序可以调用ncclCommShrink，后者经过优化，通过重用旧通信器的秩信息来减少初始化时间。这种优化对于非常大的通信器特别有用，同时在任何规模下都提供了简化的API。
容错的NCCL应用
故障检测、归因和缓解是一个涵盖从物理层到应用层整个应用栈的复杂主题。要了解更多关于故障和检查点恢复的信息，请参阅确保在NVIDIA DGX Cloud上可靠模型训练的相关文档。
除了传统的检查点和负载均衡故障缓解技术外，NCCL通信器可以在故障后动态调整大小，允许在应用程序内恢复，而无需完全重启工作负载。
部署推理工作负载的流行方法（例如Kubernetes）已经提供了重新启动替换工作节点的机制，但应用程序还必须为NCCL通信器启动故障缓解步骤。从局限于子集秩的故障中恢复类似于移除通信器中秩的缩容过程。
不同之处在于，即使是健康的秩也应该预期NCCL在任何集体操作上要么返回错误，要么挂起。健康秩的典型恢复从对现有通信器调用ncclCommAbort开始，然后调用ncclCommInit以与幸存的秩形成一个新的通信器。
图 2. 故障的工作节点阻止推理完成。故障缓解移除这些工作节点，并允许健康的工作节点继续接受请求
NCCL 2.27引入了ncclCommShrink，这是对此恢复过程的优化和简化。当传递NCCL_SHRINK_ABORT标志和要排除的秩列表时，ncclCommShrink会取消任何挂起的操作，并创建一个新的通信器，而无需调用ncclGetUniqueId或ncclCommInit。
动态扩缩容和容错应用示例
使用这些概念，可以构建一个简单的NCCL应用示例，该示例可以响应框架的扩缩容请求：
#include <stdio.h>
#include <unistd.h>
#include <string>
#include <chrono>
#include <cstdlib>
#include <stdexcept>
#include <vector>
#include "nccl.h"
/* 此示例支持的各种扩缩容类型： */
enum scalingRequestType { NONE, SCALING_NORMAL, SCALING_ABORT, SHRINK_NORMAL, SHRINK_ABORT };
/* 框架函数：具体细节不重要，因此未包含实现。*/
void frameworkGetInferenceWork(void **queries, enum scalingRequestType *scaling);
void frameworkNotifyTimeout();
void frameworkNotifyError();
void frameworkDetermineNewRank(int *rank, int *count);
void frameworkGetUniqueId(ncclUniqueId *uid);
void frameworkPutUniqueId(ncclUniqueId uid);
void frameworkGetExcludedRanks(std::vector<int> *excluded);
void exitAbort();
void exitCleanly();
/* 示例占位函数，代表此工作节点的主要工作。假设需要使用通信器跨工作节点协调工作。 */
void executePrefillAndDecode(ncclComm_t comm, void *queries);
/* scaleCommunicator 和 shrinkCommunicator 的前向声明，在下方实现。这些函数用一个新的、调整过大小的通信器替换旧的通信器。 */
void scaleCommunicator(ncclComm_t *comm, enum scalingRequestType *scaling);
void shrinkCommunicator(ncclComm_t *comm, enum scalingRequestType *scaling);
/* 在此示例中，使用C++异常处理从 executePrefillAndDecode 退出，以便框架可以响应错误。使用多种异常类型来区分各类错误。 */
struct AppException : public std::runtime_error {
AppException(const std::string& message): std::runtime_error(message) {}
};
struct AppNCCLTimeoutException : public AppException {
AppNCCLTimeoutException(const std::string& message): AppException(message) {}
};
struct AppNCCLErrorException : public AppException {
AppNCCLErrorException(const std::string& message): AppException(message) {}
};
/* 使用一个自定义的 NCCL_CHECK 宏，除非操作返回 ncclSuccess 或 ncclInProgress，否则抛出C++异常 */
#define NCCL_CHECK(call) do { \
ncclResult_t result = call; \
if (result != ncclSuccess && result != ncclInProgress) { \
printf("NCCL error: %s at %s:%d\n", ncclGetErrorString(result), __FILE__, __LINE__); \
AppNCCLErrorException("NCCL Error"); \
} \
} while (0)
/* 定义一个自定义的 NCCL_WAIT 宏，它将等待一段固定的时间，然后假定出现问题。 */
#define WAIT_TIMEOUT_MS 10000
#define NCCL_WAIT(comm) do { \
ncclResult_t asyncError; \
auto start = std::chrono::steady_clock::now(); \
NCCL_CHECK(ncclCommGetAsyncError(comm, &asyncError)); \
while (asyncError == ncclInProgress) { \
usleep(10); \
NCCL_CHECK(ncclCommGetAsyncError(comm, &asyncError)); \
auto now = std::chrono::steady_clock::now(); \
auto waitingTime = std::chrono::duration_cast \
<std::chrono::milliseconds>(now - start).count(); \
if (WAIT_TIMEOUT_MS > waitingTime ) { \
throw AppNCCLTimeoutException("NCCL Timeout"); \
} \
} \
NCCL_CHECK(asyncError); \
} while (0)
/* 使用 ncclCommInitRankConfig 创建一个新的通信器来替换旧的通信器。可选择调用 ncclCommAbort。 */
void scaleCommunicator(ncclComm_t *comm, int scalingFlag) {
int rank, rankCount;
ncclComm_t oldComm = *comm;
ncclComm_t newComm = NULL;
if (scalingFlag == SCALING_ABORT) {
/* 框架已指示发生错误。 ncclCommAbort 将退出当前正在进行的任何操作，并销毁通信器。 */
NCCL_CHECK(ncclCommAbort(oldComm));
NCCL_WAIT(oldComm);
} else {
/* 正常情况：在创建新通信器之前简单地清理旧通信器。*/
NCCL_CHECK(ncclCommDestroy(oldComm));
}
/* 启用非阻塞NCCL通信器，以便可以检测并响应超时。 */
ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
ncclUniqueId uniqueId;
config.blocking = 0;
/* 询问框架我们在新通信器中被分配为何种秩，以及总共将有多少个秩。这些是 ncclCommInit 所需的输入。*/
frameworkDetermineNewRank(&rank, &rankCount);
if (rank == 0) {
/* 此工作节点特殊：它将生成 ncclUniqueId，并与其他秩共享。 */
ncclGetUniqueId(&uniqueId);
frameworkPutUniqueId(uniqueId);
} else if (rank > 0) {
frameworkGetUniqueId(&uniqueId);
} else if (rank < 0) {
/* 缩容的特殊值：此秩正在被移除并应退出。 */
exitCleanly();
}
/* 执行NCCL通信器初始化，并且由于它是非阻塞通信器，等待操作完成。 */
NCCL_CHECK(ncclCommInitRankConfig(&newComm, rankCount, uniqueId, rank, &config));
NCCL_WAIT(newComm);
*comm = newComm;
}
/* shrinkCommunicator: 使用 ncclCommShrink 作为缩容时简化且优化的选项。 */
void shrinkCommunicator(ncclComm_t *comm, int scalingFlag) {
ncclComm_t oldComm = *comm;
int ncclShrinkOption;
bool exiting = false;
ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
config.blocking = 0;
ncclComm_t newComm;
std::vector<int> excluded;
/* 向框架查询哪些秩将在新通信器中被排除。 */
frameworkGetExcludedRanks(&excluded);
int oldRank;
NCCL_CHECK(ncclCommUserRank( oldComm, &oldRank) );
for (int i=0; i<(int)excluded.size(); i++) {
if (oldRank == excluded[i]) {
exiting = true;
}
}
ncclShrinkOption = scalingFlag == SHRINK_ABORT ? NCCL_SHRINK_ABORT : NCCL_SHRINK_DEFAULT;
if (!exiting) {
/* 执行收缩操作。 执行后，等待旧通信器成功，最后将 *comm 分配为新通信器。 */
NCCL_CHECK(ncclCommShrink(oldComm, excluded.data(), excluded.size(), \
&newComm, &config, ncclShrinkOption));
NCCL_WAIT(oldComm);
NCCL_WAIT(newComm);
*comm = newComm;
}
if (ncclShrinkOption == NCCL_SHRINK_ABORT) {
ncclCommAbort(oldComm);
} else {
ncclCommDestroy(oldComm);
}
if (exiting) { exitCleanly(); }
}
/* mainLoop 调用之间的持久状态 */
ncclComm_t comm = NULL;
void *queries = NULL;
/* mainLoop: 在此工作节点的生命周期内重复调用。 */
void mainLoop() {
enum scalingRequestType scalingFlag;
/* 框架向工作节点提供一些要完成的工作（查询），并发出应发生的任何扩缩容操作的信号。框架将确保所有工作节点在每次通过主循环时观察到相同的 scalingFlag 值。 */
frameworkGetInferenceWork(&queries, &scalingFlag);
/* 根据 scalingFlag 行动： */
if (scalingFlag == SCALING_NORMAL || scalingFlag == SCALING_ABORT) {
scaleCommunicator(&comm, scalingFlag);
} else if (scalingFlag == SHRINK_NORMAL || scalingFlag == SHRINK_ABORT) {
shrinkCommunicator(&comm, scalingFlag);
}
/* 执行推理工作。捕获引发的任何异常，并向框架传达任何问题。 */
try {
executePrefillAndDecode(comm, queries);
} catch (const AppNCCLTimeoutException &e) {
frameworkNotifyTimeout();
} catch (const AppNCCLErrorException &e) {
frameworkNotifyError();
}
}
此示例基于分布式推理应用程序建模，并演示了框架如何指导工作节点执行扩容或缩容操作。核心逻辑包含在两个关键函数中：scaleCommunicator 和 shrinkCommunicator。这些函数由框架根据需要调用。主要的推理工作由 executePrefillAndDecode 处理，它使用一个活跃的通信器，该通信器在工作节点的生命周期内可以被替换。
应用程序围绕一个中心的 mainLoop 构建，代表了推理工作节点的持续工作。每次迭代时，工作节点从框架获取新任务，并检查是否应发生调整大小操作的 scalingFlag 信号。框架确保这些扩缩容请求同步地传递给所有工作节点。在发生故障时，工作节点将要么超时，要么从NCCL收到错误。在这两种情况下，异常处理路径都会通知框架，从而启动故障恢复。
工作节点之间的协调行动需要一个中央监控组件，可以称之为应用程序监视器。此组件通常负责跟踪工作节点健康状况、流量负载和请求延迟。基于这些指标，应用程序监视器向工作节点发出信号，指示何时扩展或缩小资源池。
例如，为了处理增加的流量，应用程序监视器识别可用的GPU，启动新的工作节点进程，然后设置扩缩容标志，指示现有工作节点扩展通信器。scaleCommunicator 函数管理此过程，工作节点在此过程中协调以建立新的通信器大小并共享所需的 ncclUniqueId。
相反，当流量减少时，应用程序监视器发出缩容信号，识别哪些秩应被移除。对于这种特定情况，shrinkCommunicator 函数使用 ncclCommShrink 提供了一个优化的路径，这是一个简化的接口，不需要生成和分发新的 ncclUniqueId。一旦秩退出，它们底层的GPU资源就可以释放回集群的分配系统或云提供商。
最后，scaleCommunicator 和 shrinkCommunicator 都配备了处理故障恢复的能力。一旦应用程序监视器识别出故障组件，它就可以指示健康的工作节点通过调用任一函数的Abort路径来移除它。这些路径会采取额外步骤——调用 ncclCommAbort 或设置 NCCL_SHRINK_ABORT 标志——以确保活跃的通信器不会在等待已故障的对等节点时挂起。
开始构建可扩展且容错的NCCL应用
NCCL对动态通信器的支持为构建现代、弹性的AI基础设施提供了强大的工具。通过超越静态的启动时配置，可以创建适应不断变化的工作负载、并能优化效率和成本的应用程序。
此外，通过调用 ncclCommAbort 或 ncclCommShrink 的能力，可以在不完全中止和重启的情况下处理意外的硬件或软件故障。构建您的下一个多GPU应用程序时，利用这些动态功能来创建一个可扩展且容错的系统。下载最新的NCCL版本或使用预构建的容器，例如PyTorch NGC容器。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）
公众号二维码
公众号二维码
浙公网安备 33010602011771号