跳到正文
Databricks Blog·· 9 小时前AI 评分38

在 Databricks 上部署 RFD4-Proteina

Deploy RFD4-Proteina on Databricks

AI 导读

NVIDIA 与 Databricks 计划将蛋白质设计开放模型 RFD4-Proteina 集成进 Genesis Workbench,运行在提供 serverless NVIDIA GPU 算力的 Databricks AI Runtime 上。

正文 · AI 翻译

借助 AI 辅助开发和无服务器 NVIDIA GPU 计算,将 RFD4-Proteina 引入你的研究数据。

  • 通过 Genesis Workbench 体验首日模型支持,这是一个面向生命科学工作流的开放示例蓝图。
  • 使用 Genie Code、NVIDIA BioNeMo Agent Toolkit 技能以及编码智能体辅助来构建和调整流水线。
  • 使用 Databricks 正式可用的 AI Runtime,与企业数据、实验跟踪和治理一同执行受支持的 GPU 工作负载。

RFD4-Proteina 是 NVIDIA BioNeMo 开放模型,由 NVIDIA 与华盛顿大学蛋白质设计研究所联合开发,帮助科学家将有用蛋白质的构想转化为可测试的设计。NVIDIA 与 Databricks 正将该模型引入 AI Runtime 上的 Genesis Workbench,为开发者提供一条将蛋白质设计工作流与其研究数据连接起来的路径。

Databricks AI Runtime (AIR) 为受支持的工作负载提供无服务器 NVIDIA GPU 计算。计划中的集成将模型执行与 Databricks 上的研究数据和可重复工作流连接起来,减少评估和适配新模型所需的基础设施开销。

让模型真正投入使用,远不止下载权重。开发者需要准备输入、连接科学工具、配置执行并跟踪结果——同时还要保护专有数据。Databricks 将这些步骤整合到一个共享环境中,减少搭建和维护独立研究技术栈的需求。

“药物发现首先是一个数据问题,其次才是计算问题。NVIDIA 和蛋白质设计研究所正借助 RFD4-Proteina 等模型拓展可能性的边界,而我们携手确保生命科学团队能够在自有受治理的数据上运行这些模型,并与自己的实验相连,无需将敏感知识产权送出企业边界。科学归你所有。Databricks 平台确保这一切都在一个安全、受治理的环境中运行。”—Mike Sanky,Databricks 医疗与生命科学副总裁
“蛋白质设计的难点不在于生成候选分子,而在于挑选出少数值得在实验室中制备的候选。RFD4-Proteina 将蛋白质的序列与其完整原子结构一同设计,并且可以引导其朝向项目所需的特性,例如稳定性或结合强度。在 Databricks 上运行它,并与团队自身的实验数据放在一起,有助于科学家根据已有实验结果表明的结果来挑选这些候选分子。—Anthony Costa,NVIDIA 数字生物学总监

将模型带到你的数据旁

RFD4-Proteina 基于 Proteina-Complexa 所采用的方法构建,后者是 NVIDIA 用于设计蛋白质结合物的生成模型,可在单次运行中同时生成氨基酸序列和全原子结构。RFD4-Proteina 通过可选的奖励引导搜索扩展了这种联合生成方法,使研究人员能够将候选设计引导至热稳定性或结合亲和力等特性。  计划中的集成扩展了 Genesis Workbench 现有的 Proteina-Complexa 工作流模式,利用带有针对当前问题特定奖励函数的评分模型。Genesis Workbench 内置了两个使用 RFD4-Proteina 的工作流供您参考:抗体设计和疫苗设计。

rfd4-proteina-databricks

图示:Genesis Workbench 中使用 NVIDIA RFD4-Proteina 的抗体设计

Genesis Workbench 提供了一种简单的入门方式:其开放、模块化的蓝图展示了如何在 Databricks 上将生物模型与输入准备、候选评估和下游分析连接起来。开发者可以复用单个组件,或将实现适配到自己的科学方法中。 

团队可以将模型权重作为受治理的 MLflow 模型存储在 Unity Catalog 卷中,并在不同实验中复用。团队可以通过 MLflow 在 Unity Catalog 中注册微调后的模型,将模型版本与配置和工件关联起来。访问控制和网络配置应反映模型、评分工具和编码代理的执行位置。

通过 Genesis Workbench 使用 RFD4-Proteina,使 Databricks 用户能够更广泛地跨学科连接研究:基因组学分析可以为靶点提供信息,单细胞数据可以提供生物学背景,蛋白质设计或小分子工作流可以支持后续研究。Genesis Workbench 在共享环境中提供了这些能力的示例。 

使用科学代理技能构建

Databricks Genie Code 帮助开发者使用自然语言构建、调试和调整数据与机器学习工作流。NVIDIA BioNeMo Agent Toolkit 技能为生物模型和科学工具添加专门的指令、脚本和参考资料,包括结构预测、蛋白质设计、分子对接和基因组学分析。 

这些技能可以帮助编码代理准备输入、调用支持的工具、检查输出并组装多步骤工作流。开发者可以将这些指导与自己的数据约定、评估标准和实验跟踪要求相结合。

Genesis Workbench 还为基于 Claude 的编码代理提供技能,提供另一个入口:帮助开发者理解蓝图、为新模型(如 RFD4-Proteina)搭建集成脚手架,并将这些集成定制到特定用例。这可以加速首次实现;依赖检查、测试和科学审查仍是流程的一部分。

生命科学研究人员可以直接从 Genesis Workbench 利用预构建的 RFD4-Proteina 工作流,或使用 Genie Code 通过简单的英文提示创建 Databricks Notebook 进行临时实验,例如:

“使用 RFD4-Proteina 框架条件设计针对 PD-L1 的纳米抗体(VHH):保持真实的人源化框架(h-NbBCII10)固定,仅将 CDR 环设计到阻断 PD-1 的表位上。在 CDR 上使用 ProteinMPNN 运行 8 次迭代 × 8 个样本,并使用 Boltz 共折叠进行结合评估;奖励权重偏向结合(ipTM)、折叠置信度(pLDDT)、可开发性(溶解度、热稳定性)和低抗药抗体风险(MHC-II)。返回排名靠前的设计及其奖励轨迹和 top-K 复合物 PDB。”

Genie Code 协助实现,而开发者和研究人员通过审查代码、选择合适的工具并解读结果来保持控制权。

在托管的 NVIDIA GPU 上运行

AIR 简化了对 NVIDIA 加速的访问,团队无需自行配置和维护单独的 GPU 集群。开发者可以将 notebook 连接到 Serverless GPU 计算,选择 NVIDIA H100 加速器和支持的环境,然后运行其工作负载。文档中的设置包括在应用配置之前选择加速器和基础环境。 

这减少了配置 GPU 基础设施和维护基础软件环境的工作量,不过团队仍需验证特定于模型的库和依赖项。

对于支持的生物基础模型,团队还可以调整 BioNeMo Recipes,使用 AIR 的 serverless GPU 在专有数据上进行训练或微调。BioNeMo Recipes 提供了使用 PyTorch 和 Hugging Face 等熟悉框架进行加速训练的参考实现。 

Lakeflow 编排、MLflow 实验跟踪和 Unity Catalog 治理将执行与可重复的工作流和企业级控制连接起来。科学家可以更专注于生物学,而开发者和平台团队则减少管理基础设施的时间。 

从实验到部署

Databricks 上一个实际的模型生命周期示例可以包括:

  1. 将模型权重存储在受治理的 Unity Catalog 卷中以便复用。
  2. 在 AIR 上生成候选分子,或在模型和配方支持的情况下进行训练和微调。
  3. 使用 MLflow 跟踪参数、指标和工件,并在 Unity Catalog 中注册可部署的模型。
  4. 将兼容的模型部署到 GPU Model Serving 端点,或在更适合工作负载时编排批量执行。
  5. 通过 Genesis Workbench 示例或你自己的管道将输出连接到研究工作流。

对于支持的自定义模型端点,缩容到零可以在不活跃期间减少预置容量。冷启动会增加延迟,并且当端点重新扩容时 GPU 容量无法保证。Databricks 建议对需要一致可用性或响应时间的工作负载禁用缩容到零。

团队可以选择适合研究任务的执行模式,在运行时、成本和可用性要求之间取得平衡。

在整个工作流中,团队可以保留比较实验所需的配置、模型版本、评估结果和工件。计算评估支持候选优先级排序;实验室验证仍然必不可少。

开始使用

探索 Genesis Workbench 仓库,从一个目标、一个明确的设计目标和一个小型实验开始。使用示例实现和可用的 agent 技能来调整工作流,然后在 AIR 上执行并跟踪,之后再扩展。

来源:Databricks Blog · databricks.com