DWDP: 在NVL72上的高性能分布式权重数据并行

论文原文 DWDP: Distributed Weight Data Parallelism for High-Performance LLM Inference on NVL72
TensorRT-LLM PR#12136 代码
 TensorRT-LLM 关于DWDP的Blog

笔者注：最简单的想法做出效果很好的优化。但是，为什么别人能够想到而自己想不到呢？有时候很佩服那些能够从简单想法中延伸出来，然后不断与他人合作修正补充最终实现完美的优化的工作的人。
类似于FSDP那样的想法，按需要加载模型参数，计算完成后再释放出去。通过利用cudaMemcpyAsync的Zero-SM占用实现计算-通信的交叠，释放同步点。
如有笔误请批评指正，笔者水平可能不足，谢谢！
人工智能辅助：代码阅读与采用了GLM-5.1. 论文阅读部分参考了GLM-5.1的回答。

DWDP: 在NVL72上的高性能分布式权重数据并行

摘要

遇到了什么问题？

现有的推理并行策略需要逐层-rank间的同步，导致端到端表现对负载不均衡极度敏感。

现有方案的问题

我们知道现有的Moe-LLM推理中，存在着从稠密层到稀疏层的通信模式转换(Attn-Moe), 也存在的从稀疏层到稠密层的通信模式转换(Moe-Attn)。在Megatron-Moe中，这样的转换意味着通信模要从DP-EP-DP的转变。而转变的方式其实就是我们熟悉的all-to-all操作。

如上图所示，当我们的工作负载不均衡时，我们的跨rank的同步开销就成为了问题。为什么会出现这样的不均呢？是因为我们推理过程中，request的长度总是不同的，就很容易造成某些rank开销很长，某些rank开销很少。

上图左侧的阴影部分就是因为工作负载不均衡造成的开销，这就是本次重点解决的层间跨rank同步开销问题。

曾经的解决方法包括：缓存感知调度，负载感知调度，专家负载调度等。然而这些调度方法并没有完全地解决这样的同步开销。并随着负载的多元化而变得更加难以避免。

提出的方案

DWDP：分布式权重数据并行。一个推理并行策略，移除了在关键路径上的跨rank的集合同步。
- ranks仍然保持DP，Moe权重将被跨GPU对卸载来共同保留模型权重的同时，在推理过程中按需要取回。
- 这样，每个rank都可以独立执行推理，无需等待其他的rank，直接减轻了由于工作负载不均衡造成的性能损失。
- 同时，远程权重取回可以掩盖在计算后，在实践中需要高带宽P2P GPU通信，同时需要足够大的逐层计算窗口。

挑战

权重的卸载所导致的分离权重管理开销如何解决。
异步远程权重预取在通信侧的表现损失。

对此，文章提出了两个对应的优化策略，并分析了计算-通信交叠带来的剩余影响。

运行时方法

如上图所示，最核心的想法是：在执行第 l 层的 Moe 计算和 l+1 层的注意力机制DP时，异步加载第 l + 1 层的远程专家到本地的显存中。在执行第 l+1 层的Moe时，需要等待远程权重加载完成。

Q1: 为什么要关注专家的卸载？
A1: 因为在Moe模型中，专家模型的内存占用更大。从DeepSeek中我们可以看到，在运行过程中，参与运算的参数和不参与运算的参数达到了1:18。这大部分是因为Moe中采用TopK激活造成的。

在一个DWDP组内，注意力层将复制到每一个rank上，而专家则分配到不同的rank上。这样，每个rank就只具有自己的本地专家。剩余的专家在其他的Peer GPU上。类似FSDP，在执行到本层的l+1的时候，从远程GPU加载专家到自己的本地显存上。
需要注意的是，为了消除跨rank的同步，DWDP不再采用基于NCCL的集合通信范式（比如All gather），而是采用了 cudaMemcpyAsync 这种0-SM占用的，全部交给 asyncEngine 来执行的方式来实现计算-传输的交叠。相关的资料可以参考异步和交叠计算与数据传输.
这样，每一个rank就可以独立执行计算，不需要等待其他的rank，消除了跨rank的同步点。
除了允许异步推理之外，DWDP还提供了专家存放的灵活性主要体现在如下的两个方面：
- DWDP允许冗余专家的出现。一个专家可以出现在多个rank上并且被peer GPU灵活地拉取。在内存大小允许的情况下，我们甚至可以将经常被多个rank使用的专家冗余地配置到多个rank上，减少我们预取的开销。
- 这样，DWDP就不需要要求专家数必须能够整除rank数目（也就是DWDP组数）。同时，我们可以简单配置每个rank都具有相同的专家数。

因此我们可以总结DWDP的两大优势：

异步性：减少了同步点，每个GPU可以独立执行他们的任务。
灵活性：专家放置会变的更加灵活。

然而，这样的新策略同时也带来了新的挑战：

如何高效地分配专家，也就是高效的权重分离管理？
异步地远程专家预取引入了通信侧的开销，包括了计算-通信干扰，以及多对一问题。

分析方法与建模

分析方法也很好理解——就是分析如上图所示的时间模型，每一层执行的时间。

对于传统的模型来说，时间来源于计算（Attention，Moe）与通信（两次AlltoAll）将token发送出去。所以为 \(T_{\text{comp}} + T_{\text{comm}}\)
对于DWDP，在执行Moe的时候预取了下一层Moe需要的专家，因此执行时间模型就是 \(\max{(T_{\text{comp}}, T_{\text{prefetch}})}\)

这样我们就可以估算计算时间了。计算开销延时就是峰值运算与吞吐的最大值。这样就有 \(T_{\text{op}}=\max{(F/P_{\text{peak}}, B/BW_{\text{mem}})}\). 其中：

\(P_{\text{peak}}\): 计算吞吐峰值
\(BW_{\text{mem}}\): 内存带宽
\(F\): 操作的运算量，单位为 Flop
\(B\): 内存占用流量。

这样我们需要关注的就是两个比率：

\(T_{\text{comp}}/T_{\text{prefetch}}\): 掩盖比率。计算量越大，掩盖效果越好。比率超过一则预取不再是瓶颈。
\(T_{DEP}/T{DWDP}\): 理论加速比。新方法比旧方法好在哪里？

文章发现DEP在16K 长度的token后表现差于DWDP，因此适用于具有长序列的推理场景下。文中还提到，在batchSize增大的情况下，更小的序列长度甚至也能表现的更好（因为实验在batchSize = 1下进行）。

因为随着序列长度的增加，计算时间不断上升，因此当序列增加超过一定量时，加速比将降低，此时需要专注于计算的优化。并且，这个模型是理想化的，忽略了例如计算与权重预取的冲突，以及工作负载不均的情况。

设计挑战与优化

回顾我们的挑战一节，这一节将会讲述如何实现 高效权重分离 和 异步预取造成的多对一冲突 问题。

基线部署和性能分析

Kernel分解分析

首先文章对4-rank下的传统DP+EP与DWDP进行了分解分析。我们可以看到下表：

我们可以从表中发现DWDP带来的优劣：

Pros：消除了同步开销，消除了通信开销。
Cons：计算和通信交叠导致了算子延时的普遍下降；引入了设备内的拷贝（D2D copy，因为远程专家和本地专家需要拼接成连续的buffer），以及设备间的拷贝（也就是cudaMemcpyAsync带来的GPU拉取的开销）。

运行时tracing分析

上图通过nsight分析可以看到多对一带来的通信开销。可以看到，当多个rank同时从一个rank获取专家的时候，这个rank将会序列化执行pull请求，导致prefetch时间超出了计算时间。这样就会导致气泡的出现。

消除D2D时延开销（权重合并开销）

解决方法：通过CuTeDSL拓展Moe GroupedGEMM，使得它能够直接消费多个权重buffer。
- 增加了基于TensorList的输入。Kernel可以在内部选择本地或者预取的专家。
- 并且保留了对曾经的布局和分片的兼容。
即使引入了额外的计算，端到端的性能分析显示没有很严重的开销。

异步的通信开销的减轻

首先文章分析了越大的DWDP组越容易出现多对一的冲突。我们就不再赘述方法了。我们来看他是怎么通过 时分多路复用的复制 来实现减轻多对一通信负担的。

时分多路复用的拷贝

这里的设计其实就是利用了分块+RR调度的策略。我们简洁明了地概括如下：

对于每个等待传输的专家权重 \(E\)
- 首先，设置一个分块大小 \(s\), 这个 \(s\) 足够小，使得拷贝引擎CE能够同时支持拷贝两份分块。(也就是支持至少 \(2s\) 传输)
- 接着，将专家的权重切成 \(\lceil E/s \rceil\) 块。开始针对每个分块 chunk_i 轮询。
- 然后，内部循环中，针对这个分块，轮询每个需要这个专家的rank。将对应的分块 chunk_i 传输过去。
- 当这个分块的所有专家传输都发起后（因为是异步的我们不需要关心返回），移动到下一个分块进行传输。

也就是文中的伪代码：

copy_plan = [] 
for each parameter p: 
  M = prefetch_size(p) 
  for offset = 0 to M step s: 
    chunk = min(s, M - offset) 
    for peer in round_robin(remote_peers): 
      src = peer_ptr(peer, p) + src_offset(peer, p) + offset 
      dst = local_buffer(peer, p) + offset 
      copy_plan.append((dst, src, chunk)) 

return copy_plan

很像网络中的 packet-spray 不是吗～

在这样的设置下，如果流水线引擎能够保证两个切片能够同时传输，rank级别的数据拉取就不会在冲突度为2的情况下出现降速（也就是2个远程rank拉取同一个rank的专家）。这样也可以缓解因为网络冲突带来的性能波动。

部署实施

接下来我们将结合Blog来分析一下DWDP的运行时组件和运行时的数据流。

关键组件

目前，DWDP使用在分离式服务的上下文服务器上。

配置

dwdp_size: DWDP组中的GPU数量
num_groups: DWDP组数量
num_experts_per_worker: 本地专家数量
num_prefetch_experts: 拉取的远程专家的数量

`DwdpLayerHandleCollector`

实现在 tensorrt_llm/_torch/pyexecutor/dwdp.py 中。主要用来：

注册CUDA IPC的句柄，提供给这一层的本地Moe权重和相关的张量。主要实现在register_weights函数和_register_param中。
同时，记录张量的形状，数据类型和位置（偏移）
同时还持有 peer 指针，指向当前层的远程Moe权重，以及相关的远程张量。

`DwdpPrefetchBuffer`

这是一个运行时buffer，用来储存预取回的远程专家权重。他的角色主要是：

保证下一个专家需要时已经准备好。
保证现在的专家在计算过程中不会被覆盖。

起到提供存储，stream和事件通知的作用。

双缓冲区来保证同步过程中不需要等待WAW依赖。
专属的预取stream（也就是cudaMemcpyAsync需要的stream）
就绪信号和重用信号。前者表示数据已经准备完成，后者表示计算已经完成，可以重用这个缓冲区。

双缓冲区真的已经是高性能通信的标配了～

`DwdpManager`

这就是整个DWDP的控制中心。他将管理整个DWDP的生命周期，并在预取过程中起到关键的协调作用。他负责：

从MPI中形成整个DWDP组
创建并且追踪一个Dwdp层的句柄，每个句柄对应了一个开启DWDP的层（handle对应层在笔者前面的nccl-ep中也是一样的设计想法～）
跨Rank，all-gather所有的元数据。元数据记录了每个GPU的本地expert。
收集并且初始化预取的buffer。
在正确的时机（第l层的Moe执行完成，buffer显示“可重用信号”后）发起预取操作。

以下就是整体的工作流程：

我就不赘述流程了。

关键优化实施

消除 Split-weight 合并开销(D2D开销) -> 修改GroupedGemm，使用TensorList

在fused_moe_cute_dsl.py中，我们有：

@dataclass
class NvFp4WeightView:
    w3_w1_weight: List[torch.Tensor]
    fc1_weight_scale: List[torch.Tensor]
    fc1_global_scale: List[torch.Tensor]
    w2_weight: List[torch.Tensor]
    fc2_weight_scale: List[torch.Tensor]
    fc2_global_scale: List[torch.Tensor]
    expert_size_per_partition: int
    slot_start: int
# ...

当DWDP没有开启的时候，每个list都只有一个张量（也就是本地全部的Experts）。开启的时候，DWDP将具有本地的真实张量，以及分配好的远程预取buffer在张量列表中。

接着，我们可以看到，在run_moe函数中执行了

result = self.run_moe_nvfp4(
  x=x,
  token_selected_experts=token_selected_experts,
  token_final_scales=token_final_scales,
  x_sf=x_sf,
  moe_output=moe_output,
  enable_alltoall=enable_alltoall,
  weight_view=weight_view,
)

来执行moe。通过对源码的修改，执行引擎可以通过判断前面 w3_w1_weight 中的张量个数来判断是否启用了DWDP。

is_dwdp = len(weight_view.w3_w1_weight) > 1
forward_impl = self.run_moe_nvfp4_impl_dwdp if is_dwdp else self.run_moe_nvfp4_impl

执行 CuteDslFusedMoENvfp4Runner，我们可以看到，其中的forward_impl变成了经过dwdp的代码。于是我们可以看到执行的是：

x, x_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b(
  input=x.view(torch.float4_e2m1fn_x2),
  weight=[
    w.view(torch.float4_e2m1fn_x2) for w in weight_view.w3_w1_weight
  ],
  input_scale=x_sf.view(torch.uint8),
  weight_scale=[
    ws.view(torch.uint8) for ws in weight_view.fc1_weight_scale
  ],
  # ...
)

替换成了列表。相对称的finalize部分也修改为了列表。

torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell(
  input=x.view(torch.float4_e2m1fn_x2),
  weight=[
      w.view(torch.float4_e2m1fn_x2) for w in weight_view.w2_weight
  ],
  input_scale=x_sf.view(torch.uint8),
  weight_scale=[
      ws.view(torch.uint8) for ws in weight_view.fc2_weight_scale
  ],
  # ...
)

对应的算子实现在cute_dsl_custom_ops.py中。

对于cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b，在进入函数后，注册一个新的op，具体的注册如下：

tuner.choose_one(
            "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b",
# ...
)
#...
@torch.library.register_fake(
        "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b")
def _fake_multi_b(
  #...
)

然后旧有的ops直接包装：

def cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell(
  # ...
)
return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b(
  #...
)

对于finalize来说这也是相同的。我们也就不再赘述了。同样也是注册函数，旧有的ops包装新的op来实现。

RR+分块实现时分多路复用

实现在dwdp.py的函数prefetch_layer中。

@nvtx_range("dwdp_prefetch_layer")
    def prefetch_layer(self, layer_idx: int, wait_compute_layer_idx: Optional[int] = None):
        with torch.cuda.stream(self.prefetch_buffer.prefetch_stream):
            # 等待上一次的计算同步结束，保证不出现WAW依赖。
            if wait_compute_layer_idx is not None:
                self.prefetch_buffer.wait_compute_event(wait_compute_layer_idx)
            # 轮询每一个rank并且将对应的块传输过去。
            for peer_rank in range(self.dwdp_size):
                if peer_rank == self.dwdp_rank:
                    continue  # 跳过本地的张量
                src_expert_offset = self._get_prefetch_src_offset_from_peer(peer_rank)
                for param_name in param_names:
                    param_shape = collector.param_shapes[param_name]
                    param_dtype = collector.param_dtypes[param_name]
                    expert_size = param_shape.numel() * param_dtype.itemsize
                    # src_ptr 指向远端节点对应的权重起始地址。
                    base_ptr = collector.get_peer_ptr(peer_rank, param_name)
                    src_ptr = base_ptr + src_expert_offset * expert_size
                    # dst_tensor 本地的对应专家的存放地址。
                    dst_tensor = self.prefetch_buffer.buffers[buffer_idx][param_name][peer_rank]
                    dst_ptr = dst_tensor.data_ptr()
                    data_size = self.num_prefetch_experts * expert_size
                    # 通过cudaMemcpyAsync来实现异步拉取。
                    (err,) = cudart.cudaMemcpyAsync(
                        dst_ptr,
                        src_ptr,
                        data_size,
                        cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice,
                        self.prefetch_buffer.prefetch_stream.cuda_stream,
                    )
                    check_cuda_error(
                        err, f"prefetch layer {layer_idx} peer_rank {peer_rank} {param_name}"
                    )

buffer的同步信号管理

主要通过self.prefetch_events和self.compute_events来记录对应的stream记录下来的event。其中compute stream是默认流，prefetch则通过独有的self.prefetch_stream来记录流事件。

events在列表中如下图分布：

用于记录对应层的对应chunk块是否完成。

讨论

笔者很喜欢看这种工程上的讨论——因为他们常常深入到体系结构中。

计算——通信交叠的深度分析

为什么DWDP造成了计算效率的下降呢？这是因为硬件内部存在了冲突。我们可以看下图所示的通路示意图：

这张图显示的是目的端GPU通过CE拉取源端GPU的专家权重的过程图。CE专注于数据搬运，因此它无需占用SM。

然而，仔细分析数据流就会发现：数据会经过NoC，L2-Cache，DRAM，而GPU的kernel在执行并发的内存请求。这样，两者在硬件上存在着相互干扰的问题。

跨GPU传输的完整路径包括了：

源端：从显存（DRAM）读取 → 经过 L2 缓存 → 经过片上网络（NoC）→ 注入NVLink 控制器。
目的端：从 NVLink 接收 → 经过 NoC 和 L2 → 写入本地 DRAM。

这样发生干扰的硬件定位在：(1) NoC的仲裁：NoC要选择执行内核指令和CE指令，两者争抢带宽。(2) L2 Cache：命中率降低。一致性维护成本升高。(3) DRAM的带宽将会争用。但最重要的是 (4) 功耗冲突。GPU达到功耗墙导致其在达到内存带宽限制前主动降频。

后记

DWDP和特定的硬件平台绑定很深——他要求必须是具有CE和nvlink的节点内通信。一旦跨节点采用这种方法，很有可能导致通信时延远大于计算时延，造成严重的流水线气泡——严重影响端到端延迟。因此常见的多节点集群——通过RDMA，采用的是传输token而不是传输专家权重的方式。但是这也让我们看到了，系统优化是无法和特定硬件平台解耦的，我们必须要建立好标准的流水线模型和roofline，进行认真的权衡才能实现更优的优化。

此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。原文来自 — 版权归原作者所有。

推荐订阅源

博客园 - 木木ちゃん