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上的高性能分布式权重数据并行
- 摘要
- 运行时方法
- 分析方法与建模
- 设计挑战与优化
- 基线部署和性能分析
- 消除D2D时延开销（权重合并开销）
- 异步的通信开销的减轻
- 部署实施
- 关键组件
- 配置
- DwdpLayerHandleCollector
- DwdpPrefetchBuffer
- DwdpManager
- 关键优化实施
- 消除 Split-weight 合并开销(D2D开销) -> 修改GroupedGemm，使用TensorList
- RR+分块实现时分多路复用
- buffer的同步信号管理
- 讨论
- 后记
摘要
遇到了什么问题？
- 现有的推理并行策略需要逐层-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
在这样的设置下，如果流水线引擎能够保证两个切片能够同时传输，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中，我们有：- @dataclassclass 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# ...
接着，我们可以看到，在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,)
- is_dwdp = len(weight_view.w3_w1_weight) > 1forward_impl = self.run_moe_nvfp4_impl_dwdp if is_dwdp else self.run_moe_nvfp4_impl
- 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 ], # ...)
- 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_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( #...)
- def cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell( # ...)return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b( #...)
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}" )
- 主要通过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，进行认真的权衡才能实现更优的优化。
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！ |
提升卡
置顶卡
沉默卡
喧嚣卡
变色卡
千斤顶
照妖镜
