2020 年底,苹果发布了 M1 芯片,提供了全新的自研 GPU,并和 CPU 集成封装成 SoC,加上随之而来的 ML Compute 框架,macOS 端的深度学习训练有了显著的速度提升。
另一方面,强化学习的特性使得训练过程中需要交替使用 CPU 和 GPU,对此,M1 芯片能否提升强化学习的训练效率?据了解,此前业界并无相关测试数据。
南栖仙策(南京大学人工智能创新研究院技术孵化企业)算法小组用 M1 芯片设备组建小型服务器集群,并与基于 NVIDIA GPU 的传统服务器进行性能对比测试,结果表明,强化学习的训练效率可以得到显著提升。
什么是强化学习
强化学习(Reinforcement Learning, RL),又称增强学习,是机器学习的范式和方法论之一。强化学习智能体(Agent)在与环境的交互过程中,通过更新策略以达成回报最大化或实现特定目标。
广为人知的人脸识别等技术背后的深度学习,通常需要开发者准备大量做好标注的图像作为训练数据。
强化学习与之不同,不需要训练数据,而是设置一个环境。类似于动物处在自然环境中,强化学习需要不断与环境交互,并从中学习。强化学习研究常用的环境是各种各样的游戏。
强化学习方法简单、有效且通用,在许多游戏环境中都取得了超越人类的决策水平,被 DeepMind 认为是实现通用人工智能的一种重要途径。
强化学习是怎么训练的
图 1 强化学习示意图
如图 1 所示,在强化学习中,智能体与环境的交互,即在状态 S_t 下,执行动作 A_t,获得对应的奖励 R_(t+1) 并观测到下一个时刻的状态 S_(t+1)。智能体的目的是改进策略(即在不同状态下执行的动作),来获得最大的长期回报。在这个过程中,强化学习智能体按照当前策略(通常用神经网络实现)需要将当前状态 S_t 输入到网络中,进行一次前向推理(inference)得到 A_t,使用 A_t 在环境中执行一步来完成交互。此时,智能体会收到对应的奖励和下一个时刻的状态,然后不断循环该过程,产生新的样本用于策略训练。样本指的是交互过程中的每一个转移(transition),即四元组。
通常来说,强化学习中环境执行动作的这一步由 CPU 完成,当策略网络仅输入一个样本时,可以使用 CPU(对于浅层的策略网络)或 GPU(对于深层的复杂策略网络)进行推理,而更新策略网络模型则在 GPU 上完成。常规强化学习任务(以 Atari 游戏为例)通常需要采集十数 M( 1M 即一兆,为 10^6) 乃至数十 M 样本才能训练好一个策略,那么就需要数千万次,甚至数十亿次的交互。如果能减小时间开销,那将会给强化学习的训练速度带来前所未有的提升。
M1 芯片特殊之处
M1 架构
新的 M1 架构采用了台积电最新的 5nm 技术。M1 处理器中有 1 个 CPU, 1 个 GPU, 1 个神经引擎以及 DRAM 单元。由于 RAM 内置在处理器中,因此所有组件都会访问相同的内存。这意味着不会再因为将内存从一个组件移动到下一个组件而造成性能损失,无需再尝试弄清楚你的模型是否符合 GPU 的内存规格。
CPU
M1 芯片的 CPU 是 8 核 CPU,其中 4 核是用于数据处理和一些需要高性能任务的高性能内核,另外 4 核被称为「e 核」或高效内核,提供更高的处理效率和更低功耗的任务。
简单来说,M1 芯片不仅靠 5nm 制程使处理性能更加强大,而且里面新增的 NPU 也使 M1 更加聪明。该 NPU 采用 16 核架构,每秒进行 11 万亿次运算。M1 芯片使用共享内存,CPU 与 GPU 通信开销小,使用 Thunderbolt ,通信速度最高可达 40Gb/s,突破了限制瓶颈,可使用多台机器组建小型服务器。
南栖仙策使用四台 M1 的 Mac mini 组建了小型集群
Apple M1 集群性能测试
测试设置
M1 (单机) :使用 M1 Mac mini 单机训练;M1 集群 (4 机):4 台 M1 Mac mini 组成集群进行训练;GPU 服务器 (单机 V100×1):使用 NVIDIA V100 单 GPU 训练,选择第一张卡;GPU 服务器 (单机 V100×4):单机使用四块 NVIDIA V100 进行分布式训练;CPU 服务器: 使用 V100 服务器的 CPU 进行训练。
环境:
Pendulum、MountainCarCont、HalfCheetah
算法:
PPO, 共采集 100 次数据进行 PPO 训练,每次训练前与环境交互采集 10K step 的轨迹放入 Buffer 进行训练。SAC, 训练 100 个 Epoch,每个 Epoch 与环境交互收集 1k step 放入 Buffer。
测试结果
由于目前 Tensorflow 2.4 的 Eager 模式不能调用 M1 芯片进行训练,下面的结果是在 Mac mini 的 CPU 设备上测试获得的。
测试一:单进程采样对比
测试二:4 进程采样对比
测试三:16 进程采样对比
测试四:32 进程采样对比
功耗及价格对比
注:单个 Mac mini 功耗 39w,4 台总计 156w(瓦),传统服务器约为 1300w(瓦)。
注:单个 Mac mini 价格 6799 元 (512GB 官网价格),4 个总计 2.7 万元,传统服务器约为 30 万元。
M1 的共享内存架构适合于 RL 的训练方式,目前,使用 Apple 的 M1 芯片组建的小型服务器,以 NVIDIA GPU 服务器 10% 的价格和更低的功耗,取得超过 200% 的时间效率提升。在强化学习目前主流训练框架下,训练效率的瓶颈往往在于策略网络的通信速度不够快,而共享内存方案可以减小通信上的开销,为 RL 训练带来巨大的效率提升。