DeepSpeed-Domino 深度解析：通过通用张量切分与计算通信重叠实现单机与多节点统一的 LLM 张量并行训练

DeepSpeed-Domino 深度解析通过通用张量切分与计算通信重叠实现单机与多节点统一的 LLM 张量并行训练【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文围绕微软 DeepSpeed 团队开源的DominoProject Domino通信隐藏张量并行训练引擎展开。Domino 打破传统张量并行Tensor ParallelismTP中计算-通信的串行数据依赖通过将一次 batch 训练迭代拆分为相互独立的子任务并流水化执行在单节点与多节点场景下实现近乎完全的计算通信隐藏。读完本文你将掌握 TP 通信开销的成因、Domino 的三种张量切分策略按行切输入、按列切权重、二维切分及其数据依赖分析、前向/反向中的 NCCL 调用触发时序以及在当前 DeepSpeed 仓库中对应的模块实现与后续上手路径。目录背景为什么需要统一的单机/多节点 TP 方案Domino 亮点一览TP 通信开销的来源与数据依赖分析Domino 核心设计三种张量切分策略实现与优化前向、反向与通用内核优化仓库中的 Domino 源码实现印证如何尝试 DeepSpeed-Domino参考引用与致谢背景为什么需要统一的单机/多节点 TP 方案大语言模型LLM如 GPT、Llama 系列普遍基于 Transformer 网络结构其规模已远超单张加速卡GPU的内存容量训练与推理都需要多 GPU 乃至多节点的分布式处理。张量并行TP是训练 LLM 的主流分布式技术之一它将每个模型层按维度切分到多张 GPU 上从而聚合多卡的聚合内存容纳大模型。但 TP 对每个被切分的层在前向与反向各引入两次通信集合操作即每个 Transformer block 共 4 次 AllReduce。TP 在单机场景颇具吸引力——NVLink/NVSwitch 等节点内高带宽直连保证了优异的系统效率然而在跨节点场景下节点间互联带宽有限TP 的通信开销会显著放大。博客中引用的先前研究参见 blogs/deepspeed-domino/README.md报告通信可占端到端训练时间的 75%即便在配备 400GB/s 高速 InfiniBand 的最新 DGX-H100 节点上通信开销仍可占单次训练迭代端到端时间的 43%。同时更快的新 GPU如 DGX-B200意味着算力增长将进一步放大 TP 通信在端到端时间中的占比。需要说明的是博客同时强调一个硬件趋势NVIDIA 正在不断缩小节点内NVLink单台 DGX-H100 约 900GB/s与节点间每节点 8 张 ConnectX-7 网卡、聚合带宽 400GB/s的带宽差距。正因为二者带宽已处于同一量级针对单节点与多节点提供一套统一 TP 方案的时机已经成熟——这正是 Domino 的出发点。Domino 亮点一览Domino 的定位是TP 优化技术与统一的 TP 训练引擎其最核心的特性与收益包括近乎完全的计算通信隐藏通过把一个 batch 的训练迭代分解为更小且相互独立的子任务实现高效流水化从而在计算背后近乎完全地隐藏通信。单机与多节点的统一方案Domino 是首个同时覆盖单节点与多节点的统一 TP 方案。传统 TP 方案如 Megatron-LM受限于跨节点通信带宽而在多节点场景明显乏力。实测性能来自博客在 1 至 4 台 DGX-H100、每台 8 张 H100节点内 NVLink 900GB/s、节点间 IB 400GB/s 环境下的测试报告针对不同模型规模、序列长度与 batch 大小在 GPT 与 Llama 模型系列上Domino 相对 Megatron-LM 的端到端训练迭代吞吐分别最高提升约1.3x与1.2x在若干场景下 Domino 能达到接近最优的训练吞吐最优指假设禁用 TP 通信集合时所能达到的吞吐。需要强调的是上述数字来自博客自述的特定软硬件环境DGX-H100 400GB/s IB属于该项目的实验性报告结果并不能代表在所有硬件、模型与软件栈下的普适结论其详细实验细节请参阅博客引用的 arXiv 论文 2409.15241。TP 通信开销的来源与数据依赖分析每个 Transformer block 的 4 次 AllReduce按 Megatron-LM 论文的叙述TP 下每个 Transformer block1 个 Self-Attention 层 1 个 MLP 层在前向与反向上各执行 2 次 AllReduce共计4 次 AllReduce。对一个由 $N$ 个 Transformer block 堆叠而成的 LLMTP 训练每次迭代需要执行 $4 \times N$ 次 AllReduce。即便是 GPT-3 2.7B / 6.7B 这样的小模型32 层每次训练迭代也要执行 128 次 AllReduce且该数量随层数线性增长。实测通信占比43%47%由于通信处于关键路径且难以被计算掩盖原因见下文数据依赖分析TP 的通信开销非常突出。博客提供的实测用 Megatron-LM 在 1 到 4 台 DGX-H100、即 8 到 32 张 H100 上训练 GPT-3 与 Llama-2 各尺寸模型、不同 batch 大小显示即便使用带 400GB/s 跨节点带宽的最新 DGX-H100 硬件通信开销最高仍达 47%。博客进而推算Llama-3 405B 在 16,000 张 H100 上训练约需 54 天其中投影的通信时间可能高达约25 天——可见即便有了 NVLink/InfiniBand 这类高带宽互联TP 通信仍是端到端训练时间的巨大开销来源。TP 的数据依赖通信为何难以隐藏传统 TP 中一个 Transformer 层Attn 或 MLP的计算可抽象为 $XABY$对 Attention 层$X$ 为输入$A$ 为注意力计算如 multi-head attention$B$ 为线性层对 MLP 层$A$、$B$ 均为线性层计算完成后需要对 $Y$ 做一次 AllReduce。由于 $Y$ 上的计算$XABY$与通信AllReduce($Y$)之间存在串行数据依赖AllReduce($Y$) 会完全裸露出来、无法与计算重叠这正是 TP 在通信带宽受限场景下效率不高的根本原因。近期的 GeMMNCCL 内核融合工作受限于计算-通信重叠范围过窄通常只能与单个 GeMM 重叠也无法完全隐藏通信。Domino 核心设计三种张量切分策略针对上述 $XABY$ 的串行依赖Domino 通过以下三种方式将计算拆碎对输入 $X$ 做按行切分Row-wise Split on Inputs X对权重 $B$ 做按列切分Column-wise Split on Weights B将两者结合的2D 切分方案同时切 $X$ 与 $B$。计算被拆成互相独立的片段后Domino 即可让计算与通信分别作用于不同的独立片段上流水执行从而实现近完全的计算通信隐藏。为便于说明博客统一以把张量切成两块为例且仅以前向传播说明反向只是逆序执行。按行切分输入 X可达 100% 的通信隐藏Domino 沿行维度即 batch 维度切分输入 $X$数据依赖分析按 batch 维切分输入在层内intra-layer与层间inter-layer两个层面都不存在数据依赖。因此可以同时实现层内重叠AllReduce($Y_1$) 与 $X_2AB$ 重叠层间重叠AllReduce($Y_2$) 与下一层的 $X_1AB$ 重叠。基于这种 batch 切分Domino 能将通信的隐藏比例推到接近100%。不过仅有按行切分会把张量切得越来越窄损害内核计算效率因此还需要下述按列切分方案配合。按列切分权重 B50%70% 的通信隐藏Domino 沿列维度最后一维切分权重矩阵 $B$数据依赖分析权重 $B$ 的按列切分在层内无数据依赖但在层间存在数据依赖因此只能实现层内AllReduce($Y_1$) 与 $X_2AB$的计算-通信重叠。实践上纯按列切分权重约可隐藏 50%70% 的通信该方案之所以必要是因为它避免了窄张量对内核效率的损害。2D 切分同时切 X 与 B面向超大模型的折衷针对极大 LLMDomino 同时在行维度切输入 $X$、在列维度切权重 $B$。该策略有利于既要求低显存占用、又希望最小化通信开销的模型训练场景。由于按列切分权重会在每个 Transformer 层末尾引入一次同步AllReduce2D 方案同样只能实现层内的计算-通信重叠。Domino 相比 GeMMNCCL 融合的优势综合而言Domino 相比传统融合方案的独特价值可归纳为重叠范围更广Domino 打破了数据依赖可让一个 AllReduce 不仅与单个 GeMM 重叠还能把重叠范围扩展到多个 GeMM、LayerNorm、Dropout 等一串计算内核序列接近最优的系统吞吐通信被近完全隐藏因此在特定场景下可逼近禁用 TP 通信的理想吞吐内核层面的通用性Domino 工作在内核调度器层面任何内核优化或新内核都能无缝集成进 Domino 框架切分方案简单通用便于用户在遇到溢出、权重/梯度异常时进行端到端的正确性调试。实现与优化前向、反向与通用内核优化博客以按行切分输入的实现为主线介绍其余细节见论文。以下前向/反向描述均以把一个 batch 切成两个 $\mu$-batch$\mu$-batch0 与 $\mu$-batch1为例。前向多个计算片段轮转重叠 AllReduce在 图Transformer block1 self-attn 1 MLP前向流程上为朴素 TP、下为 Domino 实现博客图 9 所示调度中记 $\mu$-batch0 的注意力输出为 attn0、MLP 输出为 MLP0$\mu$-batch1 的对应输出为 attn1、MLP1。AllReduce(attn0)与 $\mu$-batch1 上的 self-attention 计算重叠AllReduce(attn1)通过与多个 $\mu$-batch 的 Dropout、Residual、LayerNorm 计算构成的小内核分组重叠——这个小内核分组不仅能完整隐藏 AllReduce(attn1)还为图 10 反向中的 AllReduce(MLP0) 预留了合适的重叠空间AllReduce(MLP0)藏在 $\mu$-batch1 的 MLP 计算内核序列GeMM GeLU GeMM之后AllReduce(MLP1)藏在下一层next layer的 attention 计算之后。可见 Domino 把单一 AllReduce 的重叠范围从一个 GeMM扩展到了一串内核序列。反向扩展与权重梯度计算的重叠图Transformer block 反向流程博客图 10 展示了两个 $\mu$-batch 的反向示意。除与前向类似的隐藏策略外反向还进一步把通信重叠范围扩展到同一 $\mu$-batch 内部的权重梯度计算例如 AllReduce(MLP1) 可部分与其自身 $\mu$-batch1 的权重梯度计算重叠。原因在于每一个 grad matmul 都包含输入梯度与权重梯度两个独立的 GeMM因此可以把 AllReduce(MLP1) 与 $\mu$-batch1 的权重梯度计算重叠起来。反向相比前向更棘手的一点是反向计算图由 torch.autograd() 自动生成难以精确控制 NCCL 调用的触发时刻。Domino 为此实现了一个no-operationno-op模块它在前向阶段获取通信句柄communication handle并在反向阶段继续持有使用从而在不重写自定义反向计算图的前提下与 torch.autograd() 无缝协作、精确控制每次 NCCL 的开始/结束时机该模块的实现细节见下文仓库源码印证中的NoOper。通用内核优化Domino 还采纳了通用的内核级优化手段使用CUDA Graph压缩相邻计算内核之间的空闲/bubble 时间以降低端到端延迟使用CUDA 多流multi-stream提升并行执行能力并利用torch.compile()进一步提升系统效率。仓库中的 Domino 源码实现印证Domino 的博客文档是上述设计的第一手出处而当前仓库也在deepspeed/runtime/domino/目录下提供了与之一一对应的参考实现注意该目录当前只有__init__.py、async_linear.py、transformer.py三个文件属于偏教学/原型性质的模块级实现并非完整的端到端训练流水线。deepspeed/__init__.py的第 58 行通过from .runtime import zero, domino将 domino 模块接入 DeepSpeed 包命名空间。仓库中的教程页 docs/_tutorials/domino.md 也确认 Domino 用于在张量并行训练中实现近完全的计算通信隐藏。通信句柄的中转站DominoUtil 与 NoOperdeepspeed/runtime/domino/transformer.py 中class DominoUtil: BATCH_0 BATCH0 BATCH_1 BATCH1 HANDLE_DIC {BATCH0: None, BATCH1: None}DominoUtil以两个键BATCH0/BATCH1对应两个 $\mu$-batchHANDLE_DIC用于在前向/反向间传递异步 AllReduce 的通信句柄。这正是博客中 no-operation 模块的落地实现class NoOper(torch.autograd.Function): staticmethod def forward(ctx, input_, handle_dic, h_id): ctx.handle_dic handle_dic ctx.h_id h_id return input_ staticmethod def backward(ctx, grad_output): handle ctx.handle_dic[ctx.h_id] handle.wait() return grad_output, None, NoneNoOper.forward原样透传输入不产生额外计算symbolic定义为返回输入自身以便配合 torch 图模式真正的等待动作被推迟到backward中——取出前向时存入HANDLE_DIC的异步通信句柄并执行handle.wait()。由于 PyTorch 的 autograd 只在反向需要该张量梯度时才触发 backward这恰好实现了精确控制 NCCL 完成/等待时机的效果且无需改写任何自定义反向图与博客描述完全吻合。异步 AllReduce 的发起DominoAsyncColumnParallelLineardeepspeed/runtime/domino/async_linear.py 定义了DominoAsyncColumnParallelLinearImpl一个torch.autograd.Functionstaticmethod def backward(ctx, grad_output): inp, weight, bias ctx.saved_tensors ... grad_input torch.matmul(grad_output, weight) handle dist.all_reduce(grad_input, groupTP_group, async_opTrue) ctx.handle_dic[ctx.h_id] handle ...反向中计算出的grad_input并不立即等待 AllReduce 完成而是以async_opTrue发起 AllReduce 并把返回句柄写入ctx.handle_dic[ctx.h_id]——即前向中由ShardedAttention/DominoTransformerLayer传入的DominoUtil.HANDLE_DIC。配合NoOper在后续依赖点wait()构成了异步发起 延迟等待的重叠闭环。同文件中还提供了RowParallelLinearNoComm不触发通信的按行并行线性层用于dense/linear_fc2等输出侧投影。μ-batch 流水DominoTransformer 与 DominoTransformerLayerdeepspeed/runtime/domino/transformer.py 中的DominoTransformer在前向入口通过torch.chunk(hidden_states, chunks2, dim1)把 batch 维一分为二即博客所述的 $\mu$-batch0 / $\mu$-batch1 拆分随后依据配置项选择两条不同的重叠执行路径self._forward_impl self.inter_layer_overlap_forward if config.domino_intra_layer_overlap: self._forward_impl self.intra_layer_overlap_forwardinter_layer_overlap_forward对应按行切分输入 X 的层间 层内重叠流水与博客图 6 的 100% 隐藏能力对应它把 attention、residual、MLP 等子步骤按两个 μ-batch 交错编排并通过fwd_handle0/fwd_handle1两个异步 AllReduce 句柄实现跨层轮转等待intra_layer_overlap_forward与DominoTransformerLayer.forwarddeepspeed/runtime/domino/transformer.py对应仅层内重叠的按列切分或 2D策略在每个 Transformer 层内部依次为 $\mu$-batch0、$\mu$-batch1 发起异步 AllReduce并在残差连接前wait()对应句柄。DominoTransformerLayer.forward的编排与博客图 9 的前向时序一一对应例如# Micro batch 0: attention attention_output0, attention_bias0 self.self_attention(...) # 计算 μ-batch0 注意力 fwd_handle0 dist.all_reduce(attention_output0, groupself.TP_group, async_opTrue) # Micro batch 1: attention与 μ-batch0 的 AllReduce 重叠 layernorm_output1 self.input_layernorm(hidden_states1) ... fwd_handle1 dist.all_reduce(attention_output1, groupself.TP_group, async_opTrue)从中可以看到代码对小内核分组重叠的实现μ-batch0 的 attention AllReduce 用 μ-batch1 的 LayerNorm attention 计算掩盖随后才进入 residual/LayerNorm 阶段并wait()。整体而言从源码结构可以推断Domino 在设计上刻意保持张量切分简单、通用只有列并行 Linear 与行并行无通信 Linear 两类模块其重叠调度逻辑全部集中在内核调用层由DominoTransformer/DominoTransformerLayer的显式异步句柄驱动任何 LayerNorm、Dropout、GeLU、RoPE 等算子替换都不会破坏重叠框架与博客第 3 点工作在内核调度器层面的论断互为印证。如何尝试 DeepSpeed-Domino博客在 blogs/deepspeed-domino/README.md 中说明要实际试跑 DeepSpeed-Domino请参考 DeepSpeedExamples 仓库中training/DeepSpeed-Domino目录下的 Domino tutorial示例与教程存放在独立的 DeepSpeedExamples 项目中本仓库并未附带可运行示例因此需要先 clone 对应示例仓库再按其说明执行。本仓库内的官方教程索引 docs/_tutorials/domino.md 亦给出相同指引Domino achieves near-complete communication hiding behind computation for tensor parallel training。从代码层面看DominoTransformer是面向 Megatron 风格 model parallel unitmpu与config编写的使用前需自行准备初始化好的张量并行进程组通过mpu.get_tensor_model_parallel_group()传入并准备好包含hidden_size、num_layers、num_attention_heads、ffn_hidden_size、params_dtype、init_method、add_bias_linear、domino_intra_layer_overlap等字段的配置对象——这一点也决定了尝试 Domino 通常需要结合 Megatron 风格的训练框架来完成端到端集成。使用须知Domino 的通信隐藏建立在 TP 集合AllReduce之上叠加 ZeRO、流水并行等其它并行维度时需先确认相应通信组合与切分策略的兼容性ShardedAttention的注释亦标明当前仅支持 self-attention 与 causal maskdeepspeed/runtime/domino/transformer.py。参考引用与致谢以下引用信息来自 blogs/deepspeed-domino/README.md原文 BibTeX 条目如下论文全文请按条目信息检索 arXiv:2409.15241article{wang2024-deepspeed-domino, title{{Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping}}, author{Guanhua Wang and Chengming Zhang and Zheyu Shen and Ang Li and Olatunji Ruwase}, journal{arXiv preprint arXiv:2409.15241}, year{2024} }该工作是 Microsoft DeepSpeed 团队与 University of Maryland、University of Houston 学术伙伴深入合作的成果。贡献者包括来自 Microsoft DeepSpeed 团队的 Guanhua Wang、Hongwei Chen、Olatunji Ruwase来自 University of Houston 的 Chengming Zhang以及来自 University of Maryland 的 Zheyu Shen 与 Ang Li。本文中的设计图图 3 / 6 / 7 / 8 / 9 / 10与对应分析文字均源自 blogs/deepspeed-domino/images/ 下的原图与 blogs/deepspeed-domino/README.md 的叙述仓库侧实现证据则来自 deepspeed/runtime/domino/async_linear.py 与 deepspeed/runtime/domino/transformer.py供读者在阅读论文或试跑示例时对照参考。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考