图解Megatron TP中的计算通信overlap

主要观点总结

本文探索了Megatron中实现计算通信overlap的方法，具体涉及Megatron的dp、tp和pp部分，特别是tp部分（即megatron sp-tp）。文章介绍了在tp中各个步骤的计算和通信流程，以及如何通过p2p ring exchange、pipeline chunk等方法实现计算和通信的串行overlap。此外，文章还介绍了如何通过设置计算流和通信流实现并行overlap，即bulk overlap。最后，文章总结了本文的主要内容和参考资源。

关键观点总结

关键观点1: Megatron中计算通信overlap的探索

介绍Megatron中计算通信overlap的重要性和背景。

关键观点2: Megatron的dp、tp和pp部分简介

概述Megatron中这三个部分的基本功能和在计算通信中的作用。

关键观点3: tp中的计算和通信流程

详细描述tp中各个步骤的计算和通信流程，包括all-gather、reduce-scatter等。

关键观点4: 串行overlap的实现方法

介绍如何通过p2p ring exchange和pipeline chunk等方法实现计算和通信的串行overlap。

关键观点5: 并行overlap的实现方法

介绍如何通过设置计算流和通信流实现并行overlap，即bulk overlap。

关键观点6: 总结

概括文章的主要内容和结论，以及参考资源。

文章预览

这篇文章想来探索Megatron中实现计算通信overlap的方法。 具体来说，Megatron的dp、tp和pp部分，都有可以做overlap的地方，本文探索的是tp部分（更准确地说是megatron sp-tp）。做这个探索的主要目的是：了解在哪些位置有做overlap的潜能，以及当前一些可行的实现思路。 最后，特别感谢overlap大师，megatron特级学者，大众点评美食优惠券killer：https://www.zhihu.com/people/yu-huo-er-wang 为本文提供的各类参考资料。 一、TP中哪些地方做了overlap 我们说的tp，是指“开启megatron sp做了activation显存优化”的tp，下图绘制了在megatron sp中单卡上Attn + MLP的运作流程（更多详细解读请参见 这里 ） 由此我们知道，在megatron sp中，tp部分的通讯被拆成若干个 all-gather 和 reduce-scatter ，在下文中我们会用 AG 和 RS 来简称它。 现在我们对tp中的fwd和bwd过程再做一个重新绘制，更清晰地展示通 ………………………………