How our AI Research team uses AI Runtime

Databricks AI研究团队如何使用AI Runtime | Databricks博客

跳至主要内容

工程

2026年10月15日

Databricks AI研究团队如何使用AI Runtime

作者：Ethan Tang、Ivan Zhou、Irene Dea、Steven Chen、Bruce Fontaine和Ben Hansen

摘要

• Databricks AI研究团队使用AI Runtime进行大规模模型训练和微调，从小型实验到覆盖数百块GPU的工作负载。

• 内置的GPU可靠性、快速任务启动以及与MLflow和Unity Catalog的集成，帮助研究人员更高效地进行实验。

• 基于代理的研究使训练任务数量增加了8倍，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的首个内部用户。我们使用AIR开发、训练和测试驱动Databricks AI产品的模型，依赖其GPU基础设施以及与现有工具的集成。

内置的GPU容错能力

AIR在训练开始前会对分配的GPU节点执行预检检查，确保节点状态正常且能够有效通信。在执行过程中，AIR还会持续监控节点健康状况，检测NCCL通信失败、温度或电源限制、GPU驱动问题等异常情况。如果某个节点出现故障导致作业失败，AIR会自动识别问题并将受影响节点隔离，防止后续作业被调度到相同的问题基础设施上。

这对于KARL等项目尤为重要，该项目在生成复杂模型部署方案和对超百亿参数模型进行后训练时，可能需要数千小时的GPU计算时间。通过AIR的容错机制，我们可以避免反复排查基础设施故障，从而专注于模型优化。

通过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 倍。

随着代理承担更多研究流程，我们的基础设施需要同步升级。当研究人员启动少量实验时，启动速度慢或配置错误可能尚可接受。但当代理启动数百个实验时，这些问题会迅速成为瓶颈。

代理驱动的研究需要快速启动任务、程序化接口、清晰可操作的错误提示、可靠的执行环境以及可访问的训练结果。

我们的自动研究工作流遵循持续循环：

• 提出方案：代理识别模型或训练方法的潜在改进点。
• 启动试验：它们并行向 AIR 提交多个实验。
• 学习分析：它们评估训练指标以确定哪些改进能提升性能。
• 更新计划：它们利用结果优化策略并启动新实验。

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

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

今天立即开始

我们使用 AI Runtime 的经验表明，可靠且灵活的 GPU 基础设施对研究至关重要。从大规模训练、受控数据访问到代理驱动的实验，AIR 让我们能够减少基础设施管理时间，专注于开发更优质的模型。

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

AI Runtime 还支持我们在快速搜索、强化学习和文档理解等领域的研究。更多内容请访问 Databricks AI Research。

订阅最新动态

订阅我们的博客，获取最新文章直接发送到您的邮箱。

立即注册

查看所有博客

slice-start id="_gatsby-scripts-1"

slice-end id="_gatsby-scripts-1"