Guanhua Wang, Chengming Zhang, Zheyu Shen∗, Ang Li∗, Olatunji Ruwase

Microsoft DeepSpeed

摘要

鉴于生成式AI的流行，大型语言模型（LLM）通常消耗数百或数千个GPU来并行化和加速训练过程。在大规模训练LLM时，通信开销变得更加显著。为了消除分布式LLM训练中的通信开销，我们提出Domino，这是一种通过将单个批次训练分解为更小的独立片段来隐藏通信于计算背后的通用方案。Domino对这些独立片段的训练进行流水线处理，并提供细粒度通信和计算重叠的通用策略。广泛的结果显示，与Megatron-LM相比，Domino在Nvidia DGX-H100 GPU上实现了高达1.3倍的LLM训练加速。

主要贡献

生成式AI的最新进展在各种领域启用了新的应用场景，例如聊天机器人【[53], OpenAI. ChatGPT. https://chatgpt.com/, 2024】、文本生成和摘要【[11], Tom B. Brown et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020】【[52], OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023】、图像和视频内容创建【[62], Aditya Ramesh et al. Zero-shot text-to-image generation. arXiv preprint arXiv:2102.12092, 2021】。这些生成式AI应用基于仔细训练的基础模型，如大型语言模型（LLM）。已建立的LLM是Transformer模型，例如GPT【[11], Tom B. Brown et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020】【[52], OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023】【[60], Alec Radford et al. Language models are unsupervised multitask learners. https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf, 2019】和Llama【[5], Meta AI. Introducing Llama 3.1: Our most capable models to date. https://ai.meta.com/blog/meta-llama-3-1/, 2024】【[6], Meta AI. Introducing Meta Llama 3: The most capable openly available LLM to date. https://ai.meta.com/blog/meta-llama-3/, 2024】【[72], Hugo Touvron et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023】【[73], Hugo Touvron et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023】系列。鉴于LLM模型大小通常从数十亿到数百亿参数，这远远超过了单个GPU的内存和计算限制，因此在数百到数千个GPU上进行分布式模型训练是必要的。

对于LLM Transformer训练，出现了三种突出的范式：数据并行（DP）、张量并行（TP）和流水线并行（PP）。 vanilla数据并行训练指每个GPU维护整个模型参数的完整副本，但使用不同的输入数据进行训练。在每个训练迭代结束时，需要在所有使用的GPU之间同步模型参数。为了缓解DP训练中LLM巨大参数量带来的内存压力，通常需要ZeRO【[61], Samyam Rajbhandari et al. Zero: memory optimizations toward training trillion parameter models. In ACM/IEEE SuperComputing, 2020】或FSDP【[81], Yanli Zhao et al. Pytorch fsdp: Experiences on scaling fully sharded data parallel. Proceedings of the VLDB Endowment, 16, 2023】将模型参数分布到所有GPU上，并仅收集进行计算所需的参数。TP和PP属于模型并行，其中PP【[24], Aaron Harlap et al. Pipedream: Fast and efficient pipeline parallel dnn training. arXiv preprint arXiv:1806.03377, 2018】【[26], Yanping Huang et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In NeurIPS, 2019】将不同的模型层分区到不同的GPU上。TP【[33], Vijay Anand Korthikanti et al. Reducing activation recomputation in large transformer models. In MLSys, 2023】【[67], Mohammad Shoeybi et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019】不是让不同GPU持有不同的模型层，而是将每个层分割到每个GPU上，因此每个GPU持有每个模型层的一部分。

在所有分布式模型训练范式中，张量并行（TP）越来越受欢迎，尤其是在Nvidia GPU上【[21], Introduction to the NVIDIA DGX A100 System. https://docs.nvidia.com/dgx/dgxa100-user-guide/introduction-to-dgxa100.html, 2024】【[22], NVIDIA DGX H100 Datasheet. https://resources.nvidia.com/en-us-dgx-systems/ai-enterprise-dgx, 2023】。TP是单节点多GPU训练的标准实践，鉴于其在高通信带宽（即NVlink【[50], NVIDIA NVLINK. http://www.nvidia.com/object/nvlink.html, 2017】、NVSwitch【[51], NVIDIA NVSWITCH. http://images.nvidia.com/content/pdf/nvswitch-technical-overview.pdf, 2018】）情况下的良好系统效率。然而，由于跨节点网络带宽有限，TP在多节点情况下表现不足。最近，Nvidia正在打破节点间和节点内链路之间的带宽差距。例如，最新的DGX H100【[22], NVIDIA DGX H100 Datasheet. https://resources.nvidia.com/en-us-dgx-systems/ai-enterprise-dgx, 2023】盒子配备了高带宽InfiniBand（IB）链路，聚合带宽为400 GB/s，用于跨节点通信，这与节点内NVSwitch带宽（即DGX-H100上的900GB/s【[22], NVIDIA DGX H100 Datasheet. https://resources.nvidia.com/en-us-dgx-systems/ai-enterprise-dgx, 2023】）处于同一水平。因此，现在是优化并提出一个仅TP解决方案的时候，该解决方案覆盖单节点和多节点场景。

TP的主要开销是其每层全局通信，这位于训练执行的关键路径上。如文献【[67], Mohammad Shoeybi et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019】所述，每个Transformer层需要在正向传播中通信两次，在反向传播中另外两次，使用NCCL集体操作【[30], Sylvain Jeaugey. Optimized inter-GPU collective operations with NCCL 2. https://developer.nvidia.com/nccl, 2017】（§ 2.3）。鉴于这些集体通信发生在执行的关键路径上，很难使用DP【[61], Samyam Rajbhandari et al. Zero: memory optimizations toward training trillion parameter models. In ACM/IEEE SuperComputing, 2020】或PP【[24], Aaron Harlap et al. Pipedream: Fast and efficient pipeline parallel dnn training. arXiv preprint arXiv:1806.03377, 2018】训练过程中一般的通信重叠策略将这些通信隐藏在后续计算后面。先前工作【[54], Suchita Pati et al. T3: Transparent Tracking and Triggering for Fine-grained Overlap of Compute and Collectives. In ACM ASPLOS, 2024】【[55], Suchita Pati et al. Tale of two cs: Computation vs. communication scaling for future transformers on future hardware. In IEEE International Symposium on Workload Characterization (IISWC), 2023】报告这种通信开销可高达端到端训练迭代时间的45%。如我们的测量之一在图1中所示，即使使用最新的DGX-H100节点连接400GB/s IB，通信仍占GPT-3-13B训练迭代时间的17%到43%。此外，当扩展到更多节点时，通信比例将继续增长。为了缓解TP中的这种高通信开销，先前工作【[54], Suchita Pati et al. T3: Transparent Tracking and Triggering for Fine-grained Overlap of Compute and Collectives. In ACM ASPLOS, 2024】【[78], Shibo Wang et al. Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models. In ACM ASPLOS, 2023】提供了GeMM（通用矩阵乘法）与其后续集体调用（例如NCCL【[30], Sylvain Jeaugey. Optimized inter-GPU collective operations with NCCL 2. https://developer.nvidia.com/nccl, 2017】）的内核融合，以实现细粒度计算通信重叠。然而，这种内核融合技术限制了重叠范围，并且不够通用，无法总是将通信隐藏在计算后面。特别是在集体通信比单个GeMM计算花费更长时间的情况下，大部分通信时间仍作为主要训练开销突出。此外，鉴于最新GPU（例如DGX-H100【[22], NVIDIA DGX H100 Datasheet. https://resources.nvidia.com/en-us-dgx-systems/ai-enterprise-dgx, 2023】、DGX-B200【[45], Nvidia. Dgx-b200. https://www.nvidia.com/en-us/data-center/dgx-b200/, 2024】）上的计算变得更快，通信开销在单节点和多节点情况下都更加显著。

为了在TP中提供一种隐藏通信于计算背后的通用方法，我们提出Domino，这是一种通用方法，将Transformer模型训练的数据依赖分解为片段，然后对这些片段的训练进行流水线处理，以重叠通信和计算。除了传统TP仅能在节点内使用外，Domino为单节点多GPU和多节点多GPU情况提供了统一的TP解决方案。与之前的GeMM+NCCL融合解决方案相比，Domino提供了更广泛的计算和通信重叠范围（例如，AllReduce不仅与单个GeMM重叠，还与LayerNorm、DropOut等重叠）。此外，任何内核融合和优化技术都可以轻松与Domino集成作为即插即用替换，以进一步提升整体系统效率。

广泛的基准结果从Nvidia最新的硬件DGX-H100盒子收集，这些盒子通过3200 Gb/s（即400 GB/s）InfiniBand（IB）结构【[46], Nvidia. Introduction to the NVIDIA DGX H100 System. https://docs.nvidia.com/dgx/dgxh100-user-guide/introduction-to-dgxh100.html, 2024】连接。我们使用流行Transformer模型如GPT-3【[11], Tom B. Brown et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020】和Llama-2【[73], Hugo Touvron et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023】进行基准测试。与Nvidia的状态艺术TP实现Megatron-LM相比，Domino在单节点和多节点情况下实现了高达1.3倍的加速。总体而言，Domino为Transformer训练提供了灵活重叠通信与广泛计算内核的通用方法。

我们的关键贡献总结如下：
• 据我们所知，Domino是第一个为仅张量并行训练提供端到端解决方案的工作，在单节点和多节点情况下实现通用通信-计算重叠。
• 与先前工作相比，Domino提供了更灵活和更广泛的计算和通信重叠策略。
• 在DGX-H100盒子上的实验结果显示，与Megatron-LM相比，Domino在GPT和Llama模型训练中实现了高达1.3倍的加速。
• Domino将被开源并作为https://github.com/microsoft/DeepSpeed的一部分发布。

图1: GPT-3-13B 在使用TP的1个DGX-H100节点（8 H100）、2个节点（16 H100）和4个节点（32 H100）上的每个训练迭代的计算和通信比例。

背景知识和动机

分布式训练方案。 在本节中，我们首先描述最广泛使用的分布式Transformer训练方案，即数据并行（DP）、张量并行（TP）和流水线并行（PP）如§ 2.1。然后我们说明为什么TP在这些三种方法中越来越受欢迎如§ 2.2。最后，我们分析TP在单节点多GPU和多节点多GPU情况下的通信开销如§ 2.3。

分布式训练方案概述。 有三种主要的分布式LLM训练范式，即数据并行（DP）、张量并行（TP）和流水线并行（PP）。在vanilla DP中，每个GPU维护模型权重的完整副本并消耗不同的输入数据。在每个训练迭代结束时，所有涉及的GPU执行AllReduce操作来同步模型参数。具体针对Transformer模型，ZeRO【[61], Samyam Rajbhandari et al. Zero: memory optimizations toward training trillion parameter models. In ACM/IEEE SuperComputing, 2020】和FSDP【[81], Yanli Zhao et al. Pytorch fsdp: Experiences on scaling fully sharded data parallel. Proceedings of the VLDB Endowment, 16, 2023】被广泛用于减少设备上的内存压力。在这些完全共享数据并行方案中，整个模型权重通常均匀分布到所有GPU上。当需要在特定层上进行计算时，每个GPU通过在所有GPU之间执行AllGather操作来收集该特定层的完整权重。一旦计算完成，每个GPU释放整个层权重并仅维护最初分配到每个GPU上的权重部分。因此，ZeRO/FSDP可以视为一种内存高效的数据并行方案，它用更多通信换取更少的设备内存占用。

PP和TP都是模型并行技术的代表。PP将一层或一组层分区到一个GPU上，然后从持有第一层的GPU到持有最后一层的GPU进行流水线执行正向传播，然后以反向顺序进行反向传播。与PP相比，TP在正交方向上分区模型，其中每个GPU持有每个模型层的一部分，这样每个GPU可以从第一层计算到最后一层，而不会因PP阶段之间创建的顺序依赖而阻塞【[36], Dmitry Lepikhin et al. Gshard: Scaling giant models with conditional computation and automatic sharding. In ICLR, 2021】。此外，TP似乎具有与ZeRO/FSDP类似的模型分区策略。这里的主要区别是，与ZeRO/FSDP相比，TP在正向或反向计算期间从不收集权重，而是通过AllReduce同步激活或梯度。与DP和PP相比，TP在高带宽通信链路下提供了最高的系统效率或训练吞吐量【[2], PyTorch 2.3. Tensor parallelism introduces more efficient ways to train llms. https://tinyurl.com/3fjvd88t, 2024】【[70], Meta PyTorch team. Tensor parallelism okr in pytorch distributed h2 2024 roadmap. https://drive.google.com/file/d/19cim5wCoxf8A66YZzoMrjardRntpX6gH/view, 2024】【[74], vLLM. Distributed Inference and Serving with TP. https://docs.vllm.ai/en/stable/serving/distributed_serving.html, 2024】。

图2: TP训练中每个Transformer块中的4个AllReduce。两个空白AllReduce框在正向传播中，另外两个灰色AllReduce框在反向传播中。

TP的流行趋势。 张量并行正在为LLM工作负载在Nvidia GPU上获得流行。AI从业者最近见证了TP软件和硬件栈的实质改进。在软件方面，Nvidia持续增强其Megatron-LM【[33], Vijay Anand Korthikanti et al. Reducing activation recomputation in large transformer models. In MLSys, 2023】【[67], Mohammad Shoeybi et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019】软件栈作为状态艺术TP实现。Megatron-LM通过集成更多精细调整和自定义计算内核从库如apex【[44], Nvidia. A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch. https://github.com/NVIDIA/apex, 2024】、cutlass【[17], cutlass: CUDA Templates for Linear Algebra Subroutines. https://github.com/NVIDIA/cutlass, 2024】、cublas【[15], cuBLAS library. https://docs.nvidia.com/cuda/pdf/CUBLAS_Library.pdf, 2024】、cudnn【[16], cuDNN: NVIDIA CUDA Deep Neural Network library. https://developer.nvidia.com/cudnn, 2024】来实现更高的效率。此外，Megatron-LM还涉及新功能来增强整体系统吞吐量，包括选择性激活检查点和序列并行策略【[33], Vijay Anand Korthikanti et al. Reducing activation recomputation in large transformer models. In MLSys, 2023】。

在硬件方面更重要的是，Nvidia正在努力桥接节点内和跨节点链路之间的带宽差距，这对于将TP扩展到跨节点用例至关重要。最新的DGX-H100节点配备了八个Nvidia ConnectX-7 InfiniBand（IB）卡【[42], Nvidia. NVIDIA CONNECTX-7. https://www.nvidia.com/content/dam/en-zz/Solutions/networking/infiniband-adapters/infiniband-connectx7-data-sheet.pdf, 2021】，每个提供400 Gb/s带宽。因此，每个DGX-H100盒子实现了400 GB/s跨节点通信带宽，这与节点内NVLink/NVSwitch带宽相当，如900 GB/s【[22], NVIDIA DGX H100 Datasheet. https://resources.nvidia.com/en-us-dgx-systems/ai-enterprise-dgx, 2023】。此外，Nvidia网络基础设施的进步表明，未来的DGX系统可能集成Nvidia ConnectX-8 IB卡【[48], Nvidia. NVIDIA CONNECTX-8. https://www.nvidia.com/en-us/networking/products/infiniband/quantum-x800/, 2024】，提供高达800 GB/s聚合跨节点带宽，接近节点内NVLink/NVSwitch连接的带宽。

随着软件和硬件的这些进步，PyTorch【[69], Meta PyTorch team. PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration. https://github.com/pytorch/pytorch, 2024】和新兴vLLM【[75], vLLM team. A high-throughput and memory-efficient inference and serving engine for LLMs. https://github.com/vllm-project/vllm, 2024】社区倾向于应用TP用于Transformer训练和推理。例如，PyTorch团队将TP设置为高效LLM训练的一个主要未来方向在其最近发布中【[2], PyTorch 2.3. Tensor parallelism introduces more efficient ways to train llms. https://tinyurl.com/3fjvd88t, 2024】，并将改进TP可扩展性设置为Meta PyTorch团队2024 H2路线图中的关键结果（KR）【[70], Meta PyTorch team. Tensor parallelism okr in pytorch distributed h2 2024 roadmap. https://drive.google.com/file/d/19cim5wCoxf8A66YZzoMrjardRntpX6gH/view, 2024】。同样，在推理方面，vLLM已采用TP作为分布式LLM服务的唯一选项【[74], vLLM. Distributed Inference and Serving with TP. https://docs.vllm.ai/en/stable/serving/distributed_serving.html, 2024】。

鉴于TP的日益流行和IB硬件的最近突破，现在迫切需要为单节点和跨节点场景建立统一的LLM训练TP解决方案。在深入探讨我们的Domino设计之前，我们接下来讨论与TP相关的通信开销。

图3: TP在不同模型类型和模型大小上的每个训练迭代的计算和通信比例，在1到4个DGX-H100节点（8到32 H100 GPU）上。

TP通信开销分析。 我们使用Nvidia的状态艺术TP实现Megatron-LM【[33], Vijay Anand Korthikanti et al. Reducing activation recomputation in large transformer models. In MLSys, 2023】【[47], Nvidia. Megatron-lm release ver 23.08. https://github.com/NVIDIA/Megatron-LM/tree/23.08, 2024】【[67], Mohammad Shoeybi et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019】进行测量。

TP通信与PP或DP解决方案相比具有更多独特特征，主要因为它驻留在每个输入批次训练的关键路径上。将通信隐藏在计算后面不仅是LLM训练的标准实践，而且广泛应用于所有分布式系统环境【[18], Anthony Danalis et al. Transformations to parallel codes for communication-computation overlap. In ACM/IEEE SC, 2005】【[23], Tobias Gysi et al. dcuda: Hardware supported overlap of computation and communication. In ACM/IEEE SC, 2016】【[38], Vladimir Marjanovic et al. Overlapping communication and computation by using a hybrid mpi/smpss approach. In Proceedings of the 24th ACM International Conference on Supercomputing, 2010】【[77], Guanhua Wang et al. Wavelet: Efficient DNN Training with Tick-Tock Scheduling. In MLSys, 2021】【[82], Yihao Zhao et al. Multi-resource interleaving for deep learning training. In ACM SIGCOMM, 2022】。对于使用DP或PP的Transformer训练，通信和计算的重叠相当直接，因为我们可以将通信调度到侧通道，从而绕过关键执行路径。对于像ZeRO【[61], Samyam Rajbhandari et al. Zero: memory optimizations toward training trillion parameter models. In ACM/IEEE SuperComputing, 2020】或FSDP【[81], Yanli Zhao et al. Pytorch fsdp: Experiences on scaling fully sharded data parallel. Proceedings of the VLDB Endowment, 16, 2023】这样的DP方法，预取权重启用与计算的重叠，因为权重本质上没有任何顺序数据依赖。PP通过处理不同输入批次自然地重叠通信和计算。例如，在每个GPU上，PP可以将前一批次的通信与当前批次数据的计算重叠。

如Megatron-LM【[47], Nvidia. Megatron-lm release ver 23.08. https://github.com/NVIDIA/Megatron-LM/tree/23.08, 2024】【[67], Mohammad Shoeybi et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019】所述，每个Transformer块包括一个自注意力层和一个MLP（多层感知器）层。如图2所示，自注意力和MLP层在正向和反向传播中都触发AllReduce操作。因此，每个Transformer块每个训练迭代需要总共4个AllReduce，施加了显著的通信开销。此外，如上所述，传统方法无法将此通信隐藏在计算后面，从而将所有TP通信开销置于执行关键路径上。

我们使用Megatron-LM测量TP训练中的通信开销，跨越GPT-3和Llama-2模型系列的不同模型大小。模型在不同数量的DGX-H100节点上运行，从1到4（即8到32 H100 GPU），取决于模型大小和批次大小。如图3所示，通信时间占端到端迭代时间的22%到47%。这一发现强调，尽管使用了高带宽NVLink/NVSwitch/Infiniband互连，通信开销仍占训练迭代时间的显著部分。这主要是因为每个GPU的计算能力（例如H100）比前代（例如V100、A100）有更显著的增加，从而使通信开销仍然突出。

Domino设计

系统概述。 在本节中，我们描述Domino架构的详细设计。我们首先提供系统架构概述（§ 3.1）。然后我们详细说明如何通用分区计算并将计算内核序列与通信重叠（§ 3.2,§ 3.3, § 3.4）。

整体工作流程。 我们首先描述Domino的整体工作流程。鉴于标准Transformer架构，我们将自注意力层和MLP（多层感知）层抽象为权重张量A_FULL和B_FULL，其中A_FULL代表自注意力层的注意力权重（即Wq,Wk,Wv）但对于MLP是线性权重，B_FULL对于自注意力和MLP层都是线性权重。为了便于说明，我们在正向传播中描述我们的分区策略，因为反向传播只是反向执行顺序。给定层输入数据X，自注意力和MLP层的计算都可以抽象为公式1。

$X \otimes A_{FULL} \otimes B_{FULL} = Y_{FULL}$

如图4所示，TP（例如Megatron-LM）将整个权重张量A_FULL按列分割为集合{A | Ai on GPUi}，权重张量B_FULL按行分割为集合 {B | Bi on GPUi}，用于自注意力和MLP层。在每个GPU获取自己的A和B权重分区后，TP执行X ⊗ A ⊗ B = Y，然后在集合 {Y | Yi on GPUi}上顺序执行AllReduce以恢复Y_FULL，这使得通信开销完全突出。

为了将TP通信隐藏在计算后面，Domino在每个GPU上提供额外的通用张量分区，在原始TP模型分区之上，对输入X进行行分割和对权重B进行列分割。

在高层，Domino通用地将TP的X ⊗ A ⊗ B分解为没有数据依赖的更小计算单元。然后我们将这些独立计算单元与集体通信流水线化，以实现细粒度计算和通信重叠。随着注意力计算被模块化和高度优化如flashattention【[19], Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024】【[20], Tri Dao et al. Flashattention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022】、windowed-attention【[10], Iz Beltagy et al. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020】等的最新趋势，我们保持A不变，不对其进行任何张量分区。因此，我们仅对输入张量X（§ 3.2）和第二组线性权重B（§ 3.3）进行张量切片。我们还提供X和B的混合张量分区策略（§ 3.4）。在这些张量切片后，Domino将X ⊗ A ⊗ B分解为片段并移除数据依赖。然后我们在这些独立片段上启用计算-通信重叠，以减少TP中的通信开销。

在实际模型训练之前，我们基准系统效率以使用网格搜索确定张量分区粒度。这些基准指导我们选择张量计算大小，以确保对计算内核效率的最小影响。为了进一步提升系统效率，Domino还采用最新功能如内核融合、torch.compile【[59], Meta PyTorch. torch.compile. https://pytorch.org/docs/stable/generated/torch.compile.html, 2024】和CUDAGraph【[43], Nvidia. Constructing CUDA Graphs with Dynamic Parameters. https://developer.nvidia.com/blog/constructing-cuda-graphs-with-dynamic-paramete 2022.】【[58], Meta PyTorch. Accelerating PyTorch with CUDA Graphs. https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/, 2021】技术从PyTorch和Nvidia，如§ 4.3所述。

输入上的行分割。 我们首先讨论输入数据的行分割，这指的是§ 3.1中X的张量分区。为了便于图5中的说明，我们简化并假设所有张量都是2维的。然后X可以在行维度或列维度上分割。在现实中，每层的输入张量通常是3D的，如(batch,seq,hidden)。我们将示例2D张量的行/列维度映射到实际3D张量的batch/hidden维度。

如果我们在列维度上将输入分割为N个块，通信量将是vanilla基线的N2倍。如图4所示，假设X的张量形状为(a, b)，A为(b, c)，B为(c, d)。如果我们对X进行列分割并将其分片为[X1, X2...XN]，每个形状为(a, b/N)，我们在适当重塑A后执行X ⊗ A ⊗ B = Y，将得到N2个输出张量，具有原始Y形状(a, d)。为了避免这种通信量爆炸，我们选择对X进行行分割为(a/N, b)如图5所示，这在现实中指的是输入张量在batch维度上的分区(X1, X2)。

注意我们的行分割发生在输入的batch维度上，它在数学上等价于vanilla基线。鉴于行(batch-dim)分割对X主要影响抽象计算X ⊗A，我们详细说明X ⊗A以显示我们的X行分割与基线的等价性。像GeLU()和dropout()这样的元素级操作沿X的batch维度完全独立，我们为简单起见排除它们。然后我们得到简化的X ⊗ A如公式2。

对于MLP，X⊗ A只是X和A之间的GeMM。因此，作为图5中的玩具示例，对X的行分割等价于基线如公式3。

对于自注意力，我们将其抽象为softmax(f (X))g(X)。对于第二部分g(X) = X ∗ Wv，这里的等价证明与公式3相同，因为它只是GeMM操作。对于f (X ) = (X∗Wq)(X∗Wk)T / \sqrt{d} ，其输出维度为(batch, seq, seq)。由于Softmax()在f (X)输出的最后一个维度即sequence-dim上执行，这完全独立于第一个batch维度。由于softmax(f (X))和g(X)都在batch维度上独立，它们的乘积也在batch维度上独立，因此自注意力层的X行分割也等价于没有张量切片的基线。

数据依赖： 由于输入张量的batch维度完全独立，在所有Transformer层中不需要同步。如图5所示，Domino的对输入的行分割实现了层内和层间（即与后续层计算重叠通信）的计算和通信重叠。

权重上的列分割。 与§ 3.2类似的分析，对权重张量B在行维度上分区为N个分区将导致N2倍的通信量爆炸。为了避免这一点，我们在列维度上分割权重张量B，以保持通信量与vanilla基线相同。

如图6所示，对于B，我们将其列分割为N个分区，每个部分输出将具有形状(a, d/N)。在收集所有N个块后，我们在每个X ⊗ A ⊗ B层计算结束时连接这些部分结果（(例如，图6中的Concat(Y3,Y4)）。

现在我们证明对权重B的列分割等价于没有张量分区的基线。由于dropout()发生在我们的连接之后，因为连接输出与基线相同，它可以安全地从我们的证明域中移除。通过排除元素级dropout()操作，(XA) ⊗ B对于自注意力和MLP层都只是GeMM。因此，等价性如公式4所示。

数据依赖： 鉴于对B的列分割，对于自注意力和MLP层，计算输出需要在层执行结束时同步。作为图6中列分割为2的玩具示例，Domino实现了层内计算通信重叠，但在移动到下一个自注意力或MLP层计算之前需要同步（即图6中的Concat(Y3,Y4)）。

图7: Transformer块（即1个自注意力层和1个MLP层）的正向阶段。上图是vanilla TP实现，下图是Domino实现。

混合分割。 对于极大型LLM【[52], OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023】，我们提供对输入X和最后权重张量B的混合模型分割。这种混合解决方案是必要的，因为单独的行分割或列分割会导致狭窄形状张量，这不可能实现良好的计算效率。在对X进行行分割和对B进行列分割后，Domino可以实现超细粒度计算和通信重叠。X ⊗A⊗B的聚合通信大小仍与vanilla基线相同。

数据依赖： 继承自对B的列分割，对于自注意力和MLP层，最终计算输出需要在列方向上同步（即图6中的Concat(Y3,Y4)）但在行方向上是非阻塞的。因此，混合分割只能实现层内计算和通信重叠。

实现

输入上的张量分区。 我们现在讨论实现细节，包括输入数据的行分区策略（§ 4.1）、模型权重的列分区方法（§ 4.2）以及计算内核的进一步优化（§ 4.3）。

输入分区实现。 我们首先说明我们在正向和反向传播中新型输入分区的实现，分别说明。

正向阶段。 用户可以为输入定义分区数p1，之后输入沿batch维度分为p1个分区。一个for循环顺序迭代每个分区的µ-batch。图7描绘了一个简单示例的正向阶段，其中层输入被分割为两个µ-batch（即p1 = 2）。

图8: Transformer块（1个自注意力层和1个MLP层）的反向阶段。上图是vanilla TP实现，下图是Domino实现。

在图7(a)中，为了隐藏自注意力层后的AllReduce通信（即图7(a)中的AllReduce (attn)），我们首先执行µ-batch 0的自注意力，然后异步启动其AllReduce（即图7(b)中的AllReduce(attn0)）以防止GPU在通信上阻塞。随后，我们立即继续µ-batch 1的自注意力。µ-batch 1的自注意力通信（即AllReduce(attn1)）与层归一化、残差和dropout操作重叠。将多个µ-batch的dropout、残差、layerNorm分组不仅启用隐藏图7(b)中正向传播的AllReduce(attn1)，还为图8(b)中反向传播的Allreduce(MLP0)提供适当的重叠空间。

类似地，为了隐藏图7(b)中MLP正向的AllReduce(MLP0)通信，我们在µ-batch 0的MLP计算后异步启动此AllReduce，从而立即执行µ-batch 1的MLP以实现重叠。此外，µ-batch 1的MLP后的AllReduce(MLP1)将与后续Transformer块中的µ-batch 0计算重叠。

反向阶段。 对应的反向主要由torch.autograd()生成。图8显示了一个玩具示例的反向传播，其中输入隐藏状态被分割为两个µ-batch (p1 = 2)。我们仔细组织这些µ-batch中梯度计算的执行，以重叠梯度计算和通信。

在图8中，我们首先采用类似于§ 4.1.1所述的跨µ-batch计算和通信重叠。为了进一步拓宽重叠范围，我们还采用与同一µ-batch内权重梯度计算的重叠通信。例如，图8(b)中的AllReduce(MLP1)部分与它自己的µ-batch 1的grad matmul计算重叠（即从左起的第3个橙色块）。每个grad matmul通常涉及两个独立的内核计算作为输入梯度和权重梯度计算。这种子模块重叠可以通过首先计算MLP层µ-batch 1中的输入梯度（即从左起的第3个橙色块中的第2个grad matmul），然后同时触发其权重梯度计算和输入梯度通信来实现。

然而，精确控制梯度通信行为以与梯度计算重叠是具有挑战性的，因为PyTorch自动生成梯度计算图【[69], Meta PyTorch team. PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration. https://github.com/pytorch/pytorch, 2024】。为了精确控制通信开始/结束时间，我们最初尝试手动实现自定义反向传播导致了较差的吞吐性能，因为触发的内核不如torch.autograd()高效。为了解决这个问题，我们开发了一个无操作模块。该模块在正向阶段接收通信句柄并保留它用于反向阶段。我们的无操作模块与torch.autograd()无缝集成。这种方法允许我们精确控制异步通信的完成时间，而无需复杂的代码修改。

总之，Domino通过我们的输入batch分割实现了高达∼100%通信隐藏在计算后面的能力。

权重上的张量分区。 用户也可以为权重定义分区数p2。随后，初始化p2个线性模块，每个隐藏维度缩放为1/p2。

图6的下半部分显示了一个权重分区的玩具示例，其中权重的分区数为2。具体来说，我们首先执行第一个线性模块（即X ⊗ A ⊗ B1）以生成第一半结果（即Y3）。然后我们触发对第一半结果的异步非阻塞AllReduce。之后，我们立即执行第二半线性模块(X ⊗ A ⊗ B2)。因此，AllReduce(Y3)与X ⊗ A ⊗ B2重叠。在反向中，我们采用类似于§ 4.1.2讨论的子模块重叠策略。

这里的障碍是完全恢复隐藏维度（即图6中的concat(Y3,Y4)）用于后续操作（例如layerNorm、dropout等）。torch.cat()经常分配超过需要的GPU内存【[34], Woosuk Kwon et al. Efficient memory management for large language model serving with pagedattention. In ACM SOSP, 2023】，这可能触发不必要的OOM（内存不足）错误。为了在隐藏维度上实现连接而不使用torch.cat()，我们预分配一个大缓冲区来顺序存储第一半（即Y3）和第二半（即Y4）结果如图6所示。然而，这种方法仍会因非连续内存地址而产生额外的内存拷贝（MemCpy）开销。我们相信这种MemCpy开销可以通过实现自定义内核来缓解或消除，这些内核同时从非连续内存地址读取和写入。鉴于当前这种额外MemCpy的影响最小，我们将其优化推迟到未来工作。

在实践中，Domino通过采用权重上的列分割实现了50%到70%的通信隐藏。虽然这种重叠百分比低于batch输入分割（§ 4.1），但这种方法仍然至关重要，因为单独的batch分割会导致狭窄形状张量，从而阻碍内核计算效率。

通用内核优化。 这里我们讨论使用CUDA-MultiStream和PyTorch原生编译技术的通用内核级优化。

MultiStream。 在将计算分割为更小单元后，每个内核所需的计算与原始TP基线相比显著减少。为了增加GPU利用率同时减少顺序内核启动开销，我们使用多个CUDA流并行执行独立操作。

要获取一个新的CUDA流，可以从CUDA流池中检索它。然而，这种方法生成过多的新流并以轮询方式使用它们，导致频繁流切换的高开销。为了缓解这一点，我们首先在执行前初始化并创建固定数量的全局流。然后，我们使用索引来获取特定流，从而减少与流切换相关的开销。

CudaGraph & Torch.compile。 PyTorch的torch.compile()功能通过即时（JIT）将PyTorch操作编译为优化内核来加速代码执行，从而以最小代码修改实现改进性能【[59], Meta PyTorch. torch.compile. https://pytorch.org/docs/stable/generated/torch.compile.html, 2024】。我们使用torch库的许多操作来构建我们的模块。通过融合不同的操作，我们利用torch.compile()来增强我们的计算效率。

在Domino将张量切片为多个块后，每个块所需的计算显著少于原始基线，导致相邻操作之间的GPU空闲（即气泡时间）。气泡的主要原因是不同操作的计算时间少于PyTorch调度延迟。为了解决这个问题，我们采用CudaGraph【[43], Nvidia. Constructing CUDA Graphs with Dynamic Parameters. https://developer.nvidia.com/blog/constructing-cuda-graphs-with-dynamic-paramete 2022.】【[58], Meta PyTorch. Accelerating PyTorch with CUDA Graphs. https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/, 2021】来消除相邻操作之间的间隙，从而减少整体计算时间。然而，常用的设备随机数生成器（RNG）功能与CudaGraph不兼容。作为变通方法，我们使用固定种子而不是随机数来模拟RNG的行为。

实验环境

模型配置。 评估聚焦于GPT【[11], Tom B. Brown et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020】【[52], OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023】【[60], Alec Radford et al. Language models are unsupervised multitask learners. https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf, 2019】和Llama【[5], Meta AI. Introducing Llama 3.1: Our most capable models to date. https://ai.meta.com/blog/meta-llama-3-1/, 2024】【[6], Meta AI. Introducing Meta Llama 3: The most capable openly available LLM to date. https://ai.meta.com/blog/meta-llama-3/, 2024】【[72], Hugo Touvron et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023】【[73], Hugo Touvron et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023】模型系列。具体来说，使用GPT-3【[11], Tom B. Brown et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020】模型和Llama-2【[73], Hugo Touvron et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023】模型的不同大小进行基准测试。所有模型配置细节如表1所示。对于模型大小计算，我们遵循Nvidia Megatron团队的方程【[41], Deepak Narayanan et al. fficient large-scale language model training on gpu clusters using megatron-lm. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2021】如公式5，其中h指隐藏大小，l是层数。seq_len代表序列长度，vocab是词汇大小。

硬件配置。 实验在Nvidia DGX-H100盒子【[22], NVIDIA DGX H100 Datasheet. https://resources.nvidia.com/en-us-dgx-systems/ai-enterprise-dgx, 2023】上进行，每个有8个H100 GPU。在每个DGX-H100节点内，GPU通过NVLink【[50], NVIDIA NVLINK. http://www.nvidia.com/object/nvlink.html, 2017】和NV-Switch【[51], NVIDIA NVSWITCH. http://images.nvidia.com/content/pdf/nvswitch-technical-overview.pdf, 2018】连接。每个DGX-H100节点配备8个Nvidia InfiniBand ConnectX-7网络卡【[42], Nvidia. NVIDIA CONNECTX-7. https://www.nvidia.com/content/dam/en-zz/Solutions/networking/infiniband-adapters/infiniband-connectx7-data-sheet.pdf, 2021】用于跨节点通信，提供每节点400 GB/s聚合网络带宽。有三种不同的硬件设置：1节点、2节点和4节点，代表单节点和分布式训练环境。所有节点在相同的PyTorch环境中运行，NCCL版本2.18和CUDA版本12.2。

软件配置。 使用PyTorch环境，依赖NCCL【[30], Sylvain Jeaugey. Optimized inter-GPU collective operations with NCCL 2. https://developer.nvidia.com/nccl, 2017】版本2.18和CUDA版本12.2。基准使用Megatron-LM的稳定发布作为基线。代码实现基于Python和PyTorch。

实验结果

评估指标。 与先前关于计算-通信重叠的工作类似【[32], Chenyu Jiang et al. Lancet: Accelerating Mixtureof-Experts Training via Whole Graph ComputationCommunication Overlapping. In MLSys, 2024】【[54], Suchita Pati et al. T3: Transparent Tracking and Triggering for Fine-grained Overlap of Compute and Collectives. In ACM ASPLOS, 2024】【[78], Shibo Wang et al. Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models. In ACM ASPLOS, 2023】，我们主要报告整体训练迭代时间的结果分析。在公式6中，吞吐量或TFLOPs可以推断，因为它与迭代时间测量（即iter_time）成反比。我们相信迭代时间代表更彻底的端到端结果，因为CPU侧执行（例如数据预处理、学习率适应等）也被考虑在内，这可能不包括在纯GPU TFLOPs测量中。由于我们仅使用TP进行模型分区，并且TP域中的每个GPU共享相同的输入，我们的全局批次大小等价于我们的微批次大小。

我们的基线是使用Megatron-LM的稳定发布，有两种不同设置：同步（sync）和异步（async），其中sync（即Megatron-LM(sync)）意味着所有集体操作是阻塞调用，async（即MegatronLM(async)）是启用Megatron-LM中仅反向传播的粗粒度计算和通信重叠功能【[67], Mohammad Shoeybi et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019】。默认情况下，我们将Domino和Megatron-LM (async)与vanilla Megatron-LM (sync)作为基线比较。

值得一提的是，我们的Domino方案在数学上等价于像Nvidia Megatron-LM这样的vanilla TP解决方案（如§ 3.2, 3.3中的证明）。使用固定随机种子和相同学习率调度，我们通过weights & bias工具【[1], Weights & Biases. https://wandb.ai/site, 2024】监控，显示Domino的损失曲线与Megatron-LM基线匹配。为简洁起见，我们在这里排除这些收敛结果。

GPT-3单节点结果。 对于单节点训练，测试了2.7B、6.7B和13B模型大小。总之，与Megatron-LM相比，Domino实现了高达1.3倍的吞吐量加速。此外，在多个情况下，Domino实现了接近最优或甚至超过最优设置。最优解决方案指禁用TP训练中正向和反向传播的所有通信。

一个棘手部分是是否启用CudaGraph。根据我们的实验结果，我们发现如果批次大小小（即训练作业不是计算密集型），启用CudaGraph可以压缩相邻内核之间的气泡/空闲时间，从而改进端到端性能。另一方面，如果训练作业是计算密集型并且相邻内核之间没有太多空闲时间，我们禁用CudaGraph以实现更快的模型训练初始化和更少的设备内存拷贝开销。以GPT-3 13B训练为例，使用序列长度512和微批次大小4，如果我们从cudaGraph关闭模式切换到开启模式，我们注意到显著的训练迭代时间减少（约10-15%），这显示了减少相邻计算内核之间空闲时间的好处。另一方面，如果我们将微批次大小增加到16，启用cudaGraph导致比禁用它长5-10%的迭代时间，这主要是由于CudaGraph引入的额外内存拷贝开销。

总体而言，如图9所示，启用Megatron的粗粒度计算和通信重叠（即MegatronLM (async)）与vanilla megatron基线相比实现了约2-5%的吞吐量增益。与MegatronLM基线相比，Domino在训练批次大时实现了更高的加速增益，在小批次大小情况下性能增益相对较低。

对于图9a所示的GPT-3 2.7B训练，Domino对于序列长度512和1k实现了1.14x到1.26x的加速超过Megatron基线。在图9b的GPT-3 6.7B训练中，由于我们将模型大小从2.7B增加到6.7B，最大的微批次大小与2.7B情况相比减少。然而，我们在6.7B模型中实现了比2.7B和13B情况更高的吞吐量增益。更具体地说，在图9b中，对于序列长度512和1k，我们随着微批次大小增加实现了1.15x到1.3x超过Megatron基线的加速。对于图9c的13B情况，我们有最小的微批次大小用于训练，这导致随着批次大小增加的12%到23%吞吐量加速超过Megatron基线。总之，Domino在各种批次大小和序列长度中通常优于Megatron基线。我们的性能增益随着批次大小增长而增加。

我们还描绘了Domino性能与单节点不同序列长度和批次大小的最优设置的比较。如图10所示，水平基准索引号严格遵循图9中训练设置的相同顺序。与移除MegatronLM中所有通信的最优设置相比，Domino在所有情况下达到了超过90%的最优吞吐量，并且有几个情况甚至超过了最优设置。我们进行了消融研究和性能增益分解。我们发现，对于Domino超过最优设置的情况，额外性能增益主要归因于我们如§ 4.3讨论的内核侧优化。

GPT-3多节点结果。 与单节点结果相比，多节点情况不同，因为跨节点IB带宽仍比节点内NVLink/NVSwitch低2-3倍。因此，单个NCCL集体可能比Domino可以重叠的最大计算内核数更长。

对于2和4个DGX-H100节点的实验，我们测试了6.7B、13B和30B三种不同模型大小，跨越16到32 H100 GPU，使用TP模型分区策略。如图11所示，我们报告了与Megatron基线相比的归一化吞吐量加速。对于序列长度512和1k，我们呈现了适当批次大小从4到64的最佳吞吐量结果。Megatron-LM提供的粗粒度计算和通信重叠（即图11中的Megatron-LM (async)）平均给出约2%-4%的性能增益。

如图11a所示，对于2节点情况（16 H100 GPU），Domino对于6.7B和30B模型在各种序列长度和批次大小上实现了平均约1.2x超过Megatron基线的加速。更有趣的是，对于13B训练，Domino在1k序列长度上实现了高达1.3x超过基线的吞吐量增益。我们相信GPT-3 13B在2个DGX-H100节点上的训练提供了一个甜点，即1) 在我们对输入和权重的行和列分割后，大多数计算内核仍高度高效，2) 跨节点通信可以使用Domino主要与计算重叠。

对于图11b描绘的4节点情况，Domino在GPT-3-6.7B、GPT-3-13B、GPT-3-30B的不同批次大小和序列长度上实现了1.14x到1.2x超过Megatron基线的吞吐量加速。与2节点情况相比性能增益较少的原因是跨节点通信无法与Domino的最大计算内核范围完美重叠。鉴于最新的IB与Nvidia ConnectX-8卡【[48], Nvidia. NVIDIA CONNECTX-8. https://www.nvidia.com/en-us/networking/products/infiniband/quantum-x800/, 2024】可以提供800 GB/s节点间通信带宽，我们为Megatron-LM和Domino进行了800 GB/s跨节点带宽的模拟投影。在模拟中，Domino可以实现高达1.5x超过Megatron基线的加速。在我们的模拟中，我们还注意到Domino在更大规模（例如128、256 GPU）上可能实现高于Megatron基线的更高加速。

Llama-2单节点结果。 对于单节点实验，如图12中8-H100列下的最左8个条所示，Domino对于Llama-7B训练实现了约1.16x加速，对于Llama-13B训练实现了1.1到1.15x加速。较大模型上的较低性能增益是因为我们支持较小的批次大小训练。随着较小的批次，Domino的内核启动开销变得更显著，从而导致较少的吞吐量增益。

与GPT-3的结果相比，Domino超过Megatron-LM的性能增益较少。主要问题是由于Llama-2模型中引入的旋转嵌入特征。这种旋转嵌入在我们输入batch维度分割片段之间构建了额外的数据依赖。为了更好的系统性能，我们将这个旋转嵌入问题留作未来的优化方向。

我们还将Domino和Megatron-LM与最优吞吐量场景（即无通信）比较。如图13所示，8-H100列下的最左4组数据点呈现了Domino和Megatron-LM的吞吐量性能。这里我们将最优吞吐量归一化为1，并计算Domino和Megatron的相应吞吐量。对于8-H100情况中的Domino，我们大多实现了约90+%的最优吞吐量，并且比Megatron好10%，这相当不错。

Llama-2多节点结果。 与单节点训练类似，我们在2和4节点上基准Llama-2 7B和13B训练，使用序列长度512和1k。

在图12中，2节点情况是16-H100列下的中间8个条。与Megatron-LM相比，Domino对于512和1k序列长度实现了约1.15x加速。对于4节点情况，类似结果显示为图12中32-H100列下的最右8个条。Domino在各种模型大小和4节点情况的序列长度上实现了1.08x到1.17x超过Megatron-LM的加速。

当与最优情况（即无通信）比较时，如图13所示，对于2节点和4节点情况，Domino实现了约60-80%的最优吞吐量，并且一致地比Megatron-LM好10-20%。

图9: GPT-3训练迭代时间，使用不同模型大小、序列长度（seq）和微批次（mb）大小，在单个DGX-H100节点（8 H100 GPU）上。

图10: Domino与最优解决方案（无通信）相比的归一化吞吐量加速，在单个DGX-H100（8 H100 GPU）上。基准索引严格遵循图9中水平设置（即seq, mb）的顺序。

图11: GPT-3训练吞吐量加速，使用不同模型大小、序列长度（seq）和微批次（mb）大小，在2（16-H100）和4个DGX-H100节点（32-H100 GPU）上。

图12: Llama-2训练迭代时间，使用不同序列长度和模型大小，在1个DGX-H100节点（8-H100）、2个节点（16-H100）和4个节点（32-H100）情况下。

图13: 当比较Domino、Megatron与最优设置（即无通信）时的归一化加速，用于Llama-2训练，在1个DGX-H100节点（8-H100）、2个节点（16-H100）和4个节点（32-H100）情况下。

结论

我们提出Domino，这是一种用于张量切片和分区以实现计算内核序列与通信集体的细粒度重叠的通用方法。在多个DGX-H100节点上的广泛结果显示，Domino可以实现高达1.3x超过状态艺术张量并行解决方案Megatron-LM的加速。此外，Domino在某些情况下甚至超过了最优性能（即移除Megatron-LM中的所有通信）。随着高通信带宽和每个加速器更快计算的趋势，Domino可能有益于小规模和大规模LLM训练。

对于未来工作，我们计划优化旋转嵌入以进一步改进Llama模型的性能，并探索在更大规模上的扩展，如128或256 GPU。我们还将探索与更多集体库的集成，如RCCL【[7], AMD. ROCm Collective Communication Library (RCCL). https://rocm.docs.amd.com/projects/rccl/en/latest/index.html, 2024】，以支持AMD GPU。

方法细节中的引用汇总

[11]: Language models are few-shot learners. arXiv 2020. 在Introduction中描述为GPT系列模型；在Background中描述为GPT模型；在Evaluation中描述为GPT模型基准。
[52]: Gpt-4 technical report. arXiv 2023. 在Introduction中描述为GPT系列；在Background中描述为GPT模型；在Design中描述为极大型LLM。
[60]: Language models are unsupervised multitask learners. 2019. 在Introduction中描述为GPT系列；在Evaluation中描述为GPT模型。
[5]: Introducing Llama 3.1: Our most capable models to date. 2024. 在Introduction中描述为Llama系列；在Related Work中描述为密集模型。
[6]: Introducing Meta Llama 3: The most capable openly available LLM to date. 2024. 在Introduction中描述为Llama系列；在Evaluation中描述为Llama-3类似架构；在Related Work中描述为密集模型。
[72]: Llama: Open and efficient foundation language models. arXiv 2023. 在Introduction中描述为Llama系列；在Related Work中描述为密集模型。
[73]: Llama 2: Open foundation and fine-tuned chat models. arXiv 2023. 在Introduction中描述为Llama系列；在Evaluation中描述为Llama-2基准；在Related Work中描述为密集模型。
[61]: Zero: memory optimizations toward training trillion parameter models. ACM/IEEE SuperComputing 2020. 在Introduction中描述为缓解DP内存压力；在Background中描述为ZeRO/FSDP。
[81]: Pytorch fsdp: Experiences on scaling fully sharded data parallel. VLDB 2023. 在Introduction中描述为FSDP；在Background中描述为ZeRO/FSDP。
[24]: Pipedream: Fast and efficient pipeline parallel dnn training. arXiv 2018. 在Introduction中描述为PP；在Background中描述为PP。
[26]: Gpipe: Efficient training of giant neural networks using pipeline parallelism. NeurIPS 2019. 在Introduction中描述为PP。
[33]: Reducing activation recomputation in large transformer models. MLSys 2023. 在Introduction中描述为TP；在Background中描述为Megatron-LM；在Related Work中描述为序列并行。
[67]: Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv 2019. 在Introduction中描述为TP和通信；在Background中描述为Megatron-LM和AllReduce；在Evaluation中描述为基线。
[30]: Optimized inter-GPU collective operations with NCCL 2. 2017. 在Introduction中描述为NCCL集体；在Background中描述为NCCL；在Related Work中描述为NCCL库。
[50]: NVIDIA NVLINK. 2017. 在Introduction中描述为高带宽；在Evaluation中描述为节点内连接。
[51]: NVIDIA NVSWITCH. 2018. 在Introduction中描述为高带宽；在Evaluation中描述为节点内连接。
[22]: NVIDIA DGX H100 Datasheet. 2023. 在Introduction中描述为DGX-H100；在Evaluation中描述为硬件。
[45]: Dgx-b200. 2024. 在Introduction中描述为最新GPU。
[46]: Introduction to the NVIDIA DGX H100 System. 2024. 在Introduction中描述为IB结构。
[54]: T3: Transparent Tracking and Triggering for Fine-grained Overlap of Compute and Collectives. ACM ASPLOS 2024. 在Introduction中描述为内核融合；在Related Work中描述为内核融合。
[78]: Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models. ACM ASPLOS 2023. 在Introduction中描述为内核融合；在Related Work中描述为内核融合。
[55]: Tale of two cs: Computation vs. communication scaling for future transformers on future hardware. IEEE IISWC 2023. 在Introduction中描述为通信开销。
[21]: Introduction to the NVIDIA DGX A100 System. 2024. 在Introduction中描述为Nvidia GPU；在Related Work中描述为Nvidia GPU。
[2]: PyTorch 2.3. Tensor parallelism introduces more efficient ways to train llms. 2024. 在Background中描述为TP效率；在Related Work中描述为PyTorch方向。
[70]: Tensor parallelism okr in pytorch distributed h2 2024 roadmap. 2024. 在Background中描述为TP可扩展性；在Related Work中描述为PyTorch路线图。
[74]: Distributed Inference and Serving with TP. 2024. 在Background中描述为vLLM TP；在Related Work中描述为vLLM TP。
[47]: Megatron-lm release ver 23.08. 2024. 在Background中描述为Megatron-LM。
[18]: Transformations to parallel codes for communication-computation overlap. ACM/IEEE SC 2005. 在Background中描述为隐藏通信。
[23]: dcuda: Hardware supported overlap of computation and communication. ACM/IEEE SC 2016. 在Background中描述为隐藏通信。
[38]: Overlapping communication and computation by using a hybrid mpi/smpss approach. ACM ICS 2010. 在Background中描述为隐藏通信。
[77]: Wavelet: Efficient DNN Training with Tick-Tock Scheduling. MLSys 2021. 在Background中描述为隐藏通信；在Related Work中描述为轮询。
[82]: Multi-resource interleaving for deep learning training. ACM SIGCOMM 2022. 在Background中描述为隐藏通信。
[36]: Gshard: Scaling giant models with conditional computation and automatic sharding. ICLR 2021. 在Background中描述为PP阻塞。
[44]: A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch. 2024. 在Background中描述为apex库。
[17]: cutlass: CUDA Templates for Linear Algebra Subroutines. 2024. 在Background中描述为cutlass库；在Related Work中描述为优化内核。
[15]: cuBLAS library. 2024. 在Background中描述为cublas库；在Related Work中描述为优化内核。
[16]: cuDNN: NVIDIA CUDA Deep Neural Network library. 2024. 在Background中描述为cudnn库；在Related Work中描述为优化内核。
[42]: NVIDIA CONNECTX-7. 2021. 在Background中描述为IB卡；在Evaluation中描述为网络卡。
[48]: NVIDIA CONNECTX-8. 2024. 在Background中描述为未来IB卡；在Evaluation中描述为模拟。
[69]: PyTorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration. 2024. 在Background中描述为PyTorch；在Implementation中描述为autograd。
[75]: A high-throughput and memory-efficient inference and serving engine for LLMs. 2024. 在Background中描述为vLLM。
[19]: Flashattention-2: Faster attention with better parallelism and work partitioning. ICLR 2024. 在Design中描述为注意力优化。
[20]: Flashattention: Fast and memory-efficient exact attention with IO-awareness. NeurIPS 2022. 在Design中描述为注意力优化。
[10]: Longformer: The long-document transformer. arXiv 2020. 在Design中描述为windowed-attention。
[59]: torch.compile. 2024. 在Design中描述为torch.compile；在Implementation中描述为优化。
[43]: Constructing CUDA Graphs with Dynamic Parameters. 2022. 在Design中描述为CUDAGraph；在Implementation中描述为CudaGraph。
[58]: Accelerating PyTorch with CUDA Graphs. 2021. 在Design中描述为CUDAGraph；在Implementation中描述为CudaGraph。
[68]: Roformer: Enhanced transformer with rotary position embedding. arXiv 2021. 在Evaluation中描述为RoPE。
[80]: Root mean square layer normalization. NeurIPS 2019. 在Evaluation中描述为RMSNorm。
[66]: Glu variants improve transformer. arXiv 2020. 在Evaluation中描述为SwiGLU。
[1]: Weights & Biases. 2024. 在Evaluation中描述为监控工具。
[41]: fficient large-scale language model training on gpu clusters using megatron-lm. SC 2021. 在Evaluation中描述为模型大小方程。
[32]: Lancet: Accelerating Mixtureof-Experts Training via Whole Graph ComputationCommunication Overlapping. MLSys 2024. 在Evaluation中描述为重叠；在Related Work中描述为MoE重叠。
[34]: Efficient memory management for large language model serving with pagedattention. ACM SOSP 2023. 在Implementation中描述为torch.cat问题。
[13]: Centauri: Enabling Efficient Scheduling for CommunicationComputation Overlap in Large Model Training via Communication Partitioning. ACM ASPLOS 2024. 在Related Work中描述为调度。
[83]: Alpa: Automating inter- and intra-operator parallelism for distributed deep learning. USENIX OSDI 2022. 在Related Work中描述为编译器优化。
[25]: TicTac: Accelerating Distributed Deep Learning with Communication Scheduling. MLSys 2019. 在Related Work中描述为PS调度。
[37]: Scaling distributed machine learning with the parameter server. USENIX OSDI 2014. 在Related Work中描述为PS架构。
[8]: Pathways: Asynchronous distributed dataflow for ml. MLSys 2022. 在Related Work中描述为all-worker。
[39]: Mpi: A message passing interface. ACM/IEEE SC 1993. 在Related Work中描述为MPI。
[35]: Breadth-First Pipeline Parallelism. MLSys 2023. 在Related Work中描述为优化PP。
[27]: Introduction to InfiniBand. 2007. 在Related Work中描述为InfiniBand。
[12]: Flux: Fast software-based communication overlap on gpus through kernel fusion. arXiv 2024. 在Related Work中描述为内核融合。
[29]: Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads. ACM ASPLOS 2022. 在Related Work中描述为通用融合。
[79]: Mgg: Accelerating graph neural networks with fine-grained intra-kernel communication-computation pipelining on multi-gpu platforms. USENIX OSDI 2023. 在Related Work中描述为GNN融合。
[49]: NVIDIA OpenSHMEM Library (NVSHMEM). 2024. 在Related Work中描述为NVSHMEM。
[57]: Optimizing distributed ml communication with fused computation-collective operations. arXiv 2023. 在Related Work中描述为AMD融合。
[4]: Collective communications library with various primitives for multi-machine training. 2024. 在Related Work中描述为Gloo。
[76]: Blink: Fast and Generic Collectives for Distributed ML. MLSys 2020. 在Related Work中描述为Blink。
[64]: Horovod: fast and easy distributed deep learning in TensorFlow. arXiv 2018. 在Related Work中描述为Horovod。
[28]: InfiniBand In-Network Computing With NVIDIA SHARP. 2021. 在Related Work中描述为SHARP。
[31]: NCCL Collnet, CollnetDirect and CollnetChain. 2024. 在Related Work中描述为CollNet。
[63]: Enabling compute-communication overlap in distributed deep learning training platforms. ACM/IEEE ISCA 2021. 在Related Work中描述为ACE。
[56]: A generic communication scheduler for distributed dnn training acceleration. ACM SOSP 2019. 在Related Work中描述为ByteScheduler。
[9]: Message Passing Interface. 2018. 在Related Work中描述为MPI。
[71]: Optimization of collective communication operations in mpich. IJHPCA 2005. 在Related Work中描述为MPI优化。
[14]: MSCCLang: Microsoft Collective Communication Language. ACM ASPLOS 2023. 在Related Work中描述为MSCCL。
[40]: MSCCL++: A GPU-driven communication stack for scalable AI applications. 2024. 在Related Work中描述为MSCCL++。
[65]: TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches. USENIX NSDI 2023. 在Related Work中描述为TACCL。
[7]: ROCm Collective Communication Library (RCCL). 2024. 在Related Work中描述为RCCL；在Conclusion中描述为AMD支持。
[3]: Phi-3 technical report: A highly capable language model locally on your phone. arXiv 2024. 在Related Work中描述为密集模型。
[53]: ChatGPT. 2024. 在Introduction中描述为聊天机器人。
[62]: Zero-shot text-to-image generation. arXiv 2021. 在Introduction中描述为图像生成。

Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping

摘要