层间线性层分片

层间线性层分片#

概述#

Layer Shard Linear 是一种专为大语言模型 (LLM) 推理设计的内存优化特性。它旨在解决由于跨多层的重复线性算子（结构相同但权重不同）所带来的巨大内存压力。

Layer Shard Linear 不再在每个设备上复制所有权重，而是将一系列此类算子的权重分片存储在通信组内的各个 NPU 设备上：

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

如下图所示，该设计实现了权重的预取：在计算当前层（例如 MLA 或 MOE）时，系统会在后台异步广播下一层的权重。由于 MLA 模块中的 Attention 计算具有足够的延迟带宽，o_proj 的权重传输可以完全与计算重叠，从端到端推理的角度来看，通信达到了零延迟的效果。

这种方法在保持完全一致的计算语义的同时，显著降低了 NPU 显存占用，这对于以下情况至关重要：

极深的模型架构（例如拥有 61 层的 DeepSeek-V3/R1）；
使用 DSA-CP 或 FlashComm2 的模型，这些场景要求每一层的完整 O (输出) 投影矩阵必须驻留在内存中；
Attention 计算延迟能够完全覆盖（隐藏）权重广播通信开销的场景。

流程图#

层间分片

图示： Layer Shard Linear 工作流。上方显示权重按层分片存储在不同设备上；下方显示在前向执行期间，异步广播在当前层计算时预取下一层的权重，从而实现零开销的权重加载。

快速上手#

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

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

支持的场景#

该特性可以在任何场景下启用，但在以下情况下收益最大：

启用 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 层架构）加载到有限的设备显存中。