Databricks研究团队如何使用AI Runtime

在 Databricks，我们持续开发和评估新的 AI 研究技术，使我们的产品对客户更智能、更高效、更有用。

我们的研究涵盖 Genie、Smart Routing 和 AI Functions 等产品。最近，我们发表了如何使搜索速度提升 3 倍，开发了离线强化学习的新技术，并推动了复杂文档提取的前沿。

为了构建出色的 AI 模型并管理数百个并行研究工作流，我们需要能够跟上步伐的基础设施。这就是我们使用 Databricks AI Runtime（AIR）的原因。

在这篇文章中，我们将分享 AIR 如何帮助我们快速实验、扩展到大规模训练工作负载，并适应智能体主导的训练，以及我们在跨研究项目管理 GPU 资源方面所学到的经验。

挑战：大量实验，计算需求却大相径庭

我们的研究团队开发的模型为 Databricks AI Functions 提供支持，包括 ai_parse_document 和 ai_extract。例如，Databricks AI Functions 提供最先进的文档解析质量，其背后是经过训练以理解企业文档的模型。要达到这样的质量，需要在不同架构、数据集、训练技术和超参数方面进行大量研究，涉及数万次训练和评估运行。

我们使用监督微调（SFT）、离线强化学习（RL）和在线 RL 来训练从几十亿参数到近一万亿参数的模型。根据工作负载的不同，我们使用 NVIDIA H100 或 B300 GPU，两者均可通过 AIR 获得。我们最大的训练运行可能需要数百块 GPU 协同工作数天或数周，并依靠高带宽网络来保持训练效率。

支持这些工作负载带来了三大挑战：

1. 长时间运行训练任务的弹性。 训练可能需要数千个 GPU 小时。单个节点或 RDMA 网络故障就可能让数百块 GPU 闲置，并迫使我们从较早的检查点重新开始，损失数小时的进度。对于像 KARL 这样的强化学习项目，可靠的基础设施至关重要。
2. 快速实验和原型设计。 改进 AI Functions 意味着要测试数千种方法，从不同的模型和训练技术到数据配方。我们需要快速启动实验，而无需反复配置基础设施或重建环境。
3. GPU 集群管理。 我们要在多个具有不同计算需求、优先级和截止日期的研究项目之间取得平衡，同时保持足够的灵活性，以应对新的研究技术或开源模型。

采用 AI Runtime 来解决这些挑战

自 AI Runtime 去年问世以来，我们的研究团队一直是其客户零号。我们使用 AIR 开发、训练和测试为 Databricks AI 产品提供支持的模型，依赖其 GPU 基础设施以及与现有工具的集成。

内置 GPU 弹性

AIR 会对分配的 GPU 节点执行预检，以确保它们在训练开始前健康且能有效通信。它还会在执行期间监控节点健康状况，检测 NCCL 通信故障、热节流或功耗节流以及 GPU 驱动程序问题等。如果某个节点变得不健康并导致作业失败，AIR 会检测到该问题并隔离受影响的节点，这样后续作业就不会被调度到同一有问题的基础设施上。

这对 KARL 等项目尤为重要，因为复杂的 rollout 生成以及超过 1000 亿参数模型的后期训练可能需要数千个 GPU 小时。与其反复排查基础设施故障，我们可以专注于改进模型。

使用 AIR CLI 和自定义容器实现快速迭代

研究要求我们频繁更改模型配置、训练方法和代码。借助 AIR CLI，我们可以以编程方式启动工作负载，指定所需的 GPU，并让 AIR 处理配置和调度，包括跨云分布式训练所需的高带宽网络。

我们还依赖自定义 Docker 容器进行本地开发。AIR 让我们能够使用相同的镜像进行训练，确保兼容性，而无需重建环境。我们可以上传代码库，在这些容器中执行自定义代码，并使用 Databricks Artifact Registry 注册和共享镜像以用于各项实验。AIR 在启动延迟和容器优化方面的投入帮助我们在几秒钟内启动训练作业，从而在测试新想法时缩短反馈循环。

使用 MLflow 和日志进行实验跟踪

我们使用 MLflow 跟踪研究实验。AIR 会自动为我们的训练作业创建和管理 MLflow 实验，无需单独配置。

我们可以监控作业状态、流式传输和检查训练日志，并检索指标以了解实验进展。这也让我们的 AI 智能体能够获取评估结果并确定下一步尝试所需的信息。

无缝、受治理的研究数据访问

我们所有的研究数据都存储在 Unity Catalog 中。通过 Unity Catalog FUSE 挂载，AIR 让我们可以像访问本地文件系统一样访问这些数据，即使跨不同的 Databricks 工作区工作也是如此。

这使我们的数据保持可访问、可跟踪、可共享和受治理，同时无需改变训练代码与文件交互的方式。结合 MLflow 和基于 Docker 的环境，它为我们提供了从本地开发到分布式训练的一致工作流。

跨研究项目管理局 GPU 容量

我们的工作负载资源需求差异很大。调试实验可能只需要几块 GPU，而大型训练运行可能需要数百块 GPU 持续数天或数周。不同项目还有相互竞争的时间表和优先级。

AIR 的工作负载优先级排序和抢占能力帮助我们在各项目之间分配资源、优先处理时间敏感的实验，并在有空闲容量时利用这些富余容量。例如，空闲容量可以在周末支持额外的探索性实验，而不会干扰更高优先级的工作。我们的目标是让 GPU 资源始终可用于真正重要的实验，而无需研究人员手动协调基础设施。

在 AIR 上借助智能体主导的训练进行扩展

由 Andrej Karpathy 推广的 autoresearch 正在改变我们开发 AI 模型的方式。研究人员不再手动启动和评估每一个实验，而是越来越多地使用智能体来提出变更、运行实验、评估结果，并决定下一步尝试什么。

这一转变极大地增加了我们的实验量。在我们的 AI Functions 研究工作负载中，训练作业的数量增长了 8 倍。

随着智能体承担越来越多研究流程中的工作，我们的基础设施也需要跟上。当研究人员只启动少量实验时，启动缓慢或令人困惑的配置错误或许还可以应付。但当智能体启动数百个实验时，这些同样的问题很快就会成为瓶颈。

智能体驱动的研究需要快速的作业启动、程序化接口、清晰且可操作的错误信息、可靠的执行，以及可访问的训练结果。

我们的 autoresearch 工作流遵循一个持续循环：

1. 提出： 智能体识别模型或训练方法的潜在改进。
2. 启动试验： 它们并行向 AIR 提交多个实验。
3. 学习： 它们评估训练指标，以确定哪些改进能提升性能。
4. 更新计划： 它们利用结果来改进方法并启动新的实验。

AIR 特别适合这一工作流，因为智能体可以通过程序化方式与其交互，并受益于其快速的作业启动、内置的可靠性以及扩展能力。通过 Databricks CLI 和 databricks-air-skills，我们的智能体可以启动并发工作负载、监控作业状态、检索 MLflow 指标，并根据结果进行迭代。

这让我们能够运行更多实验，同时研究人员可以专注于开发新想法和提升模型质量。随着我们进一步迈向智能体驱动的研究，AIR 为我们提供了基础设施，以支持更大的模型和更快的实验节奏。

立即开始

我们使用 AI Runtime 的经验表明，可靠、灵活的 GPU 基础设施对研究有多么重要。从大规模训练和受治理的数据访问，到智能体驱动的实验，AIR 让我们能够花更少时间管理基础设施，花更多时间开发更好的模型。

如果你正在训练或微调自己的模型，请查看 AI Runtime 文档 以开始使用。

AI Runtime 还为我们多个领域的研究提供了支持，例如 更快的搜索、强化学习 和文档理解。在 Databricks AI Research 探索我们更多的研究成果。