其他
资源依赖的“诅咒” | 原有深度学习框架的缺陷①
一个显而易见的出发点是,我们看到了原有的主流深度学习框架的本质不足。尤其在抽象层面和API层面,它们的设计有种种不足,导致开发者在使用时带来了极大不便,尽管他们正在试图解决一些缺陷,但有些重要问题依然被忽视了。
为此,我们将推出三篇系列文章,详细论述原有主流深度学习框架的运行时系统的三大不足,此为第1篇内容。本文将介绍资源依赖的重要性,以及原有框架在应对该难题上的局限性,最后,本文将介绍OneFlow框架如何在设计之初就简单、优雅地解决了这一难题。
撰文 | 袁进辉
1 算子之间的三种依赖关系
忽略因共享资源而形成的依赖是深度学习框架设计上的致命不足,这个问题会降低系统的安全性和稳定性;
在现有框架的设计体系内解决这一问题不是不可能,但难度比较大,且会破坏系统抽象的优雅性;
OneFlow中基于actor机制可以比较简单的解决这个问题。
如果O2先被调度,它可以成功执行并消费内存中M2的输出。O2完成之后,O2的输入和输出就可以被释放 (状态4)。然后O1就可以被调度并成功执行。
如果是O1先被调度,那么剩余内存对O1来说是不够的,但调度器在发射O1指令的时刻并不关心也不知道当前的内存是否能满足O1的需求。
3
借助控制边规定算子的执行顺序并非易事
4 通过定制内存分配器支持双缓冲流水线
图3就展示了在TensorFlow里实现双缓冲的一个潜在方案,也就是添加一个自定义的内存分配器,这个内存分配器限制每个算子只能分配到两份内存。
当一个算子使用分配器分配内存时,分配器会查询一个计数器,查询该算子是否有空闲的缓冲区可以使用。如果有,它就为该算子分配一个缓冲区,让该算子继续执行,并在该算子完成计算后释放缓冲区。如果这个算子的两个缓冲区都已被占用,分配器会将步骤2 (do compute)和步骤3 (release) 放入一个等待列表中。
当一个算子向分配器释放其内存时,分配器更新与该算子对应的空闲缓冲区的计数器,并检查等待列表中是否有请求该缓冲区的算子。如果有,就从等待列表中弹出处于就绪状态的算子,执行步骤2和步骤3。
如何优雅解决资源依赖问题?
Existing TF kernels encapsulate shape computation and memory allocation within the kernel implementation, making some graph compiler optimizations challenging or infeasible, such as reusing buffers across kernels. In TFRT kernels, shape computation and memory allocation will be hoisted out of the opaque C++ kernel implementations. A core design principle of TFRT is that kernel executions are never allowed to block, as this allows us to have a fine-grained control over the number of compute threads and the thread switching behavior, which is important for achieving high CPU utilization.