近日，上海数字大脑研究院（简称“数研院”）提出分布式经验回放系统GEAR（又名PipeDT），优化了现代GPU系统在数据储存、并行计算和高速通信的软硬件特性，更适用于RL大模型的经验回放过程。该项论文成果已被机器学习顶会ICML2023收录。随着近年来大型序列模型在自然语言处理（NLP）和计算机视觉（CV）的成功应用，它的强大建模能力也得到了强化学习领域研究人员越来越多的关注。2022年以DeepMind Gato模型的推出也标志着RL模型进入到了多领域任务泛化、规模化的快车道。数字大脑研究院在这一新兴方向也有大量早期布局 ，例如数研院提出的MAT模型是世界首个用Transformer结构解决多智能体强化学习问题的研究工作；数研院开源的数字大脑多模态决策大模型DB 1.0参数量达到十亿级别，在870个不同的决策任务上训练并在76%的任务上取得良好表现，填补了国内在此方面的空白。
经验回放是提升强化学习模型训练过程中数据效率乃至最终训练结果的重要一环。在回放过程中，专家策略或是训练策略和外部环境交互生成的经验以轨迹数据的方式被储存下来，并且在训练过程中通过先进先出(FIFO)、优先级回放(Prioritized Reply)等选择机制采样聚合，成为了后续并行训练过程中的数据批次。然后适用于RL大模型训练所需的经验回放存在如下挑战：1. 轨迹数据储存成本高：类似于语言模型，大型决策模型饱和训练需要海量的数据；同时受任务影响，其输入数据同时包含了原生状态向量、图片等数据，相较于文本数据有着更高的空间占用。举例而言，DB1模型训练采用的数据集由超过350B tokens组成，实际储存空间占用高达110TB，这带来了大量的储存成本。
2. 轨迹数据选择的高计算开销：轨迹数据选择过程可被堆、前缀和树等特定数据结构加速，然而即便如此，面对超大规模的数据集和大数据批次的组合器性能开销依旧十分惊人。例如前缀和树优化后的优先级回放(Priority Reply)的时间复杂度为，其中B为数据批次大小，N为数据集总量。对于一个十亿级别的数据集和十万甚至百万级别的数据批次来说，轨迹选择会带来大量的时间开销和运行空泡。3. 轨迹数据传输的高通信开销：RL任务中轨迹数据常常因为时间跨度长、多模态等因素带来巨大的通信开销。
图1: GEAR储存架构，GEAR能和现有的并行等方案良好兼容
为了优化RL模型在经验回放过程中存在的上述挑战，数研院提出了分布式经验回放系统GEAR，围绕现代GPU系统在数据储存、并行计算和高速通信的软硬件特性优化适用于RL大模型的经验回放过程，助力其规模化方向的研究探索。具体改进如下：
1. 分布式轨迹数据储存：如图1所示，根据强化学习模型训练中的访存特性，GEAR优化了轨迹数据储存格式。并结合模型并行方案和数据优先级等模型相关和任务相关信息进一步优化数据分布式切分、储存方案。这一设计使得GEAR能和现有主流的模型并行方法（例如数据并行、流水线并行）兼容的同时，最大化系统整体吞吐量。同时通过Trajectory Manager等组件，GEAR允许本地客户端读写本地内存中的数据分片或是从远程数据分片加载数据，提供了一个逻辑上简单统一的经验回放系统储存抽象。2. 使用GPU加速轨迹数据选择：GEAR也探索了使用高性能的CUDA算子加速轨迹选择过程。中心化采样方式由于有着更多同步、通信开销，在大规模并行训练中瓶颈更明显，但同等条件使用GPU加速后相较于于sum-tree等CPU采样实现也有一定优势。3. 以GPU为中心的轨迹收集：GEAR充分利用了当代GPU直接读取宿主机共享内存这一特性。面对RL训练任务中频繁进行离散轨迹选择和跨设备拷贝的工作流，GPU这一访存特性可以极大的提升整体吞吐量，同时释放了大量CPU资源，而后者可以被其他的RL任务比如外部环境模拟计算等利用。同时GEAR也充分利用了高性能GPU集群上的InfiniBand网络加速数据跨节点转移。
图二：（左）中心化GPU采样；（右）去中心化GPU采样为了对比GEAR和Reverb的实际性能，我们在DGX-A100集群上测试了两者轨迹选择的吞吐量。在测试流程中，类似于SPMD并行程序，多个 Client进程形成了一个同步数据拉取的数据并行进程组向中心Server进程发起数据请求。结果如图四所示，Reverb在64个并行client进程下取得了其单节点的最高吞吐量6.73GB.s，而GEAR使用8个client进程取得了最高17.1GB/s的吞吐。在后续拓展到多节点的实验中，GEAR也展现出了更好的性能拓展表现。
图三：DGX集群上多机吞吐量实验
GEAR是目前世界上运行最快的决策大模型训练使用的经验回放系统，在同等条件下，相较于DeepMind Reverb这一常用经验回放系统可提供最高达六倍的数据吞吐量，这对于决策模型在复杂任务或多模态输入等场景下的泛化能力的探索研究提供了系统侧支持。总体而言作为该领域的探索性工作，GEAR提出了以GPU为中心构建经验回放系统这一全新的设计可能，充分利用分布式集群中的GPU加速轨迹选择、收集过程，并且对于DeepSpeed等常用框架有较好兼容性。另一方面，决策大模型的端到端训练和评估管线也包含了经验回放之外的大量数据处理和计算，这些问题是GEAR目前无法涵盖的，如何更好的和现有系统做到全流程、端到端的整合加速仍是一个需要大量的研究工作投入的问题。数研院将持续关注该领域的前沿研究，并改进完善GEAR系统能力。数研院后续将在系列决策大模型（DB2.0等）的开发过程中应用GEAR系统加速训练评估流程，也将继续探索和改进GEAR对于海量数据集、更大规模决策模型以及模型训练推理混合模式等需求的支持。同时GEAR作为一个经验回放系统的独立工作，也将在整理后向社区开源。
相关链接
MAT模型：https://arxiv.org/abs/2205.14953
上海数字大脑研究院上海数字大脑研究院汇聚决策智能领域全球领先的科研力量，面向中国及全球产业智能化升级需求，致力于为全球企业提供决策智能解决方案，驱动各产业全面智能化升级。其孵化的独立初创企业谜题科技（Enigma Tech.）是一家致力于破解AGI谜题的创新型公司，公司有着探索未知领域、破解复杂问题、构建AI与人类和谐共处新世界的决心。数研院将携谜题科技一起，继续为人类的生活带来更多美好的可能性，共同开创一个充满智慧与梦想的未来。
