1. 背景概述

随着深度学习模型规模的不断扩大，单设备训练已难以满足实际需求，分布式训练成为提升训练效率、支撑大规模模型的关键技术。在华为昇腾 NPU 平台上，torch_npu 分布式训练模块为开发者提供了完整的分布式训练支持，使训练任务能够在多个 NPU 设备及多台服务器之间高效协同运行。

该模块深度集成 PyTorch 分布式框架，在保持标准 torch.distributed API 兼容性的同时，充分发挥昇腾 NPU 的硬件特性，为分布式训练提供高性能、高可靠的基础设施支撑。

本文在剖析该模块架构设计与核心实现的基础上，结合 FSDP、子进程组、RPC 及完整训练脚本等实践示例，帮助开发者系统掌握昇腾 NPU 上的分布式训练方法。

2. 模块概述

2.1 主要功能

• HCCL 后端：基于华为集合通信库（HCCL）实现的高性能分布式通信后端
• LCCL 后端：轻量级集合通信库，用于本地设备间通信
• 进程组管理：支持创建和管理多个进程组，实现灵活的通信拓扑
• 集合通信操作：支持 all_reduce、broadcast、all_gather、reduce_scatter、send、recv 等操作
• FSDP 支持：为全切片数据并行（Fully Sharded Data Parallel）提供 NPU 适配
• RPC 支持：基于 TensorPipe 的远程过程调用框架
• DDP Reducer：分布式数据并行的梯度同步机制

2.2 核心组件职责

组件	职责
ProcessGroupHCCL	HCCL 后端进程组实现，处理所有集合通信操作
ProcessGroupLCCL	LCCL 后端进程组实现，用于本地通信
Reducer	DDP 梯度聚合和同步
HCCLComm	HCCL 通信器的 RAII 封装
ParallelTcpStore	高性能分布式键值存储
TensorPipeAgent	RPC 通信代理

3. 架构设计

3.1 整体架构

torch_npu 分布式模块采用分层架构设计，自顶向下依次为 Python API 层、后端注册层、C++ 后端实现层、通信工具层以及 CANN/HCCL/LCCL 原生库层，各层职责清晰、解耦合理。

+---------------------------------------------------------------------------------+
|                              Python API Layer                                    |
|  +-------------------+  +-------------------+  +-------------------+             |
|  | torch.distributed |  | torch_npu.distrib |  | FSDP / RPC APIs  |             |
|  +-------------------+  +-------------------+  +-------------------+             |
+---------------------------------------------------------------------------------+
                                     |
                                     v
+---------------------------------------------------------------------------------+
|                           Backend Registration Layer                             |
|  +------------------------+  +-------------------------+  +------------------+  |
|  | _new_process_group_    |  | Backend._plugins        |  | rendezvous       |  |
|  | helper (patched)       |  | ["HCCL"] registration   |  | handlers         |  |
|  +------------------------+  +-------------------------+  +------------------+  |
+---------------------------------------------------------------------------------+
                                     |
                                     v
+---------------------------------------------------------------------------------+
|                           C++ Backend Implementation                             |
|  +---------------------+     +---------------------+     +---------------------+ |
|  | ProcessGroupHCCL    |     | ProcessGroupLCCL    |     | TensorPipeAgent    | |
|  | - WorkHCCL          |     | - WorkLCCL          |     | (RPC)              | |
|  | - Watchdog          |     +---------------------+     +---------------------+ |
|  | - Options           |                                                        |
|  +---------------------+                                                        |
+---------------------------------------------------------------------------------+
                                     |
                                     v
+---------------------------------------------------------------------------------+
|                           Communication Utilities                                |
|  +------------------+  +------------------+  +------------------+                |
|  | HCCLUtils        |  | LCCLUtils        |  | TraceUtils       |                |
|  | - HCCLComm       |  | - LcclComm       |  | - HCCLTraceBuffer|                |
|  | - Data Type Map  |  | - Data Type Map  |  | - Debug Info     |                |
|  +------------------+  +------------------+  +------------------+                |
+---------------------------------------------------------------------------------+
                                     |
                                     v
+---------------------------------------------------------------------------------+
|                           CANN/HCCL/LCCL Native Libraries                        |
|  +------------------+  +------------------+  +------------------+                |
|  | HCCL API         |  | LCCL API         |  | ACL Runtime      |                |
|  | (hccl.h)         |  |                  |  |                  |                |
|  +------------------+  +------------------+  +------------------+                |
+---------------------------------------------------------------------------------+

3.2 HCCL 后端设计

HCCL（Huawei Collective Communication Library）后端是 torch_npu 分布式训练的核心组件，负责管理通信器生命周期、调度集合通信操作并监控执行状态。

+------------------------------------------------------------------+
|                      ProcessGroupHCCL                             |
+------------------------------------------------------------------+
|  Members:                                                         |
|  - store_: c10::intrusive_ptr      // 分布式存储           |
|  - devHCCLCommMap_: map // 通信器缓存           |
|  - hcclStreams_: map   // HCCL专用流           |
|  - workMetaList_: list          // 待处理工作列表       |
|  - watchdog_: unique_ptr        // 看门狗线程           |
+------------------------------------------------------------------+
|  Core Methods:                                                    |
|  - allreduce()      - broadcast()      - reduce()                 |
|  - allgather()      - reduce_scatter() - barrier()                |
|  - send()           - recv()           - alltoall()               |
+------------------------------------------------------------------+
|  Internal Components:                                             |
|  +------------------+  +------------------+  +------------------+ |
|  | WorkHCCL         |  | Watchdog         |  | Options          | |
|  | - hcclComms_     |  | - hcclCommWatch  |  | - timeout        | |
|  | - hcclEvents_    |  | - runLoop()      |  | - high_priority  | |
|  | - wait()         |  | - heartbeat      |  | - hccl_config    | |
|  +------------------+  +------------------+  +------------------+ |
+------------------------------------------------------------------+

3.3 LCCL 后端设计

LCCL（Local Collective Communication Library）后端专注于同一节点内的设备间通信，提供轻量级的本地通信能力。

+------------------------------------------------------------------+
|                      ProcessGroupLCCL                             |
+------------------------------------------------------------------+
|  Members:                                                         |
|  - store_: intrusive_ptr                                   |
|  - devLCCLCommMap_: map                         |
|  - lcclStreams_: map                           |
+------------------------------------------------------------------+
|  Supported Operations:                                            |
|  - allreduce()      - allgather()                                 |
|  - broadcast()      - reduce_scatter()                            |
+------------------------------------------------------------------+

3.4 进程组管理

进程组是分布式通信的基本组织单元。torch_npu 支持全局进程组和多个子进程组的灵活管理，满足不同并行策略的通信需求。

+------------------------------------------------------------------+
|                    Process Group Management                       |
+------------------------------------------------------------------+
|                                                                   |
|   Global Process Group (default)                                  |
|   +-----------------------------------------------------------+  |
|   | ProcessGroupHCCL (rank 0-N)                                |  |
|   | - All ranks participate                                    |  |
|   | - global_ranks_in_group = []                               |  |
|   +-----------------------------------------------------------+  |
|                              |                                    |
|              +---------------+---------------+                    |
|              v                               v                    |
|   Sub Process Group 1              Sub Process Group 2            |
|   +------------------------+      +------------------------+      |
|   | ProcessGroupHCCL       |      | ProcessGroupHCCL       |      |
|   | ranks [0,1,2,3]        |      | ranks [4,5,6,7]        |      |
|   | global_ranks = [0-3]   |      | global_ranks = [4-7]   |      |
|   +------------------------+      +------------------------+      |
|                                                                   |
+------------------------------------------------------------------+

3.5 与 PyTorch 分布式框架的集成

torch_npu 通过插件机制无缝集成到 PyTorch 分布式框架中，用户无需修改现有代码即可使用标准的 torch.distributed API。

+------------------------------------------------------------------+
|                 PyTorch Integration Architecture                  |
+------------------------------------------------------------------+
|                                                                   |
|   torch.distributed.init_process_group(backend="hccl")            |
|                              |                                    |
|                              v                                    |
|   +----------------------------------------------------------+   |
|   | _new_process_group_helper (Patched in distributed_c10d.py)|   |
|   +----------------------------------------------------------+   |
|                              |                                    |
|   Check backend_str == "HCCL"?                                    |
|                              |                                    |
|              +---------------+-----------------+                   |
|              |                                 |                   |
|              v                                 v                   |
|   Backend._plugins["HCCL"]          其他后端                          |
|              |                                                    |
|              v                                                    |
|   creator_fn -> ProcessGroupHCCL                                  |
|              |                                                    |
|              v                                                    |
|   pg._register_backend(device, backend_type, backend_class)       |
|                                                                   |
+------------------------------------------------------------------+

4. 核心流程

4.1 后端注册流程

后端注册发生在 import torch_npu 时，通过一系列初始化步骤将 HCCL 后端注册到 PyTorch 分布式框架中。

启动时序:
+--------------------------------------------------------------------------+
|  1. import torch_npu                                                      |
|         |                                                                 |
|         v                                                                 |
|  2. torch_npu._C._c10d_npu_init()                                         |
|         |                                                                 |
|         v                                                                 |
|  3. 注册 HCCL 后端到 Backend._plugins                                     |
|         |                                                                 |
|         v                                                                 |
|  4. 注册 rendezvous handlers ("parallel", "env://")                       |
|         |                                                                 |
|         v                                                                 |
|  5. 应用 FSDP 补丁 (_apply_fsdp_patch)                                    |
|         |                                                                 |
|         v                                                                 |
|  6. Monkey-patch _new_process_group_helper                                |
+--------------------------------------------------------------------------+

4.2 进程组创建流程

进程组的创建遵循延迟初始化策略，通信器在首次执行集合操作时才真正创建，以节省系统资源。

+--------------------------------------------------------------------------+
|  dist.init_process_group(backend="hccl", ...)                            |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  _new_process_group_helper()                                              |
|  - 检查 backend_str == "HCCL"                                             |
|  - 从 Backend._plugins["HCCL"] 获取 creator_fn                            |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  creator_fn(dist_backend_opts, backend_options)                           |
|  - dist_backend_opts 包含 store, rank, size, timeout, group_id            |
|  - backend_options 包含 global_ranks_in_group, hccl_config                |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  ProcessGroupHCCL 构造函数                                                |
|  1. 初始化成员变量 (store_, options_, rank_, size_)                       |
|  2. 配置超时参数 (kOpWaitTimeout, HCCL_EXEC_TIMEOUT)                      |
|  3. 创建 Watchdog 实例                                                    |
|  4. 启动 Watchdog 线程 (watchdog_->start())                               |
|  5. 如果是全局组且启用 HCCL_ZERO_COPY, 创建零拷贝通信器                   |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  HCCL Communicator 创建 (延迟初始化,首次集合操作时)                       |
|  1. broadcastMasterID() - 广播 HcclRootInfo                               |
|  2. HCCLComm::create_config() - 调用 HcclCommInitRootInfo                 |
|  3. 缓存到 devHCCLCommMap_                                                |
+--------------------------------------------------------------------------+

4.3 集合通信流程（以 all_reduce 为例）

集合通信操作遵循统一的执行流程：参数校验 → 通信器获取 → 流同步 → 异步执行 → 工作项入队。

+--------------------------------------------------------------------------+
|  dist.all_reduce(tensor)                                                  |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  ProcessGroupHCCL::allreduce(tensors, opts)                               |
|  1. 检查张量有效性 (check_npu_tensors_different_devices)                  |
|  2. 获取设备列表和 devicesKey                                             |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  collective() 模板函数                                                    |
|  1. getHCCLComm() - 获取或创建 HCCL 通信器                                |
|  2. syncStreams() - 同步当前流和 HCCL 流                                  |
|  3. initWork() - 创建 WorkHCCL 对象                                       |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  执行 HCCL 操作                                                           |
|  1. 记录开始事件 (hcclStartEvents_)                                       |
|  2. 调用 HcclAllReduce()                                                  |
|  3. 记录结束事件 (hcclEndEvents_)                                         |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  workEnqueue(work)                                                        |
|  1. 将 work 添加到 workMetaList_                                          |
|  2. 更新 pgStatus_ 统计信息                                               |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|  返回 Work 对象给用户                                                     |
|  用户可以调用 work.wait() 等待完成                                        |
+--------------------------------------------------------------------------+

4.4 FSDP 工作流程

FSDP（Fully Sharded Data Parallel）通过参数分片和按需全收集的方式，在保证训练效果的同时显著降低显存占用。

+--------------------------------------------------------------------------+
|                        FSDP Forward Pass                                  |
+--------------------------------------------------------------------------+
|  1. FSDPParamGroup 调用 all_gather 收集完整参数                           |
|     - _patched_all_gather_copy_in() 优化数据拷贝                          |
|     - _patched_get_param_all_gather_inputs() foreach_copy 优化            |
+--------------------------------------------------------------------------+
         |
         v
+--------------------------------------------------------------------------+
|                        FSDP Backward Pass                                 |
+--------------------------------------------------------------------------+
|  1. 计算梯度                                                              |
|  2. _patched_foreach_reduce() 执行梯度规约                                |
|     - 使用 reduce_scatter 进行梯度聚合                                    |
|     - 支持 HSDP 的 all_reduce                                             |
|  3. _patched_finalize_backward() 完成后向传播                             |
|     - 等待 all_gather 完成                                                |
|     - 同步 grad_offload_event                                             |
+--------------------------------------------------------------------------+

5. 核心类设计

5.1 类关系总览

分布式训练模块的核心类包括进程组实现（ProcessGroupHCCL/LCCL）、通信工作项（WorkHCCL）、梯度规约器（Reducer）、RPC 代理（TensorPipeAgent）以及调试追踪缓冲区（HCCLTraceBuffer）。其中 ProcessGroupHCCL 继承自 Backend 基类，与 WorkHCCL 为组合关系，并持有 Watchdog 看门狗线程与 Options 配置对象；HCCLComm 作为通信器的 RAII 封装被进程组持有，Reducer 依赖 ProcessGroupHCCL 完成梯度同步，HCCLTraceBuffer 作为调试工具被进程组引用。各类之间的关系如下：

5.2 关键类职责

HCCLComm：HCCL 通信器的 RAII 封装类，负责通信器的创建、配置、销毁以及异步错误查询。支持全局通信器和子通信器的创建。

ProcessGroupLCCL：轻量级集合通信库后端，用于本地设备间通信，提供与 HCCL 类似的接口。

WorkLCCL：LCCL 操作的工作项，管理设备列表、通信器列表和结束事件。

Reducer：DDP 梯度聚合的核心组件，管理参数分桶、梯度累积和通信钩子注册。

TensorPipeAgent：RPC 通信代理，负责跨节点的消息发送、设备映射管理和组成员更新。

HCCLTraceBuffer：Flight Recorder 调试缓冲区，记录所有通信操作的元数据，支持超时后的调试信息转储。

5.3 关键设计要点

• WorkHCCL 封装了异步操作的完整生命周期，通过 NPUEvent 实现高效的流同步和状态查询
• Watchdog 独立线程监控通信操作的健康状态，支持超时检测和去同步调试
• HCCLComm 采用 RAII 模式管理底层通信器资源，确保异常安全
• Options 提供灵活的配置入口，支持超时、流优先级、通信配置等参数定制

6. 关键实现细节

6.1 ProcessGroupHCCL 核心数据结构

class ProcessGroupHCCL : public c10d::Backend {
protected:
    // 分布式存储，用于跨进程同步
    c10::intrusive_ptr store_;

    // HCCL 通信器缓存（设备key -> 通信器列表）
    std::unordered_map>> devHCCLCommMap_;

    // HCCL 专用流（设备key -> 流列表）
    std::unordered_map> hcclStreams_;

    // 待处理的 Work 列表（由 Watchdog 监控）
    std::list workMetaList_;

    // 看门狗实例
    std::unique_ptr watchdog_;

    // 进程组状态统计
    std::shared_ptr pgStatus_;
};

6.2 collective 模板函数

所有集合通信操作（allreduce、broadcast、allgather 等）均通过统一的模板函数实现，保证了操作流程的一致性：

template 
c10::intrusive_ptr ProcessGroupHCCL::collective(
    std::vector& input,
    std::vector& output,
    Fn fn,                          // HCCL 操作函数
    PreProcess pre,                 // 前处理回调
    PostProcess post,               // 后处理回调
    c10d::OpType opType,
    bool asyncOp) {

    // 1. 获取设备列表和缓存 key
    const auto devices = getDeviceList(input);
    auto key = getKeyFromDevices(devices);

    // 2. 获取或创建 HCCL 通信器
    auto& hcclComms = getHCCLComm(key, devices);

    // 3. 同步流，确保数据就绪
    syncStreams(devices, hcclEvents_[key], hcclStreams_[key]);

    // 4. 创建 Work 对象
    auto work = initWork(devices, rank_, opType, ...);

    // 5. 执行前处理（如输入拷贝）
    pre(hcclStreams_[key], work);

    // 6. 执行 HCCL 操作
    for (size_t i = 0; i < input.size(); ++i) {
        fn(input[i], output[i], hcclComms[i]->getHcclComm(), hcclStreams_[key][i]);
    }

    // 7. 执行后处理
    post(hcclStreams_[key], work);

    // 8. 记录事件并入队监控
    work->hcclEndEvents_->record(hcclStreams_[key][i]);
    workEnqueue(work);

    return work;
}

6.3 Watchdog 工作循环

Watchdog 线程持续监控待处理的通信操作，执行以下逻辑：

1. 等待通知或超时唤醒
2. 检查 HCCL 通信器错误状态
3. 遍历工作队列，判断各操作是否完成
4. 对已完成操作更新状态并出队
5. 对超时操作触发调试信息转储

void ProcessGroupHCCL::Watchdog::runLoop() {
    while (!pg_->terminateProcessGroup_.load()) {
        std::unique_lock lock(pg_->workMetaListMutex_);

        // 等待通知或超时
        pg_->workMetaListCV_.wait_for(lock,
            std::chrono::milliseconds(kWatchdogThreadSleepMillis), ...);

        // 检查 HCCL 通信器错误
        pg_->checkHcclComms();

        // 遍历待处理工作
        for (auto it = pg_->workMetaList_.begin();
             it != pg_->workMetaList_.end(); ) {
            auto& work = *it;

            // 检查是否完成
            if (work.finishedNPUExecution()) {
                pg_->pgStatus_->lastCompletedSeq = work.seq_;
                it = pg_->workMetaList_.erase(it);
            } else if (work.checkTimeout()) {
                // 超时处理：转储调试信息
                pg_->shouldDump_.store(true);
                pg_->dumpDebuggingInfo();
                throw std::runtime_error("HCCL timeout");
            }
        }
    }
}

6.4 后端注册机制

Python 层注册

通过 monkey-patch _new_process_group_helper 函数，将 HCCL 后端无缝集成到 PyTorch 分布式框架中。用户无需修改代码即可使用标准的 torch.distributed API。

C++ 层 PyBind 绑定

auto processGroupHCCL = intrusive_ptr_no_gil_destructor_class_(
    module, "ProcessGroupHCCL", dist.attr("Backend"))
    .def(py::init&,
            int,    // rank
            int,    // size
            c10::intrusive_ptr>())
    .def("get_hccl_comm", &ProcessGroupHCCL::getHcclComm)
    .def("_set_hccl_comm_name", &ProcessGroupHCCL::setHcclCommName)
    .def("batch_isend_irecv", ...)
    // ... 其他方法绑定 ...

6.5 FSDP 补丁优化

FSDP 补丁针对 NPU 特性进行了以下优化：

1. 非阻塞拷贝优化：使用 non_blocking=True 进行异步拷贝，提高数据传输效率
2. 事件同步优化：使用 NPU 特定的事件同步机制，确保计算与通信的正确衔接
3. foreach_reduce 优化：处理 PREMUL_SUM 操作符不支持的情况，自动转换为 SUM 操作

7. 关键文件路径

组件	文件路径
ProcessGroupHCCL 头文件	torch_npu/csrc/distributed/ProcessGroupHCCL.hpp
ProcessGroupHCCL 实现	torch_npu/csrc/distributed/ProcessGroupHCCL.cpp
ProcessGroupLCCL 头文件	torch_npu/csrc/distributed/ProcessGroupLCCL.hpp
ProcessGroupLCCL 实现	torch_npu/csrc/distributed/ProcessGroupLCCL.cpp
HCCLUtils	torch_npu/csrc/distributed/HCCLUtils.hpp
TraceUtils	torch_npu/csrc/distributed/TraceUtils.h
Reducer	torch_npu/csrc/distributed/reducer.hpp
Python 分布式接口	torch_npu/distributed/distributed_c10d.py
FSDP 补丁	torch_npu/distributed/fsdp/_add_fsdp_patch.py
RPC 后端注册	torch_npu/distributed/rpc/backend_registry.py
C++ Init (PyBind)	torch_npu/csrc/distributed/Init.cpp

8. 架构特性分析

8.1 设计优势

与 PyTorch 深度集成

• 通过 monkey-patch 无缝集成到 PyTorch 分布式框架
• 用户可使用标准 torch.distributed API，无需修改代码
• 完整支持 DDP、FSDP、RPC 等分布式特性

完善的错误处理和监控

• Watchdog 线程：持续监控通信操作，检测超时和错误
• 心跳监控：独立线程检测 Watchdog 卡死状态
• Flight Recorder：记录所有通信操作，便于问题定位
• Desync Debug：支持检测和分析进程间同步问题

灵活的通信器管理

• 延迟初始化：通信器在首次使用时创建，节省资源
• 缓存机制：缓存已创建的通信器，避免重复创建开销
• 子组支持：支持创建子进程组，实现灵活通信拓扑

性能优化

• 独立 HCCL 流：使用专用流执行通信操作，与计算流并行
• 事件同步：使用 NPUEvent 实现高效流同步
• 零拷贝支持：通过环境变量配置减少内存拷贝
• FSDP 优化：针对 NPU 的异步拷贝优化

8.2 设计约束与应对

单设备限制

HCCL 按每个进程一个 NPU 设备的模型进行适配。对于需要多设备管理的场景，可通过多进程方式实现。

部分操作兼容性

• gather 操作基于 all_gather 组合实现，在特定场景下可根据实际需求选择合适的通信原语
• recvAnysource 场景可通过其他通信原语组合实现等价功能

框架版本兼容性

monkey-patch 方式与 PyTorch 内部实现存在耦合，建议在 PyTorch 版本升级时做好兼容性测试。

9. 优化建议

9.1 架构演进方向

抽象通信后端接口

建议创建统一的 NPU 通信后端接口，将 HCCL 和 LCCL 实现统一抽象，便于扩展和维护。

减少框架耦合

探索使用 PyTorch 官方的后端注册机制替代 monkey-patch，降低对内部实现的依赖。

9.2 性能优化方向

批量小张量通信优化

实现 coalesced 操作的优化路径，将多个小张量合并到单个缓冲区执行通信，减少通信次数：

c10::intrusive_ptr ProcessGroupHCCL::allreduce_coalesced(
    std::vector& tensors,
    const AllreduceCoalescedOptions& opts) {
    // 将小张量合并到单个缓冲区
    auto flat_tensor = flatten_tensors(tensors);
    // 执行单次 allreduce
    auto work = allreduce({flat_tensor}, opts);
    // 返回后展开
    return work;
}

利用 HCCL 高级特性

充分利用 HCCL 的异步错误检测能力，减少 Watchdog 的轮询开销：

void ProcessGroupHCCL::enableAsyncErrorDetection() {
    // 启用 HCCL 内部的异步错误检测
    // 减少 Watchdog 的轮询开销
}

9.3 可靠性增强

细粒度超时控制

实现分场景的超时配置，针对不同操作类型设置独立的超时阈值：

struct TimeoutConfig {
    std::chrono::milliseconds init_timeout;       // 初始化超时
    std::chrono::milliseconds collective_timeout; // 集合操作超时
    std::chrono::milliseconds p2p_timeout;        // 点对点超时
    std::chrono::milliseconds barrier_timeout;    // 屏障超时
};

增强 Flight Recorder

支持实时导出和分析，提供更丰富的性能统计信息：

class HCCLTraceBuffer {
public:
    // 添加实时导出功能
    void export_to_file(const std::string& path, bool include_stack_traces);

    // 添加性能统计
    struct Statistics {
        double avg_allreduce_latency_ms;
        double avg_broadcast_latency_ms;
        uint64_t total_bytes_transferred;
    };
    Statistics get_statistics() const;
};

9.4 易用性改进

改进错误消息

提供更详细的错误诊断信息，包括操作类型、设备信息、张量形状等上下文：

void ProcessGroupHCCL::throwDetailedError(
    const std::string& operation,
    HcclResult error_code,
    const WorkHCCL& work) {
    std::string msg = fmt::format(
        "HCCL {} failed on rank {} (group: {}, seq: {})\n"
        "Error code: {} ({})\n"
        "Input shapes: {}\n"
        "Possible causes: ...",
        operation, rank_, pg_name_, work.seq_,
        error_code, hcclErrorString(error_code),
        format_tensor_shapes(work.inputs_));
    throw std::runtime_error(msg);
}

添加健康检查 API

提供进程组健康状态查询接口，便于运维监控：

def check_process_group_health(pg) -> HealthStatus:
    """
    检查进程组的健康状态
    返回: HealthStatus 包含:
        - is_healthy: bool
        - pending_operations: int
        - last_error: Optional[str]
        - communication_latency_ms: float
    """
    pass

10. 实践示例

10.1 基础分布式训练

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
import torch_npu

def setup(rank, world_size):
    """初始化分布式环境"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'

    # 设置 NPU 设备
    torch.npu.set_device(rank)

    # 初始化进程组，使用 HCCL 后端
    dist.init_process_group(
        backend="hccl",
        rank=rank,
        world_size=world_size
    )

def cleanup():
    """清理分布式环境"""
    dist.destroy_process_group()

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)

def train(rank, world_size):
    setup(rank, world_size)

    # 创建模型并移动到 NPU
    model = SimpleModel().to(f'npu:{rank}')

    # 使用 DDP 包装模型
    ddp_model = DDP(model, device_ids=[rank])

    # 创建优化器
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)

    # 创建示例数据
    inputs = torch.randn(32, 10).to(f'npu:{rank}')
    labels = torch.randn(32, 10).to(f'npu:{rank}')

    # 训练循环
    for epoch in range(10):
        optimizer.zero_grad()
        outputs = ddp_model(inputs)
        loss = nn.functional.mse_loss(outputs, labels)
        loss.backward()
        optimizer.step()

        if rank == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

    cleanup()

if __name__ == "__main__":
    world_size = torch.npu.device_count()
    torch.multiprocessing.spawn(
        train,
        args=(world_size,),
        nprocs=world_size,
        join=True
    )

10.2 集合通信操作示例

import torch
import torch.distributed as dist
import torch_npu

def collective_operations_example():
    """演示各种集合通信操作"""
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = f'npu:{rank}'

    # All-Reduce：所有进程的张量求和
    tensor = torch.ones(4, 4).to(device) * rank
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    print(f"[Rank {rank}] After all_reduce: {tensor[0, 0].item()}")

    # Broadcast：从 rank 0 广播数据
    if rank == 0:
        tensor = torch.tensor([1, 2, 3, 4], dtype=torch.float32).to(device)
    else:
        tensor = torch.zeros(4, dtype=torch.float32).to(device)
    dist.broadcast(tensor, src=0)
    print(f"[Rank {rank}] After broadcast: {tensor}")

    # All-Gather：收集所有进程的数据
    tensor = torch.tensor([rank], dtype=torch.float32).to(device)
    tensor_list = [torch.zeros(1).to(device) for _ in range(world_size)]
    dist.all_gather(tensor_list, tensor)
    print(f"[Rank {rank}] After all_gather: {[t.item() for t in tensor_list]}")

    # Reduce-Scatter：规约后分散
    input_tensor = torch.ones(world_size, 4).to(device) * rank
    output_tensor = torch.zeros(4).to(device)
    dist.reduce_scatter_tensor(output_tensor, input_tensor, op=dist.ReduceOp.SUM)
    print(f"[Rank {rank}] After reduce_scatter: {output_tensor}")

    # Barrier：同步屏障
    dist.barrier()
    print(f"[Rank {rank}] Passed barrier")

10.3 FSDP 分布式训练实现

FSDP（Fully Sharded Data Parallel）是 PyTorch 提供的一种模型并行训练策略，通过将模型参数、梯度和优化器状态分片到多个设备上，显著降低单设备显存占用。以下是在昇腾 NPU 上实现 FSDP 训练的完整示例。

模型定义

首先定义基础的 Transformer 模型结构，包括单个 Transformer 块和由多个块堆叠而成的大型 Transformer 模型：

class TransformerBlock(nn.Module):
    def __init__(self, d_model=512, nhead=8):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model),
        )
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        x = self.norm1(x + self.self_attn(x, x, x)[0])
        x = self.norm2(x + self.ffn(x))
        return x

class LargeTransformer(nn.Module):
    def __init__(self, num_layers=12, d_model=512):
        super().__init__()
        self.embedding = nn.Embedding(10000, d_model)
        self.layers = nn.ModuleList([
            TransformerBlock(d_model) for _ in range(num_layers)
        ])
        self.head = nn.Linear(d_model, 10000)

    def forward(self, x):
        x = self.embedding(x)
        for layer in self.layers:
            x = layer(x)
        return self.head(x)

FSDP 训练流程

在昇腾 NPU 上使用 FSDP 进行训练，需要完成环境初始化、模型包装、混合精度配置和训练循环等步骤：

def train_fsdp(rank, world_size):
    # 初始化分布式环境
    dist.init_process_group(backend="hccl", rank=rank, world_size=world_size)
    torch.npu.set_device(rank)

    # 创建模型
    model = LargeTransformer(num_layers=24, d_model=1024)

    # 配置混合精度
    mixed_precision = MixedPrecision(
        param_dtype=torch.float16,
        reduce_dtype=torch.float16,
        buffer_dtype=torch.float16,
    )

    # 配置自动包装策略
    auto_wrap_policy = partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={TransformerBlock},
    )

    # 使用 FSDP 包装模型
    model = FSDP(
        model,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        mixed_precision=mixed_precision,
        auto_wrap_policy=auto_wrap_policy,
        device_id=torch.npu.current_device(),
    )

    # 创建优化器
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    # 训练循环
    for step in range(100):
        # 创建示例数据
        inputs = torch.randint(0, 10000, (8, 128)).to(f'npu:{rank}')
        labels = torch.randint(0, 10000, (8, 128)).to(f'npu:{rank}')

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = nn.functional.cross_entropy(
            outputs.view(-1, outputs.size(-1)),
            labels.view(-1)
        )
        loss.backward()
        optimizer.step()

        if rank == 0 and step % 10 == 0:
            print(f"Step {step}, Loss: {loss.item():.4f}")

    dist.destroy_process_group()

10.4 创建和使用子进程组

在复杂的分布式训练场景中，往往需要将设备划分为多个子组，实现灵活的通信拓扑。例如，在流水线并行或 2D 并行中，需要按行或按列组织通信组。

基础子进程组

以下示例演示如何根据 rank 将设备划分为两个子组，并在子组内执行集合通信：

def subgroup_example():
    """演示如何创建和使用子进程组"""
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = f'npu:{rank}'

    # 创建子进程组
    # 假设 world_size = 8, 创建两个子组: [0,1,2,3] 和 [4,5,6,7]
    if rank < world_size // 2:
        group_ranks = list(range(world_size // 2))
    else:
        group_ranks = list(range(world_size // 2, world_size))

    subgroup = dist.new_group(ranks=group_ranks)

    # 在子组内执行通信
    tensor = torch.ones(4).to(device) * rank

    if rank in group_ranks:
        dist.all_reduce(tensor, group=subgroup)
        print(f"[Rank {rank}] Subgroup all_reduce result: {tensor}")

    # 清理子组
    dist.destroy_process_group(subgroup)

2D 并行中的行列进程组

在 2D 并行场景中，需要同时创建行组和列组，实现矩阵式的通信拓扑：

def create_2d_process_groups():
    """创建行/列进程组 (用于 2D 并行)"""
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = f'npu:{rank}'

    # 假设 world_size = 4, 排列为 2x2 网格
    if world_size == 4:
        # 行组: [0,1], [2,3]
        row_group = dist.new_group(ranks=[rank // 2 * 2, rank // 2 * 2 + 1])

        # 列组: [0,2], [1,3]
        col_group = dist.new_group(ranks=[rank % 2, rank % 2 + 2])

        tensor = torch.tensor([rank], dtype=torch.float32).to(device)

        # 在行组内广播
        dist.broadcast(tensor, src=rank // 2 * 2, group=row_group)
        print(f"[Rank {rank}] After row broadcast: {tensor.item()}")

10.5 使用 RPC 进行分布式计算

RPC（Remote Procedure Call）框架允许在分布式环境中的不同节点上执行函数，适用于参数服务器架构、模型并行等场景。昇腾 NPU 通过 TensorPipe 后端提供了完整的 RPC 支持。

注册 NPU RPC 后端

在使用 RPC 之前，需要先注册昇腾 NPU 的 RPC 后端：

import torch
import torch.distributed.rpc as rpc
import torch_npu
from torch_npu.distributed.rpc.backend_registry import _rpc_backend_registry

# 注册 NPU RPC 后端
_rpc_backend_registry()

远程计算示例

以下示例演示如何在远程 worker 上执行计算任务：

def remote_computation(tensor):
    """在远程 worker 上执行的函数"""
    # 移动到本地 NPU 并计算
    local_device = torch.npu.current_device()
    tensor = tensor.to(f'npu:{local_device}')
    result = tensor * 2 + 1
    return result.cpu()

def rpc_example(rank, world_size):
    # 初始化 RPC
    rpc.init_rpc(
        name=f"worker_{rank}",
        rank=rank,
        world_size=world_size,
        backend=rpc.BackendType.NPU_TENSORPIPE,
    )

    if rank == 0:
        # 从 rank 0 发起远程调用
        tensor = torch.randn(4, 4)

        # 异步调用 worker_1 上的函数
        future = rpc.rpc_async(
            "worker_1",
            remote_computation,
            args=(tensor,)
        )

        # 获取结果
        result = future.wait()
        print(f"RPC result shape: {result.shape}")
        print(f"RPC result: {result}")

    # 同步所有 worker
    rpc.shutdown()

if __name__ == "__main__":
    import torch.multiprocessing as mp
    world_size = 2
    mp.spawn(rpc_example, args=(world_size,), nprocs=world_size, join=True)

10.6 完整的分布式训练脚本

以下是一个完整的分布式训练脚本示例，使用 torchrun 启动，适用于 MNIST 图像分类任务。该脚本涵盖了数据加载、模型定义、DDP 包装、训练循环和模型保存等完整流程。

启动方式

torchrun --nproc_per_node=8 --nnodes=1 train_distributed.py

完整脚本

#!/usr/bin/env python3
"""
使用 torchrun 启动:
torchrun --nproc_per_node=8 --nnodes=1 train_distributed.py
"""

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms
import torch_npu

def get_args():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch-size', type=int, default=64)
    parser.add_argument('--epochs', type=int, default=10)
    parser.add_argument('--lr', type=float, default=0.01)
    return parser.parse_args()

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = nn.functional.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)

def train_epoch(model, device, train_loader, optimizer, epoch, rank):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = nn.functional.nll_loss(output, target)
        loss.backward()
        optimizer.step()

        if batch_idx % 100 == 0 and rank == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

def test(model, device, test_loader, rank):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += nn.functional.nll_loss(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)

    if rank == 0:
        print(f'\nTest set: Average loss: {test_loss:.4f}, '
              f'Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

    return accuracy

def main():
    args = get_args()

    # 初始化分布式环境
    dist.init_process_group(backend='hccl')

    rank = dist.get_rank()
    world_size = dist.get_world_size()
    local_rank = int(os.environ.get('LOCAL_RANK', 0))

    # 设置 NPU 设备
    torch.npu.set_device(local_rank)
    device = torch.device(f'npu:{local_rank}')

    if rank == 0:
        print(f"Training with {world_size} NPU devices")

    # 数据加载
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('./data', train=False, transform=transform)

    # 使用分布式采样器
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        sampler=train_sampler,
        num_workers=4,
        pin_memory=True
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=args.batch_size * 2,
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )

    # 创建模型并使用 DDP 包装
    model = Net().to(device)
    model = DDP(model, device_ids=[local_rank])

    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.7)

    # 训练循环
    best_accuracy = 0
    for epoch in range(1, args.epochs + 1):
        train_sampler.set_epoch(epoch)
        train_epoch(model, device, train_loader, optimizer, epoch, rank)
        accuracy = test(model, device, test_loader, rank)
        scheduler.step()

        # 保存最佳模型 (仅 rank 0)
        if rank == 0 and accuracy > best_accuracy:
            best_accuracy = accuracy
            torch.save(model.module.state_dict(), 'best_model.pth')

    if rank == 0:
        print(f"Best accuracy: {best_accuracy:.2f}%")

    dist.destroy_process_group()

if __name__ == '__main__':
    main()

10.7 torch_npu.distributed 特有 API

torch_npu 在标准 PyTorch 分布式 API 的基础上，提供了一些针对昇腾 NPU 平台的特有功能，用于处理不均匀张量通信、获取 HCCL 通信器信息以及实现容错恢复等场景。

不均匀张量通信

在真实训练场景中，不同设备上的数据量可能不一致。torch_npu 提供了处理不均匀张量的 reduce_scatter 和 all_gather 操作：

def torch_npu_specific_apis():
    """演示 torch_npu.distributed 特有的 API"""
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = f'npu:{rank}'

    # 检查 HCCL 是否可用
    print(f"HCCL available: {is_hccl_available()}")

    # ============ 不均匀大小的 reduce_scatter ============
    # 每个 rank 的输入大小不同
    input_split_sizes = [rank + 1 for rank in range(world_size)]
    total_size = sum(input_split_sizes)

    input_tensor = torch.ones(total_size, 4).to(device) * rank
    output_tensor = torch.zeros(input_split_sizes[rank], 4).to(device)

    reduce_scatter_tensor_uneven(
        output_tensor,
        input_tensor,
        input_split_sizes=input_split_sizes,
        op=dist.ReduceOp.SUM
    )
    print(f"[Rank {rank}] Uneven reduce_scatter output shape: {output_tensor.shape}")

    # ============ 不均匀大小的 all_gather ============
    output_split_sizes = [rank + 1 for rank in range(world_size)]
    total_output_size = sum(output_split_sizes)

    input_tensor = torch.ones(output_split_sizes[rank], 4).to(device) * rank
    output_tensor = torch.zeros(total_output_size, 4).to(device)

    all_gather_into_tensor_uneven(
        output_tensor,
        input_tensor,
        output_split_sizes=output_split_sizes
    )
    print(f"[Rank {rank}] Uneven all_gather output shape: {output_tensor.shape}")

HCCL 通信器信息获取与容错

    # ============ 获取 HCCL 通信器名称 ============
    pg = dist.distributed_c10d._get_default_group()
    backend = pg._get_backend(torch.device('npu'))

    # 获取 HCCL comm 名称 (用于调试)
    hccl_comm_name = backend.get_hccl_comm_name(rank)
    print(f"[Rank {rank}] HCCL comm name: {hccl_comm_name}")

    # ============ 重新初始化进程组 (用于容错) ============
    # 在出现通信错误后，可以尝试重建连接
    # reinit_process_group(rebuild_link=False)  # 仅恢复通信器
    # reinit_process_group(rebuild_link=True)   # 完全重建

附录：环境变量配置

在使用昇腾 NPU 进行分布式训练时，可以通过环境变量对 HCCL 通信库进行调优。下表列出了常用的环境变量及其说明：

环境变量	说明
RANK_TABLE_FILE	使用 rank table 文件初始化集群
HCCL_ASYNC_ERROR_HANDLING	异步错误处理模式
HCCL_BLOCKING_WAIT	是否阻塞等待操作完成
HCCL_EXEC_TIMEOUT	HCCL 操作超时时间（秒）
HCCL_DESYNC_DEBUG	启用去同步调试
HCCL_ZERO_COPY	启用零拷贝优化
HCCL_DETERMINISTIC	启用确定性计算
TORCH_HCCL_ENABLE_MONITORING	启用心跳监控线程

总结

本文系统解析了 torch_npu 分布式训练模块的架构设计与核心实现，涵盖模块组成、进程组管理、通信器生命周期、Watchdog 监控机制、后端注册流程等关键机制，并给出了相应的优化建议；同时，通过 FSDP 实现、子进程组通信、RPC 分布式计算以及完整的 DDP 训练脚本等实践示例，展示了在昇腾 NPU 上开展分布式训练的具体方法。理解底层机制并合理利用 torch_npu 提供的 HCCL 后端与特有 API，可以帮助开发者高效开展大规模分布式训练任务。在实际应用中，建议根据具体场景选择合适的并行策略，并通过环境变量对通信性能进行调优，以获得最佳的训练效率。