为什么要重新设计一个像OneFlow这样的分布式深度学习框架？
一个显而易见的出发点是，我们看到了原有的主流深度学习框架的本质不足。尤其在抽象层面和API层面，它们的设计有种种不足，导致开发者在使用时造成极大不便，尽管他们正在试图解决一些缺陷，但有些重要问题依然被忽视了。
为此，我们将推出三篇系列文章，详细论述原有主流深度学习框架的运行时系统的三大“诅咒”，第1篇《资源依赖的“诅咒”》已经发布，此为第2篇内容。本文将探讨数据搬运放进计算图的关键问题，以及使用回调函数的缺陷，最后，本文将介绍OneFlow的数据搬运是“一等公民”的理念。
撰文 | 袁进辉
1
基于计算图的依赖管理
在上一篇《资源依赖的“诅咒”》中，我们讨论了深度学习框架中算子间的三种典型的依赖关系，数据依赖、控制依赖和资源依赖。
妥善而高效的管理这些依赖关系对深度学习框架至关重要，就像常见的大数据系统一样，深度学习框架的从业者发现计算图（computation graph) 是一种表达和管理复杂依赖的有效工具。
理论上，如果把所有算子和所有依赖都表示在计算图之中，那么，只需要图执行器（graph executor）这一种机制就可以很好的管理好算子的执行顺序。
看上去，只需要管好计算图就万事大吉了，然而，实现上并非如此。
数据搬运是否放进计算图？
问题的关键是，计算图是不是只描述“计算”？是不是只有在CPU或GPU上执行的计算才被当成算子？
不幸的是，原有的深度学习框架都是这么认为的。
图1：单卡训练下，数据搬运非常简单，所有的计算算子都在GPU上一鼓作气地执行
如图1 所示，在非分布式场景，通常只要把输入数据拷贝到GPU上，所有的计算操作都会在一个GPU内部执行，计算完成后只需要把结果从GPU上拷贝走就可以了。
在这种情景下，数据搬运非常简单（从磁盘加载，预处理并拷贝到GPU)，不会带来太多的烦恼。正如现有深度学习框架那样，计算图把所有“计算”类型的算子都囊括进去并管理起来，基于计算图的执行引擎很优雅地解决好大部分问题。计算图只管“计算”，不管数据搬运的事，数据搬运都是其他功能模块去管。
但是，在分布式场景中，情况发生了变化。
在分布式训练中，多个GPU通过高速互联设备连接，无论用什么并行方式（数据并行、模型并行还是流水并行），每处理一个批次的数据，GPU之间就需要进行好多次数据搬运（跨机器或者跨设备）。
图2：一个三层神经网络在4张GPU上分布式执行的逻辑图和物理图图2 展示了一个简单的分布式训练深度学习模型的例子。这是一个3层神经网络，包括3个算子，即f1,f2,f3，假设该任务申请到了4张GPU卡，即d1,d2,d3,d4。f1和f2被放置到d1和d2这两张GPU上，通过数据并行的方式执行，而f3被放置到d3和d4这两张GPU上，通过模型并行的方式执行。那么，无论是手工还是编译器自动生成的物理图（或执行计划），就像图2的下半部分展示的那样，需要在计算算子之间插入一些数据搬运和数据路由的操作，譬如，在前向计算时，从f2到f3 的数据并行切换到模型并行处，需要插入一个all-gather的集群通信原语g。而反向计算时，从b3 (f3的反向）到b2（f2的反向）的模型并行到数据并行切换时，需要插入一个reduce-scatter 的集群通信原语s。因为f1和f2采用数据并行，在反向时，也分别需要插入all-reduce的集群通信原语r1和r2。可以想见，从逻辑图生成物理图的过程比较复杂。与原有深度学习框架不同的是，OneFlow是通过编译器自动完成从逻辑图到物理图映射，而其它框架还需要手工编程插入这些数据搬运和数据路由的操作。有时，数据搬运的时间和计算的时间差不多在一个数量级上（例如，10GB带宽下搬运100MB的数据需要100毫秒），如果不能通过流水线来掩盖数据搬运的时间，那么横向扩展性就会严重受限，即使加进来更多的GPU，带来的加速比也会远远低于理想的线性加速比。(不过，数据搬运到底何等重要，如何通过编译器自动实现从逻辑图到物理图的转换，如何提高整个分布式系统的计算效率都不是本文计划深入讨论的内容，感兴趣的朋友请查阅OneFlow团队分享的其它相关技术资料。这篇文章要探讨的问题是，在已知数据搬运至关重要的情况下，如何更优雅的设计和实现分布式深度学习系统。)总结一下：尽管在分布式场景下，数据搬运的操作非常多，而且消耗的时间和计算时间相差无几，原有深度学习框架的计算图只把在单GPU内部执行的计算逻辑当成算子，并通过图执行器去调度和管理“计算”的逻辑。也就是，尽管数据搬运和计算之间也有数据的生产和消费关系，但数据搬运没有被当成算子放在计算图里面的，那么数据搬运和计算之间的依赖关系是怎么管理的？答案是，通过另外一套机制去管理，也就是整个系统中存在两种依赖管理机制。一方面，这会使得整个系统的概念复杂化，另一方面，在实现上也带来了相当的复杂性，特别是在两种机制交界和交互的地方。熟悉 TensorFlow 论文和代码的读者应该见过下面这张图：图3：TensorFlow 通过在设备的边界插入一对 send 和 recv op来完成跨设备数据搬运如图3 所示，当需要在设备之间搬运数据时，TensorFlow 会在生产者和消费者两侧分别插入一个 send 和 recv 算子，由这一对算子来完成跨设备的数据搬运。你可能马上会产生一个疑问：TensorFlow 明明也用算子 op 来表示数据搬运了，为什么还说现有的框架有两套依赖管理机制？回答是，TensorFlow 只是在逻辑层面这么做了，在物理层面并没有做彻底。设想一下，如果这两个设备A和B在一台机器内部，且支持peer to peer传输，那么只需要一个memCopyAsync就可以把数据从A搬到B。如果两个设备A和B在一台机器内部，且不支持peer to peer传输，那么，在物理层面，首先需要把数据从A上拷贝到主存（这个操作记为m1)，再从主存拷贝到设备B（这个操作记为m2），实际上m1和m2 并不是以算子的形式被TensorFlow放在计算图中的，m1和m2之间的依赖关系也没有在计算图层面表达，他们的触发也不是graph executor管理的。在具体实现中，m1和m2是以回调函数的链条(chain)的方式建立起联系的，也就是，m2作为输入给m1的回调函数，当m1执行完毕时调用m2。如果两个设备A和B不在同一台机器上，那么在send和recv op背后的数据搬运链条就更长，可能包含从设备A搬运到设备A所在机器的主存（m1)，再从设备A所在机器的主存搬运到设备B所在机器的主存（m2)，再从设备B所在机器的主存搬运到设备B（m3)。可以想见，实际执行过程是由一系列回调函数构成的链条来执行的。在TensorFlow中，回调函数构成的链条就是另一套依赖管理机制，这种机制和计算图的依赖管理互相协作完成了分布式训练。搞两套依赖管理系统的确不够“简单”，这是一处不好。另外一处不好就是，一般回调函数处理“单一依赖”比较自然，不方便表示和管理“多依赖”。譬如O1和O2两个函数，如果O2依赖于O1，那么通过把O2当成O1 里执行的回调函数来处理是比较方便的。假如有O1,O2,O3三个函数，O3 依赖于O1和O2，也就是只有当O1和O2都执行完毕时，O3才能开始执行。如果我们确定的知道O1和O2哪个会先执行，哪个后执行，那用回调函数仍然好解决，譬如O1 -->O2 --> O3，可以构造一个回调函数的链条来执行。但如果O1和O2没有依赖关系，也不确定哪一个会先执行完，编程时就不知道O3应该放在O1的回调函数里，还是放在O2的回调函数里。这种情况，必须引入一个计数器来记录O1和O2是否都执行完毕了，当计数器显示这俩函数都执行完毕之后，O3才可以执行。这种借助计数器的办法，实质上和计算图的原理几乎是一样的。而实际的系统中，“多依赖”是很普遍的，譬如数据依赖，额外添加的控制依赖，甚至是资源共享造成的依赖，而使用回调函数来表达这种”多依赖关系“并不方便。如果同一个op的”多个依赖关系“中的一部分用回调函数表示，一部分用计算图来表示，情况就会变得更加复杂，下面让我们看一个例子。正如上文讨论的那样，TensorFlow的send和recv op仅仅描述了逻辑关系，物理上的数据搬运操作被放在回调函数里，通过在必要的地方插入回调函数来体现操作之间的依赖关系。看图4中的例子，O2依赖于O1, O2 被封装在一个回调函数中，我们期待它在O1完成后被调用。然而，如果O2有其他依赖关系是被调度器管理的，比如其它算子的输出数据或框架插入的控制依赖关系，那么O1的完成并不一定足以解决O2的所有依赖关系。为了正确和安全地调度O2，回调函数应该通知调度器O1已经完成，以便调度器更新O1的所有依赖关系。如果调度器返回O2所有的依赖关系都已解决，O2可以立即被调度。否则，O2将被插入到一个等待列表中，当其它依赖关系被解决时，才能被调度。实际情况是，现有的深度学习框架还没有为这种”多依赖“的回调函数提供支持。当一些算子之间的依赖关系并没有全部被调度器跟踪和管理（例如部分在回调函数里，部分在调度器里时），就只能通过图4类似的思路才能解决。这就要求深度学习框架把调度器的模块暴露给外界，因为有时候回调函数是用户手写的，而不是框架自动插入的（譬如Megatron-LM，DeepSpeed 基于PyTorch所做的二次开发中经常有这种情况），框架要支持用户的代码和调度交互，这在软件工程上是不好的实践。当然，因为现有深度学习框架没有考虑过这种需求，也并没有把底层调度器暴露出来，这也就意味着图4所述的思路不能轻易地实现。综上，我们希望无论是数据搬运还是计算，所有的依赖关系都通过一种机制来管理和调度。OneFlow常说的数据搬运为“一等公民”就是这样一种实现，OneFlow的编译器不仅像传统深度学习框架那样仅把计算类型的操作当作算子，而是也显式的把所有数据搬运操作当作算子表示在计算图里，算子之间所有的依赖关系也都表示在计算图里。换句话说，在传统框架里，所有由系统插入的、硬编码的、在运行时才会触发的回调函数都被提取出来，并转换成计算图里的普通算子。这也可以理解为，回调函数的链条被展开并静态表示为计算图。这样做的好处是，整个系统只有一套”依赖管理“机制，那就是计算图（当然在OneFlow的运行时体现为actor之间的生产和消费关系）。抽象和实现都会更简洁，避免了回调陷阱（callback hell)，同时也可能为运行效率带来好处（编译器可以分析和优化算子之间的执行顺序，而回调函数描述的执行顺序是固定的），也自然地支持了一些高级功能（譬如流控）。
点击“阅读原文”，欢迎下载体验OneFlow新一代开源深度学习框架