Databricks AI 研究团队如何用 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 协同工作数天或数周，并需要高带宽网络来保持训练效率。

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

长时间训练任务的韧性。 训练可能需要数千个 GPU 小时。单个节点或 RDMA 网络故障就可能让数百个 GPU 闲置，并迫使我们从较早的检查点重新开始，损失数小时的进度。对于像 KARL 这样的强化学习项目，可靠的基础设施至关重要。
快速实验和原型设计。 改进 AI Functions 意味着测试数千种方法，从不同的模型和训练技术到数据配方。我们需要快速启动实验，而无需反复配置基础设施或重建环境。
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 工作流遵循一个持续循环：

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

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

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

立即开始

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

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

AI Runtime 还为我们在更快的搜索、强化学习和文档理解等领域的研究提供了支持。欢迎在Databricks AI Research了解更多我们的工作，