跳转至

层间线性层分片指南

概述

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

与在每个设备上复制所有权重不同,层间线性层分片将一系列此类算子的权重分片存储在通信组内的各个 NPU 设备上

  • The i-th layer's linear weight is stored only on device i % K, where K is the number of devices in the group.
  • 其他设备在初始化期间持有一个轻量级的**共享虚拟张量**,并在前向传播过程中通过**异步广播**按需获取真实权重。

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

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

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

流程图

层间分片

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


快速上手

要启用 层间线性层分片,请在启动推理任务时使用 --additional-config 参数指定目标线性层。例如,要对 o_projq_b_proj 层进行分片,请使用:

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

Restriction 层间分片只能在 PD 解耦架构的 P 节点 中启用。 Layer Sharding is not supported by RFork weight transfer. If --load-format rfork is used with layer_sharding, RFork transfer is bypassed and the model is loaded through the default model loader.


支持的场景

该特性在以下情况下收益最大:

启用 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_projo_proj 层都要求在每一层存储巨大的权重矩阵。通过在 NPU 之间对这些层进行分片,有助于将极深的模型(如 61 层架构)加载到有限的设备显存中。

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

配置示例:

export VLLM_ASCEND_ENABLE_FLASHCOMM1=1
vllm serve \
  --model DeepSeek-V3.2 \
  --additional-config '{
    "layer_sharding": ["q_b_proj", "o_proj"]
  }'