在混合云上大规模运行 AI 推理

  • Roblox 应用 AI(尤其是生成式 AI)的领域在过去几年中发展迅猛。

  • 我们启动了一个三阶段流程来构建和优化支持这种级别的 AI 工具所需的基础设施,现在我们已经进入最后阶段。

  • 我们正在共享构建混合云基础设施的步骤,这种基础设施能够支持大规模 ML 推理。

在上周的 RDC 上,我们公布了最新的 AI 孵化项目:开发一种可支持生成式创作的多模态 3D 基础模型。 为数百万人使用的始终在线、沉浸式 3D 全球平台使用的 AI 提供支持需要大量算力。 2023 年初,我们支持的机器学习 (ML) 推理管道不到 50 条。 目前,我们的基础设施支持大约 250 条这样的管道。 我们在两个数据中心和混合云基础设施上维护着数万个 CPU 和一千多个 GPU 来运行所有这些模型。 但我们的工作还没有结束。

我们之前曾分享过我们如何思考面向创作者的生成式 AI我们如何应用 AI 保障用户在 Roblox 上的安全以及 AI 翻译如何帮助全世界的人们进行沟通。 但这些只是不多的几个范例:目前大约有 250 个模型应用于生产环境,Roblox 上的几乎每一次交互都由某种形式的 AI 提供支持。 当某人第一次来到 Roblox 并考虑加入哪种体验时,AI 就会通过我们的推荐和搜索系统发挥作用。 当他选择一种体验并点击游玩按钮时,我们的匹配算法就会确定他可以加入的最优服务器。

数百万创作者已经可以使用我们生成式 AI 工具的强大功能。 利用助手,他们可以输入简单的提示来生成脚本和情节,加快体验创作。 借助我们的纹理材质生成器工具,他们可以快速更改和迭代对象的外观和样式。 随着最近推出虚拟形象自动设置,我们现在正在进入 4D 生成式 AI时代,该功能简化了虚拟形象创作过程,为创作者节省了数小时的工作时间。 截至 2024 年 8 月,在 Roblox 上发布的用户生成内容 (UGC) 虚拟形象身体中,约有 8% 是使用虚拟形象自动设置制作的。

我们现在已进入一项已持续多年的三阶段流程的最后阶段。 这次流程始于 2021 年末。 当时,由于缺乏统一的 Roblox AI 平台,工程团队不得不构建自己的小平台并选择完全不同的框架。 我们看到团队在开发关键组件,包括我们的虚拟形象商城、主页和搜索功能,每个团队都在构建自己的自定义特征工程。 团队没有利用集中式特征平台,而是将临时解决方案拼凑在一起。 此外,由于缺乏核心平台的支持,每个团队都需要承担繁重的工作,即独立开发自己的优化并应对推理扩展挑战。 这种碎片化方式凸显了我们迫切需要一个统一的、集中化平台,来简化我们的流程并全面提高效率。

第一阶段:为机器学习 (ML) 奠定坚实基础

我们很早就采用了 Kubeflow,以利用其打包的 ML 核心建造单位,包括笔记本、管道、离线实验和模型部署。 特征平台仍然是必要的,因此最初我们采用了第三方解决方案。 为了让 Roblox 的工程师更容易使用 ML,我们开发了 roblox-ml,这是一个 Python 库,它进一步降低了将模型部署到生产环境的复杂性。

我们利用 Jupyter 笔记本提供针对模型迭代优化过的开发环境,并且服务器经过配置以提供必要的数据访问和 GPU 资源。 扩展或定期运行训练作业以重新训练模型,通常需要我们编写管道。 借助我们的 roblox-ml库,工程师能够通过对运行时环境和源代码拍摄快照(无需构建 Docker 镜像),并通过选择具有优先级的计算资源、设置通知和处理身份验证,轻松地将笔记本代码转换为 Kubeflow 管道。

只有在正确的时间具备正确的特征,模型才会有效。 我们的特征平台简化了定义新特征的过程,同时促进了 100 多个特征服务中 900 多个特征的共享。 随着我们的特征集合不断增长,团队能够更快地创建和部署新模型。

在我们的 ML 管道平台正常运行并保持稳定后,我们看到对在线推理支持的需求有所增加,其中个性化、搜索和商城处于前列。 虽然我们建议将批量推理作为成熟 ML 运维的导航点,但我们开发了模型注册和部署平台来支持实时推理。 通过我们的模型注册,Roblox 工程师可以使用 roblox-ml 上传和下载模型,这些模型会被标记并自动进行版本控制,以促进可追溯性、回滚和 A/B 测试。 例如,我们的个性化模型每天都会进行训练和部署,并且我们经常同时运行大约 20 个 A/B 测试。 对于我们的部署平台,我们使用包含 Triton Inference Server 的 KServe 作为底层模型部署运行时,因为它具有强大的性能,并且支持使用 GPU 和 CPU 的多种 ML 框架。

无论是批量操作还是在线操作,Roblox 上的模型在发布之前都会经过大量的测试。 包括离线实验、影子测试和 A/B 测试。 发布后,模型会受到持续监控,以确保其在运行(例如推理延迟)和准确性方面均达到预期性能。 作为我们对安全和文明的承诺,人工审核员也会评估任何报告的推理不一致,这有助于确保我们做出正确的重大决策,并帮助改进模型的训练数据集。

第二阶段:为推理扩展做好准备

2023 年初,我们看到了生成式 AI 在加速 Roblox 上创作活动方面的巨大潜力。 为了充分利用这一潜力,我们在 2023 年的大部分时间里,都在优化 ML 训练和推理基础设施的性能和效率。 由于这些优化,我们显著降低了 CLIP 嵌入式创建的计算成本。 首先,我们扩展了分布式训练系统,以支持使用大型数据集进行训练以及在多个工作节点上运行具有数十亿个参数的模型。

当我们开始构建分布式工作流时,我们意识到现有的离线推理设置无法支持我们看到的长期增长率。 我们的最初设置是为实时推理设计的,其中输入和输出数据是有序的。 虽然由于我们早期的努力,它的效果很好,但它不能方便地支持任务并行或多阶段处理,而且资源效率也不够高,无法支持我们目前需要的规模。 此外,工程师还需要编写自己的数据分块和错误处理逻辑,随着我们的推理需求扩展,这些工作变得越来越耗时。

为了应对这些挑战,我们增加了对 Ray 的支持,Ray 是一个开源计算框架,可以轻松扩展批量推理工作量。 通过构建基于 Ray 的批量推理分布式任务管道,我们能够优化资源利用率、支持多阶段处理并提供强大的任务并行和更高的容错能力。 此外,Ray Data 库让工程师仅用几行代码就可以定义具有流式执行的管道,这有助于提高开发人员的速度和效率。 到目前为止,我们已经看到将 Ray 用于批量推理带来的巨大效率提升。

随着我们的推理需求不断增长,我们将所有的 CPU 推理迁移到我们自己的数据中心,我们就能够更直接地控制延迟和隐私设置。 我们每天为 7950 万日活跃用户处理大约 10 亿条个性化请求(截至 2024 年 6 月 30 日)。 将这些工作负载迁移到我们自己的数据中心,有助于我们保持效率,同时又不影响用户体验。 为了节省推理成本,许多系统会缓存请求——这会导致推荐过时,因为许多用户会在一天中多次访问 Roblox 主页。 然而,这也提高了我们的效率,使我们能够更好地优化推理的运行位置并分配工作量,从而减少所需的计算资源。

随着我们不断扩展,我们意识到需要一个自定义特征平台解决方案,它能够支持高吞吐量、低延迟和成本效益,同时还能实现各种服务的快速迭代。 现有的第三方解决方案无法满足这些要求,因此我们开发了自己的自定义特征平台,它建立在开源项目 Feast 之上。 我们的特征平台提供了一种自定义领域特定语言,用于定义批次特征和流处理特征的转换。 我们采用 Flink 作为流处理引擎来支持实时特征,这对于需要整合最新信息的模型至关重要。 另一方面,需要通过在分布式环境中重新运行 Roblox 游戏引擎,在批次中处理大量 3D 资产来获得一些特征。 目前,我们的特征平台每天摄取大约 300 亿条记录,提供大约 700 亿条记录,P99 延迟为 50 毫秒—并且支持超过 100 种特征服务。

随着语义理解需求不断增长,模型对嵌入的使用也迅速增长,无论是通过 NLP、计算机视觉还是推荐系统。 这促使我们建立一个矢量数据库,以便高效地作为高维点存储和检索矢量。 矢量数据库实现了快速最近邻查找,为多模式搜索和内容违规检测等功能提供支持。

随着越来越多的团队开始使用 ML 模型,我们希望找到规模效率并帮助工程师更快地取得成功,因此我们建立了自己的真实数据团队。 该团队帮助工程师设计自己的数据集生产管道,通过人工评估员训练和验证数据,并提供高质量数据。 这帮助我们实现了数据管道构建和数据集验证流程的标准化,并统一了数据传输、跟踪和监控的格式。

第三阶段:实施大规模推理

随着 Roblox 助手的推出,我们看到每周处理的令牌数量增加到 15 亿个。 我们还发布了新功能,包括实时 AI 聊天翻译和我们的语音安全模型(目前已经开源),令推理功能的需求激增。 我们启动了两个核心项目来增强 AI 应用开发:我们的 ML 网关和基于 vLLM 项目的大型语言模型运维 (LLMOps) 平台。 这两个项目将共同为 Roblox 的下一代 ML 奠定基础。

我们构建了统一的 ML 网关,以集中访问各种环境中的所有(开源和内部开发的)大模型,包括云端和本地的 CPU 和 GPU。 我们的目标是创建一个高效、精简的系统,用来管理全公司的 AI 资源。 在后端,该网关提供公共 API 接口、用户友好的配置选项并在我们部署的所有模型之间实现高效资源共享。

该网关通过对生成式 AI 工作量的令牌计数进行集中控制,并提供区域之间的延迟感知负载平衡,增强了我们的推理服务的弹性。 此外,它通过集中 API 密钥管理增强了安全性,支持全面的使用情况跟踪和潜在的权限实施,还可与监控工具集成以提高可观测性。 所有这些功能将帮助我们优化大模型的使用,降低成本,并为 Roblox 的工程师提供宝贵的见解。

此外,我们采用 vLLM 作为 LLM 的主要推理引擎,利用 vLLM 的高性能功能为 Roblox 上的 AI 应用提供支持。 自从采用 vLLM 以来,我们看到延迟和吞吐量都改进了近 2 倍,目前我们每周提供约 40 亿个令牌。

我们选择 vLLM 符合我们利用开源尖端技术的承诺,这些技术可以高效扩展,满足我们庞大的用户群和多样化体验的需求。 Roblox 是开源 vLLM 项目的积极贡献者,带头开发 vLLM 的多模式支持,这使得该引擎不仅可以处理文本还可以处理图像,并且将来有可能处理其他类型的数据。 我们还实现了推测解码技术来进一步提高推理性能,从而更快、更高效地处理语言任务。

借助 ML 网关和 vLLM,我们可以高效支持 Roblox 上使用的数百条 ML 管道—并随着对 AI 功能需求的不断增长而继续扩展推理。 而我们的这项工作还远没有完成。 我们对 Roblox 的 AI 未来有着宏伟的计划。 我们正在开发创新的 AI 工具,使新手和专家创作者都能更高效地创作。 一如既往,我们在想方设法提高基础设施的性能和效率,以更好地支持我们和我们的创造者每天都在使用的 AI 工具。

我们对开源的承诺

我们是在多个成功开源项目的基础上取得今天这一成就的。 我们的许多技术堆栈都是使用上面提到的开源技术构建的。

我们致力于成为开源 AI 社区的坚实的合作伙伴,并贡献我们自己的一些开源技术。 我们最近发布了我们的第一个开源模型,我们的语音安全分类器,而且我们目前正在开发我们的 ML 网关,并希望也能将其开源。 我们相信 AI 的未来应该是开放和透明的,我们很高兴成为这个社区的活跃成员。