为什么要重新设计一个像OneFlow这样的分布式深度学习框架?
一个显而易见的出发点是,我们看到了原有的主流深度学习框架的本质不足。尤其在抽象层面和API层面,它们的设计有种种不足,导致开发者在使用时造成极大不便,尽管他们正在试图解决一些缺陷,但有些重要问题依然被忽视了。
为此,我们将推出数篇系列文章,详细论述原有主流深度学习框架的运行时系统的“诅咒”,此为第3篇内容(点击查看第1、2篇)。本文从分配线程池这一关键问题出发,介绍了计算图调度机制,重点探讨原有框架使用动态调度的缺陷以及设置最优线程数的烦恼,并阐述了OneFlow实现的更优雅的方案。
撰文 | 袁进辉
不同的硬件资源应该使用不同的stream;
有依赖关系的算子最好发射到同一个stream,无依赖关系的算子最好使用不同的stream;
同一个硬件资源创建多个stream利弊各半,同一个硬件资源只创建一个stream利弊各半,那怎么办呢?
使用单个线程同时来做计算图调度和算子执行,这个线程访问和修改计算图的状态,但算子可能被派发到多个不同的设备上去。需要注意的是,每次向一个设备派发任务时,需要提前为这个设备初始化设备上下文,当向另一个设备派发任务时,就需要相应的切换设备上下文,切换设备环境会带来一定的开销。
把计算图调度和算子执行的角色分开,有一个线程专门用来管理计算图,但该线程不允许执行算子,也就是并不直接启动CUDA核函数。相反地,算子的启动和执行被派发到不同的工作者线程(worker thread)上去,而且每个工作者线程都只为一个特定的设备服务。
在这种情况下,每个工作者线程只为一个设备服务,不需要切换设备上下文。同时,计算图只被调度线程访问,不需要考虑多线程并发读写的问题。
但是,算子之间的状态更新必须以调度线程为中介才能完成,譬如情况b所示,生产者算子和消费者算子在不同的工作者线程上,生产者算子执行完成后,会通知调度线程更新计算图,调度者线程更新完计算图之后才会影响消费者算子所在的工作者线程。
如果生产者算子和消费者算子可以直接对话,就不需要额外地引入和中介角色的调度线程来回交互的开销。
为了消除情况b引入的代理开销,可以允许同一个线程既做计算图的调度也做算子执行,每个线程仅调度和执行某一个特定设备上的算子(从而避免设备上下文的切换),每个线程都可以访问和修改计算图的状态,处于不同线程上的、具有上下游生产消费关系的算子之间可直接对话。
这种方法可消除生产者、消费者之间的调度线程作为中介引入的额外开销,但是,因为多个线程并发访问计算图的状态,这个全局状态必须使用锁来保护,需要尽可能减小并发访问引起的竞争开销。如果锁的临界区界定的范围过大(譬如整个计算图被一个粗粒度锁保护),可能会严重损害系统的性能。
为了尽量减少c情况中的竞争开销,观察到计算图的算子状态是局部的,可以通过算子对计算图的状态解耦,从而可以使用细粒度的锁(例如,每个算子的状态有一个专门的锁来保护)降低竞态条件引入的开销。但在一般的并发系统里,通过缩小锁的粒度(临界区的范围)来优化程序性能是很复杂的问题。