面向AI应用分布式执行框架

导读:下一代人工智能应用程序需要不断地与环境交互,并从这些交互中学习 。这对系统的性能和灵活性提出了新的要求,而现有的机器学习计算框架大多无法满足这些要求 。为此,UC Berkeley 项目组开发了一个新的分布式框架 Ray,并于近日在 Arvix 上发表了相关论文:《Ray: A Distributed Framework for Emerging AI Applications》 。
论文第一作者为 Philipp Moritz 及 Robert Nishihara,是 UC Berkeley AMP Lab 的博士生,而 Michael I. Jordan 和 Ion Stoica 的名字也赫然列于其中 。
Michael I. Jordan :UC Berkeley 电气工程与计算机科学系和统计系杰出教授,是美国国家科学院、美国国家工程院、美国艺术与科学院三院院士,是机器学习领域唯一获此成就的科学家 。2016 年,他被 Semantic Scholar 评为“最具影响力的计算机科学家” 。
Ion Stoica :UC Berkeley 计算机系教授,AMPLab 共同创始人,弹性 P2P 协议 Chord、集群内存计算框架 Spark、集群资源管理平台 Mesos 核心作者 。
目前的计算框架存在的短板
如今大部分人工智能应用都是基于局限性较大的监督学习的范式而开发的,即模型在线下进行训练,然后部署到服务器上进行线上预测 。随着该领域的成熟,机器学习应用需要更多地在动态环境下运行,响应环境中的变化,并且采用一系列的动作来完成既定目标 。这些要求自然地建立在增强学习(Reinforcement Learning,RL)范式中,即在不确定的环境中连续学习 。
RL 应用与传统的监督学习应用有三个不同之处:
1)RL 应用严重依赖仿真来探索所在状态及操作结果 。这需要大量的计算,现实情况下,一个应用大概需要进行亿万次仿真 。
2)RL 应用的计算图是异质的、动态变化的 。一次仿真可能会花掉几毫秒到几分钟的时间,仿真的结果又决定未来仿真的参数 。
3)许多 RL 应用程序,如机器人控制或自主驾驶,需要迅速采取行动,以响应不断变化的环境 。
因此,我们需要一个能支持异质和动态计算图,同时以毫秒级延迟每秒处理数以百万计任务的计算框架 。而目前的计算框架或是无法达到普通 RL 应用的延迟要求(MapReduce、Apache Spark、CIEL),或是使用静态计算图(TensorFlow、Naiad、MPI、Canary) 。
RL 应用对系统提出了灵活性、表现性能以及易开发的要求,Ray 系统则是为满足这些要求而设计的 。
示例

面向AI应用分布式执行框架

文章插图
 
【面向AI应用分布式执行框架】经典RL 训练应用伪代码
面向AI应用分布式执行框架

文章插图
 
用Ray 实现的Python 代码样例
在Ray 中,通过@ray.remote 声明remote 函数和actor 。当调用remote 函数和actor methods 时会立即返回一个future(对象id),使用ray.get()可以同步获取该id 对应的对象,可以传递给后续的remote 函数和actor methods 来编码任务依赖项 。每个actor 有一个环境对象 self.env,在任务之间共享状态 。
面向AI应用分布式执行框架

文章插图
 
上图是调用train_policy.remote() 对应的任务图 。remote 函数和actor methods 调用对应任务图中的任务 。图中有2 个actor,每个actor 之间的状态边(stateful edges)意味着他们共享可变状态 。从train_policy 到它所调用的任务之间有控制边(control edges) 。要并行训练策略(policy),可以多次调用train_policy.remote() 。
原理
为了支持RL 应用所带来的异质和动态工作负荷要求,Ray 采用与CIEL 类似的动态任务图计算模型 。除了CIEL 的任务并行简化外,Ray 在执行模型顶层提供了代码简化,能够支持诸如第三方仿真的状态结构 。
Ray 系统结构
为了在支持动态计算图的同时满足严格的性能要求,Ray 采取一种新的可横向扩展的分布式结构 。Ray 的结构由两部分组成:application 层和 system 层 。Application 层实现 API 和计算模型,执行分布式计算任务 。System 层负责任务调度和数据管理,来满足表现性能和容错的要求 。
面向AI应用分布式执行框架

文章插图
 
Ray 系统结构
该结构基于两个关键想法:
1)全局状态存储 GSC(Global Control Store) 。系统所有的控制状态存储在 GSC 中,这样系统其他组件可以是无状态的 。不仅简化了对容错的支持(出现错误时,组件可以从 GSC 中读取最近状态并重新启动),也使得其他组件可以横向扩展(该组件的复制或碎片可以通过 GSC 状态共享) 。


推荐阅读