我们的AI研究团队如何使用AI Runtime

我们的AI研究团队如何使用AI Runtime

💡 原文英文,约1200词,阅读约需5分钟。
📝

内容提要

Databricks研究团队利用AI Runtime(AIR)支持AI Functions等模型研发,覆盖数十亿至近万亿参数模型的训练。AIR提供GPU健康检查与故障隔离、快速启动、MLflow实验追踪和Unity Catalog数据治理,并支持智能体驱动研究,使训练任务量增长8倍,提升实验效率。

🔎

延伸解读

GPU故障隔离如何减少研究中断

AIR在训练前对GPU节点进行预检,确保节点健康且通信正常;训练中监控NCCL通信失败、热或功耗节流、驱动问题等。一旦节点异常导致任务失败,AIR会隔离该节点,避免后续任务调度到同一问题基础设施。对于需要数千GPU小时的大规模训练,这能减少反复排查基础设施故障的时间,让团队更专注于模型改进。

快速迭代:CLI与自定义容器

研究需要频繁调整模型配置、训练方法和代码。通过AIR CLI,团队可以编程方式启动工作负载、指定GPU,由AIR处理资源调配和调度,包括跨云分布式训练的高带宽网络。同时,AIR支持使用本地开发的自定义Docker镜像进行训练,无需重建环境,并可通过Databricks Artifact Registry注册和共享镜像。AIR在启动延迟和容器优化上的投入,使训练任务能在数秒内启动,缩短反馈循环。

智能体驱动研究对基础设施的新要求

受Andrej Karpathy提出的autoresearch启发,团队越来越多使用智能体提出变更、运行实验、评估结果并决定下一步。这使AI Functions研究负载的训练任务数量增长了8倍。当智能体启动数百个实验时,缓慢的启动或配置错误会迅速成为瓶颈。因此,智能体驱动研究需要快速任务启动、编程接口、清晰可操作的错误、可靠执行和可访问的训练结果。AIR通过Databricks CLI和databricks-air-skills支持智能体并发启动工作负载、监控状态、获取MLflow指标并迭代。

跨项目GPU容量管理策略

研究团队的工作负载资源需求差异很大:调试实验可能只需几个GPU,而大型训练运行可能需要数百个GPU持续数天或数周。不同项目还有竞争性的时间线和优先级。AIR的工作负载优先级和抢占能力帮助跨项目分配资源,优先处理时间敏感的实验,并在有空闲容量时加以利用。例如,周末的闲置容量可以支持额外的探索性实验,而不会干扰更高优先级的工作。目标是让GPU资源始终用于重要的实验,而无需研究人员手动协调基础设施。

❓

Q&A

Databricks AI Runtime (AIR) 是什么?

Databricks AI Runtime (AIR) 是 Databricks 提供的 AI 基础设施,用于支持 AI 模型的研究、训练和测试。它提供 GPU 基础设施、与现有工具的集成,并支持智能体驱动的研究。

AIR 如何确保 GPU 的可靠性?

AIR 在训练开始前对分配的 GPU 节点进行预检,确保它们健康且能有效通信。在执行过程中监控节点健康,检测 NCCL 通信失败、热或电源节流、GPU 驱动问题等。如果节点不健康导致作业失败,AIR 会检测并隔离该节点,防止后续作业调度到同一问题基础设施上。

AIR 如何加速实验迭代?

AIR 提供 CLI 以编程方式启动工作负载,指定 GPU 并处理供应和调度。支持自定义 Docker 容器,允许使用本地开发镜像进行训练,确保环境兼容。通过 Databricks Artifact Registry 注册和共享镜像。AIR 优化启动延迟和容器,使训练作业在几秒内启动,缩短反馈循环。

AIR 如何管理实验跟踪和数据访问?

AIR 自动为训练作业创建和管理 MLflow 实验,无需单独配置。可以监控作业状态、流式传输和检查训练日志、检索指标。所有研究数据存储在 Unity Catalog 中,通过 Unity Catalog FUSE 挂载,AIR 允许像本地文件系统一样访问数据,即使跨工作区,保持数据可访问、可跟踪、可共享和受治理。

AIR 如何帮助管理跨项目的 GPU 资源?

AIR 提供工作负载优先级和抢占功能,帮助跨项目分配资源,优先处理时间敏感的实验,并在有空闲容量时使用。例如,周末空闲容量可以支持额外的探索性实验,而不干扰高优先级工作。目标是保持 GPU 资源可用于重要实验,无需研究人员手动协调基础设施。

什么是智能体驱动研究,AIR 如何支持它?

智能体驱动研究(autoresearch)是由 Andrej Karpathy 推广的,使用智能体提出变更、运行实验、评估结果并决定下一步尝试。AIR 特别适合这种工作流,因为智能体可以以编程方式与 AIR 交互,受益于快速作业启动、内置可靠性和扩展能力。通过 Databricks CLI 和 databricks-air-skills,智能体可以启动并发工作负载、监控作业状态、检索 MLflow 指标并根据结果迭代。

🏷️

标签

➡️

继续阅读