层分片线性算子指南

层分片线性算子指南#

概述#

层分片线性算子是一项专为大语言模型推理设计的内存优化特性。它解决了因多层中重复出现结构相同但权重不同的线性算子所导致的高内存压力问题。

层分片线性算子并非在每个设备上复制所有权重，而是将“一系列”此类算子的权重分片存储到通信组内的多个NPU设备上：

第 i 层的线性权重仅存储在设备 i % K 上，其中 K 为通信组内的设备数量。
其他设备在初始化时持有一个轻量级的共享虚拟张量，并在前向传播过程中通过异步广播按需获取真实权重。

如下图所示，该设计使得广播能够覆盖权重：当当前层（例如MLA或MOE）正在计算时，系统在后台异步广播下一层的权重。由于MLA模块中的注意力计算具有足够的延迟限制，o_proj的权重传输与计算完全重叠，从而使得通信从端到端推理的角度来看是无延迟的。

该方法在显著减少NPU内存占用的同时保持了精确的计算语义，尤其适用于以下场景：

极深架构（例如具有61层的DeepSeek-V3/R1）；
使用**DSA-CP或FlashComm2**的模型，其中每层必须将完整的O（输出）投影矩阵驻留在内存中；
注意力计算延迟完全重叠（隐藏）权重广播通信成本的场景。

流程图#

层分片

图. 层分片线性算子工作流程：权重按层分片到各个设备上（上图），在前向执行期间（下图），异步广播在当前层计算时预取下一层的权重，从而实现零开销的权重加载。

快速开始#

要启用层分片线性算子，请在启动推理任务时使用--additional-config参数指定目标线性层。例如，要对o_proj和q_b_proj层进行分片，请使用：

--additional-config '{
  "layer_sharding": ["o_proj", "q_b_proj"]
}'

限制层分片只能在PD分离架构的P节点中启用。RFork权重传输不支持层分片。如果在使用layer_sharding的同时使用了--load-format rfork，则会绕过RFork传输，并通过默认模型加载器加载模型。

支持场景#

该特性在以下场景中效果最佳：

启用FlashComm2#

使用FlashComm2时，每层必须将完整的输出投影（o_proj）矩阵驻留在内存中。层分片通过将这些权重分布到多个设备上，显著降低了内存压力。

示例配置：

export VLLM_ASCEND_FLASHCOMM2_PARALLEL_SIZE=1
vllm serve \
  --model DeepSeek-V3/R1 \
  --additional-config '{
    "layer_sharding": ["o_proj"]
  }'

启用DSA-CP#

使用DSA-CP时，q_b_proj和o_proj层都需要为每层存储大型权重矩阵。将这些层分片到多个NPU上，有助于将极深模型（例如61层架构）适配到有限的设备内存中。

层分片只能在PD分离架构的P节点中启用。