企业官网建设流程全解析

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 团队开源的Domino（Project Domino）通信隐藏张量并行训练引擎展开。Domino 打破传统张量并行（Tensor Parallelism，TP）中"计算-通信"的串行数据依赖，通过将一次 batch 训练迭代拆分为相互独立的子任务并流水化执行，在单节点与多节点场景下实现近乎完全的计算通信隐藏。读完本文，你将掌握 TP 通信开销的成因、Domino 的三种张量切分策略（按行切输入、按列切权重、二维切分）及其数据依赖分析、前向/反向中的 NCCL 调用触发时序，以及在当前 DeepSpeed 仓库中对应的模块实现与后续上手路径。

目录

1. 背景：为什么需要统一的单机/多节点 TP 方案
2. Domino 亮点一览
3. TP 通信开销的来源与数据依赖分析
4. Domino 核心设计：三种张量切分策略
5. 实现与优化：前向、反向与通用内核优化
6. 仓库中的 Domino 源码实现印证
7. 如何尝试 DeepSpeed-Domino
8. 参考引用与致谢

背景：为什么需要统一的单机/多节点 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 环境下的测试报告）：
  1. 针对不同模型规模、序列长度与 batch 大小，在 GPT 与 Llama 模型系列上，Domino 相对 Megatron-LM 的端到端训练迭代吞吐分别最高提升约1.3x与1.2x；
  2. 在若干场景下 Domino 能达到接近最优的训练吞吐（"最优"指假设禁用 TP 通信集合时所能达到的吞吐）。

需要强调的是，上述数字来自博客自述的特定软硬件环境（DGX-H100 + 400GB/s IB），属于该项目的实验性报告结果，并不能代表在所有硬件、模型与软件栈下的普适结论；其详细实验细节请参阅博客引用的 arXiv 论文 2409.15241。

TP 通信开销的来源与数据依赖分析

每个 Transformer block 的 4 次 AllReduce

按 Megatron-LM 论文的叙述，TP 下每个 Transformer block（1 个 Self-Attention 层 + 1 个 MLP 层）在前向与反向上各执行 2 次 AllReduce，共计4 次 AllReduce。对一个由 $N$ 个 Transformer block 堆叠而成的 LLM，TP 训练每次迭代需要执行 $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）的计算可抽象为 $XAB=Y$：

• 对 Attention 层，$X$ 为输入，$A$ 为注意力计算（如 multi-head attention），$B$ 为线性层；
• 对 MLP 层，$A$、$B$ 均为线性层；
• 计算完成后需要对 $Y$ 做一次 AllReduce。

由于 $Y$ 上的计算（$XAB=Y$）与通信（AllReduce($Y$)）之间存在串行数据依赖，AllReduce($Y$) 会完全"裸露"出来、无法与计算重叠，这正是 TP 在通信带宽受限场景下效率不高的根本原因。近期的 GeMM+NCCL 内核融合工作受限于计算-通信重叠范围过窄（通常只能与单个 GeMM 重叠），也无法完全隐藏通信。

Domino 核心设计：三种张量切分策略

针对上述 $XAB=Y$ 的串行依赖，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%。不过仅有按行切分会把张量切得越来越"窄"，损害内核计算效率，因此还需要下述按列切分方案配合。

按列切分权重 B：50%～70% 的通信隐藏

Domino 沿列维度（最后一维）切分权重矩阵 $B$：

数据依赖分析：权重 $B$ 的按列切分在层内无数据依赖，但在层间存在数据依赖，因此只能实现层内（AllReduce($Y_1$) 与 $X_2AB$）的计算-通信重叠。实践上，纯按列切分权重约可隐藏 50%～70% 的通信；该方案之所以必要，是因为它避免了窄张量对内核效率的损害。

2D 切分（同时切 X 与 B）：面向超大模型的折衷

针对极大 LLM，Domino 同时在行维度切输入 $X$、在列维度切权重 $B$。该策略有利于既要求低显存占用、又希望最小化通信开销的模型训练场景。由于按列切分权重会在每个 Transformer 层末尾引入一次同步（AllReduce），2D 方案同样只能实现层内的计算-通信重叠。

Domino 相比 GeMM+NCCL 融合的优势

综合而言，Domino 相比传统融合方案的独特价值可归纳为：

1. 重叠范围更广：Domino 打破了数据依赖，可让一个 AllReduce 不仅与单个 GeMM 重叠，还能把重叠范围扩展到多个 GeMM、LayerNorm、Dropout 等一串计算内核序列；
2. 接近最优的系统吞吐：通信被近完全隐藏，因此在特定场景下可逼近"禁用 TP 通信"的理想吞吐；
3. 内核层面的通用性：Domino 工作在内核调度器层面，任何内核优化或新内核都能无缝集成进 Domino 框架；
4. 切分方案简单通用：便于用户在遇到溢出、权重/梯度异常时进行端到端的正确性调试。

实现与优化：前向、反向与通用内核优化

博客以按行切分输入的实现为主线介绍（其余细节见论文）。以下前向/反向描述均以"把一个 batch 切成两个 $\mu$-batch（$\mu$-batch0 与 $\mu$-batch1）"为例。

前向：多个计算片段轮转重叠 AllReduce

在 图：Transformer block（1 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-operation（no-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 与 NoOper

deepspeed/runtime/domino/transformer.py 中：

class DominoUtil: BATCH_0 = "BATCH0" BATCH_1 = "BATCH1" HANDLE_DIC = {"BATCH0": None, "BATCH1": None}

DominoUtil以两个键BATCH0/BATCH1对应两个 $\mu$-batch，HANDLE_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, None

NoOper.forward原样透传输入（不产生额外计算），symbolic定义为返回输入自身以便配合 torch 图模式；真正的等待动作被推迟到backward中——取出前向时存入HANDLE_DIC的异步通信句柄并执行handle.wait()。由于 PyTorch 的 autograd 只在反向需要该张量梯度时才触发 backward，这恰好实现了"精确控制 NCCL 完成/等待时机"的效果，且无需改写任何自定义反向图，与博客描述完全吻合。

异步 AllReduce 的发起：DominoAsyncColumnParallelLinear

deepspeed/runtime/domino/async_linear.py 定义了DominoAsyncColumnParallelLinearImpl（一个torch.autograd.Function）：

@staticmethod def backward(ctx, grad_output): inp, weight, bias = ctx.saved_tensors ... grad_input = torch.matmul(grad_output, weight) handle = dist.all_reduce(grad_input, group=TP_group, async_op=True) ctx.handle_dic[ctx.h_id] = handle ...

反向中计算出的grad_input并不立即等待 AllReduce 完成，而是以async_op=True发起 AllReduce 并把返回句柄写入ctx.handle_dic[ctx.h_id]——即前向中由ShardedAttention/DominoTransformerLayer传入的DominoUtil.HANDLE_DIC。配合NoOper在后续依赖点wait()，构成了"异步发起 + 延迟等待"的重叠闭环。同文件中还提供了RowParallelLinearNoComm（不触发通信的按行并行线性层）用于dense/linear_fc2等输出侧投影。

μ-batch 流水：DominoTransformer 与 DominoTransformerLayer

deepspeed/runtime/domino/transformer.py 中的DominoTransformer在前向入口通过torch.chunk(hidden_states, chunks=2, dim=1)把 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_forward

• inter_layer_overlap_forward对应按行切分输入 X 的层间 + 层内重叠流水（与博客图 6 的 100% 隐藏能力对应），它把 attention、residual、MLP 等子步骤按两个 μ-batch 交错编排，并通过fwd_handle0/fwd_handle1两个异步 AllReduce 句柄实现跨层轮转等待；
• intra_layer_overlap_forward与DominoTransformerLayer.forward（deepspeed/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, group=self.TP_group, async_op=True) # Micro batch 1: attention（与 μ-batch0 的 AllReduce 重叠） layernorm_output1 = self.input_layernorm(hidden_states1) ... fwd_handle1 = dist.all_reduce(attention_output1, group=self.TP_group, async_op=True)

从中可以看到代码对"小内核分组重叠"的实现：μ-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 unit（mpu）与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 mask（deepspeed/runtime/domino/transformer.py）。

参考引用与致谢

以下引用信息来自 blogs/deepspeed-domino/README.md，原文 BibTeX 条目如下（论文全文请按条目信息检索 arXiv:2409.15241）：

@article{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