登录/ 注册
 找回密码
 立即注册
关闭

CSDN热搜

程序园 精品问答 技术交流 资源下载
返回列表 发新帖
首页 › 业界区 › 安全 ›DWDP: 在NVL72上的高性能分布式权重数据并行 ...

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

格恳绌 昨天 20:28
猛犸象科技工作室:网站开发，备案域名，渗透，服务器出租，DDOS/CC攻击，TG加粉引流
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上的高性能分布式权重数据并行

    • 摘要

      • 遇到了什么问题？

        • 现有方案的问题
        • 提出的方案

      • 挑战

    • 运行时方法
    • 分析方法与建模
    • 设计挑战与优化

      • 基线部署和性能分析

        • Kernel分解分析
        • 运行时tracing分析

      • 消除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开销很少。
1.png

上图左侧的阴影部分就是因为工作负载不均衡造成的开销，这就是本次重点解决的层间跨rank同步开销问题。
曾经的解决方法包括：缓存感知调度，负载感知调度，专家负载调度等。然而这些调度方法并没有完全地解决这样的同步开销。并随着负载的多元化而变得更加难以避免。
提出的方案


  • DWDP：分布式权重数据并行。一个推理并行策略，移除了在关键路径上的跨rank的集合同步。

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

挑战


  • 权重的卸载所导致的分离权重管理开销如何解决。
  • 异步远程权重预取在通信侧的表现损失。
对此，文章提出了两个对应的优化策略，并分析了计算-通信交叠带来的剩余影响。
运行时方法

2.png


  • 如上图所示，最核心的想法是：在执行第 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可以独立执行他们的任务。
  • 灵活性：专家放置会变的更加灵活。
然而，这样的新策略同时也带来了新的挑战：

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

3.png

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

  • 对于传统的模型来说，时间来源于计算（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进行了分解分析。我们可以看到下表：
4.png

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

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

5.png

上图通过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 传输过去。
    • 当这个分块的所有专家传输都发起后（因为是异步的我们不需要关心返回），移动到下一个分块进行传输。

也就是文中的伪代码：
  1. copy_plan = []
  2. for each parameter p:
  3.   M = prefetch_size(p)
  4.   for offset = 0 to M step s:
  5.     chunk = min(s, M - offset)
  6.     for peer in round_robin(remote_peers):
  7.       src = peer_ptr(peer, p) + src_offset(peer, p) + offset
  8.       dst = local_buffer(peer, p) + offset
  9.       copy_plan.append((dst, src, chunk))
  11. 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显示“可重用信号”后）发起预取操作。
以下就是整体的工作流程：
6.png

我就不赘述流程了。
关键优化实施

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

在fused_moe_cute_dsl.py中，我们有：
  1. @dataclass
  2. class NvFp4WeightView:
  3.     w3_w1_weight: List[torch.Tensor]
  4.     fc1_weight_scale: List[torch.Tensor]
  5.     fc1_global_scale: List[torch.Tensor]
  6.     w2_weight: List[torch.Tensor]
  7.     fc2_weight_scale: List[torch.Tensor]
  8.     fc2_global_scale: List[torch.Tensor]
  9.     expert_size_per_partition: int
  10.     slot_start: int
  11. # ...
复制代码
当DWDP没有开启的时候，每个list都只有一个张量（也就是本地全部的Experts）。开启的时候，DWDP将具有本地的真实张量，以及分配好的远程预取buffer在张量列表中。
接着，我们可以看到，在run_moe函数中执行了
  1. result = self.run_moe_nvfp4(
  2.   x=x,
  3.   token_selected_experts=token_selected_experts,
  4.   token_final_scales=token_final_scales,
  5.   x_sf=x_sf,
  6.   moe_output=moe_output,
  7.   enable_alltoall=enable_alltoall,
  8.   weight_view=weight_view,
  9. )
复制代码
来执行moe。通过对源码的修改，执行引擎可以通过判断前面 w3_w1_weight 中的张量个数来判断是否启用了DWDP。
  1. is_dwdp = len(weight_view.w3_w1_weight) > 1
  2. forward_impl = self.run_moe_nvfp4_impl_dwdp if is_dwdp else self.run_moe_nvfp4_impl
复制代码
执行 CuteDslFusedMoENvfp4Runner，我们可以看到，其中的forward_impl变成了经过dwdp的代码。于是我们可以看到执行的是：
  1. x, x_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b(
  2.   input=x.view(torch.float4_e2m1fn_x2),
  3.   weight=[
  4.     w.view(torch.float4_e2m1fn_x2) for w in weight_view.w3_w1_weight
  5.   ],
  6.   input_scale=x_sf.view(torch.uint8),
  7.   weight_scale=[
  8.     ws.view(torch.uint8) for ws in weight_view.fc1_weight_scale
  9.   ],
  10.   # ...
  11. )
复制代码
替换成了列表。相对称的finalize部分也修改为了列表。
  1. torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell(
  2.   input=x.view(torch.float4_e2m1fn_x2),
  3.   weight=[
  4.       w.view(torch.float4_e2m1fn_x2) for w in weight_view.w2_weight
  5.   ],
  6.   input_scale=x_sf.view(torch.uint8),
  7.   weight_scale=[
  8.       ws.view(torch.uint8) for ws in weight_view.fc2_weight_scale
  9.   ],
  10.   # ...
  11. )
复制代码
对应的算子实现在cute_dsl_custom_ops.py中。
对于cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b，在进入函数后，注册一个新的op，具体的注册如下：
  1. tuner.choose_one(
  2.             "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b",
  3. # ...
  4. )
  5. #...
  6. @torch.library.register_fake(
  7.         "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b")
  8. def _fake_multi_b(
  9.   #...
  10. )
复制代码
然后旧有的ops直接包装：
  1. def cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell(
  2.   # ...
  3. )
  4. return torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_swiglu_blackwell_multi_b(
  5.   #...
  6. )
复制代码
对于finalize来说这也是相同的。我们也就不再赘述了。同样也是注册函数，旧有的ops包装新的op来实现。
RR+分块实现时分多路复用

实现在dwdp.py的函数prefetch_layer中。
  1. @nvtx_range("dwdp_prefetch_layer")
  2.     def prefetch_layer(self, layer_idx: int, wait_compute_layer_idx: Optional[int] = None):
  3.         with torch.cuda.stream(self.prefetch_buffer.prefetch_stream):
  4.             # 等待上一次的计算同步结束，保证不出现WAW依赖。
  5.             if wait_compute_layer_idx is not None:
  6.                 self.prefetch_buffer.wait_compute_event(wait_compute_layer_idx)
  7.             # 轮询每一个rank并且将对应的块传输过去。
  8.             for peer_rank in range(self.dwdp_size):
  9.                 if peer_rank == self.dwdp_rank:
  10.                     continue  # 跳过本地的张量
  11.                 src_expert_offset = self._get_prefetch_src_offset_from_peer(peer_rank)
  12.                 for param_name in param_names:
  13.                     param_shape = collector.param_shapes[param_name]
  14.                     param_dtype = collector.param_dtypes[param_name]
  15.                     expert_size = param_shape.numel() * param_dtype.itemsize
  16.                     # src_ptr 指向远端节点对应的权重起始地址。
  17.                     base_ptr = collector.get_peer_ptr(peer_rank, param_name)
  18.                     src_ptr = base_ptr + src_expert_offset * expert_size
  19.                     # dst_tensor 本地的对应专家的存放地址。
  20.                     dst_tensor = self.prefetch_buffer.buffers[buffer_idx][param_name][peer_rank]
  21.                     dst_ptr = dst_tensor.data_ptr()
  22.                     data_size = self.num_prefetch_experts * expert_size
  23.                     # 通过cudaMemcpyAsync来实现异步拉取。
  24.                     (err,) = cudart.cudaMemcpyAsync(
  25.                         dst_ptr,
  26.                         src_ptr,
  27.                         data_size,
  28.                         cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice,
  29.                         self.prefetch_buffer.prefetch_stream.cuda_stream,
  30.                     )
  31.                     check_cuda_error(
  32.                         err, f"prefetch layer {layer_idx} peer_rank {peer_rank} {param_name}"
  33.                     )
复制代码
buffer的同步信号管理


  • 主要通过self.prefetch_events和self.compute_events来记录对应的stream记录下来的event。其中compute stream是默认流，prefetch则通过独有的self.prefetch_stream来记录流事件。
events在列表中如下图分布：
7.png

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

笔者很喜欢看这种工程上的讨论——因为他们常常深入到体系结构中。
计算——通信交叠的深度分析

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

这张图显示的是目的端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，进行认真的权衡才能实现更优的优化。

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！
DWDPNVL72高性能分布式权重

使用道具 举报

相关推荐
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册
发帖
格恳绌
昨天 20:28

0

粉丝关注

22

主题发布

板块介绍填写区域，请于后台编辑
微信扫一扫