“黑科技”RDMA如何帮助深度学习模型优化通信网络？

很多人有疑问：为什么深度学习模型需要优化通信网络？

事实上，越是高性能算力的平台，越离不开通信网络的优化。近日，OneFlow创始人袁进辉与熊先奎先生（中兴通讯股份有限公司无线首席架构师）、宋庆春先生（国际高性能计算和人工智能咨询委员会亚太区主席）在《中兴通讯技术》上联合发表署名论文，探讨关于深度学习中通信网络优化的相关课题。让我们先来一起速读这篇论文的“精华”吧。

深度学习模型为何需要通信网络优化？

作为深度学习实现工程化的基础，以图形处理单元（GPU）多指令多数据流（MIMD）计算架构提供的高性能算力训练平台离不开计算范式的定制化设计和计算通信的网络优化，尤其是最新深度学习模型的训练需要消耗更多的算力。

鉴于深度学习模型计算主要是稠密型计算，业界广泛采用GPU等协处理器进行并行加速，但单个协处理器的算力仍无法满足日益增长的算力需求。通过高速互联技术把更多协处理器连接起来，能够协同输出大规模算力，可实现点对点数十至上百吉比特的传输带宽，例如节点内可以使用外设部件互连标准（PCIe）或NvLink等技术。当单个节点仍无法满足需求时，可通过高速网络实现多节点分布式计算。此时，由于传输带宽过低，普通以太网络（千兆网或万兆网）会出现多节点扩展效率极低、计算资源严重浪费的现象。因此，在分布式深度学习训练场景中，基于远程直接内存访问（RDMA）的网络通信成为最佳选择。

分布式深度学习训练任务必须使用RDMA技术，这是因为：
（1）深度学习训练任务普遍使用随机梯度下降算法。每处理一小片数据就需要更新模型参数，计算粒度很细，对网络传输的延迟容忍度非常低。
（2）深度学习普遍使用的GPU加速卡，吞吐率非常高。如果数据搬运速度跟不上计算速度，就容易造成计算资源浪费。
（3）在深度学习训练中，系统调度、数据加载和预处理均需要使用CPU资源，而基于传统以太网的网络传输也需要消耗很多CPU计算资源，这会影响整个系统的效率。RDMA的内核旁路技术可降低CPU利用率，提高整个系统效率。

分布式AI中使用RDMA面临的问题

现有深度学习框架极少直接基于RDMA技术开发网络通信库，通常都是通过NvidiaNCCL来间接使用RDMA，这可能导致两个潜在的问题：一方面，非Nvidia通用图形处理器（GPGPU）的加速器厂商如果想要实现分布式深度学习训练功能，就必须自己开发类似于NvidiaNCCL的通信库来使用RDMA；另一方面，深度学习框架与底层通信库相隔离，无法进行一体优化。

对于基于RDMA技术实现底层网络通信，最主要的挑战是内存管理。RDMA传输需要使用注册内存（锁页内存），如果每一次网络传输都根据传输量的实际需求当场申请并注册，则会导致开销显著，还会增加网络传输的延迟。如果提前申请并注册好内存块，则可能因无法准确预知程序运行时每次的实际传输需求而过多申请内存，这会造成内存资源的浪费。

如何解决？OneFlow已经想到

针对这一问题，一般有如下几个解决方法：（1）提前申请并注册大块内存，在实际传输时使用内存池技术，并从已注册好的大块内存上分配需要的内存；（2）提前申请并注册一些固定大小的内存块，每次在进行数据传输时，首先把需要传输的数据拷贝到这些已注册内存块上，然后再通过RDMA进行传输；（3）对于静态结构的神经网络，每一次迭代网络传输量是固定不变且可提前预知的，在系统运行前申请并注册内存，然后重复使用这些注册内存块。

OneFlow软件是目前唯一一个内部原生集成RDMA网络传输功能的深度学习框架。用户既可通过NvidiaNCCL使用RDMA传输，也可直接使用OneFlow基于RDMA自研网络通信库去实现分布式训练。针对动态形状和静态形状的网络传输，可分别使用上文描述的第2、3种技术解决方案。动态形状的网络传输须引入一次额外的拷贝，而静态形状的网络传输则实现了真正的零拷贝，达到了同时兼顾内存利用率和运行效率需求的目的。

OneFlow软件还充分发挥了RDMA的优点，采用Actor软件机制实现了一个简洁的去中心化调度系统框架。该软件先在编译阶段生成静态计算图，然后生成包含Actor实例的分布式环境描述信息的计划。分布式系统最终根据计划生成Actor实例运行态（各Actor间的生产、消费数据会被存储在Register中），通过RDMA低延迟网络消息传递协作来完成计算流水。

想要了解更多关于RDMA的内容，请点击这里阅读论文全文。

延伸阅读：关于《中兴通讯技术》
《中兴通讯技术》创办于1995年，是原广电总局认定的第一篇学术期刊。该刊主要介绍、探讨通信领域的新理论、新技术，传播电信知识，交流推广新业务、新经验，报道通信行业各类信息以加强通信界的学术与技术交流，展现中国国内外通信领域发展动态。现有审稿团队超过100人，覆盖了国内外30多所知名高校、研究院所等，其中不乏知名院士、专家。读者群为高等院校、科研院所及电信运营商的科研人员。