分布式通信 (distributed/)

分布式推理和多NPU间通信

一、模块概述

distributed目录实现多NPU间的通信和协调，支持张量并行、流水线并行和数据并行。

图1：分布式推理架构
graph LR
    A[分布式推理] --> B[张量并行 TP]
    A --> C[流水线并行 PP]
    A --> D[数据并行 DP]
    B --> E[NPU通信器]
    C --> E
    D --> E

二、文件结构

distributed/
├── __init__.py
├── kv_transfer.py          # KV缓存传输
├── weight_transfer.py      # 权重传输
├── parallel.py             # 并行工具
├── device_communicators/    # 设备通信器
│   └── npu_communicator.py
└── utils.py                # 工具函数

三、通信架构

图2：NPU通信器架构
graph TD
    A[NPU通信器] --> B[HCCL后端]
    A --> C[集合通信]
    A --> D[点对点通信]
    B --> E[昇腾NPU]

kv_transfer.py

KV缓存传输

  • P到D传输
  • KV序列化
  • 异步传输

weight_transfer.py

权重传输

  • 模型权重同步
  • 分片传输
  • 压缩传输

npu_communicator.py

NPU通信器

  • HCCL封装
  • AllReduce
  • AllGather

parallel.py

并行工具

  • 张量并行
  • 流水线并行
  • 数据并行

四、KV缓存传输流程

图3：PD分离架构下的KV传输
graph TD
    A[P节点-预填充] --> B[生成KV缓存]
    B --> C[序列化KV数据]
    C --> D[传输到D节点]
    D --> E[D节点-解码]
    E --> F[返回结果]
首页