AI 分布式框架 Ray

随着机器学习算法和技术的进步,出现了越来越多需要在多台机器并行计算的机器学习应用。然而,在集群计算设备上运行的机器学习算法目前仍是专门设计的。尽管对于特定的用例而言(如参数服务器或超参数搜索),这些解决方案的效果很好,同时 AI 领域之外也存在一些高质量的分布式系统(如 Hadoop 和 Spark),但前沿开发者们仍然常常需要从头构建自己的系统,这意味着需要耗费大量时间和精力。

UC Berkeley AI 研究院内的研究员,发布的分布式系统 Ray。开发者称,Ray 专门为人工智能应用设计,通过这款框架,运行于笔记本电脑上的原型算法仅需加入数行代码就可以转化为高效的分布式计算应用。近日,该框架已被开源。Ray 的目标之一在于:让开发者可以用一个运行在笔记本电脑上的原型算法,仅需添加数行代码就能轻松转为适合于计算机集群运行的(或单个多核心计算机的)高性能分布式应用。这样的框架需要包含手动优化系统的性能优势,同时又不需要用户关心那些调度、数据传输和硬件错误等问题。

与深度学习框架的关系:Ray 与 TensorFlow、PyTorch 和 MXNet 等深度学习框架互相兼容,在很多应用上,在 Ray 中使用一个或多个深度学习框架都是非常自然的(例如,UC Berkeley 的强化学习库就用到了很多 TensorFlow 与 PyTorch)。

与其他分布式系统的关系:目前的很多流行分布式系统都不是以构建 AI 应用为目标设计的,缺乏人工智能应用的相应支持与 API,UC Berkeley 的研究人员认为,目前的分布式系统缺乏以下一些特性:

  • 支持毫秒级的任务处理,每秒处理百万级的任务;
  • 嵌套并行(任务内并行化任务,例如超参数搜索内部的并行模拟,见下图);
  • 在运行时动态监测任意任务的依赖性(例如,忽略等待慢速的工作器);
  • 在共享可变的状态下运行任务(例如,神经网络权重或模拟器);
  • 支持异构计算(CPU、GPU 等等)。
开源的人工智能框架
开源的人工智能框架

一个嵌套并行的简单例子。一个应用运行两个并行实验(每个都是长时间运行任务),每个实验运行一定数量的并行模拟(每一个同时也是一个任务)。

Ray 有两种主要使用方法:通过低级 API 或高级库。高级库是构建在低级 API 之上的。目前它们包括 Ray RLlib,一个可扩展强化学习库;和 Ray.tune,一个高效分布式超参数搜索库。