DRUGAI
人工智能模型正在将生物学和化学编码带入高通量、高质量的计算机辅助药物开发新阶段。然而，这些模型的训练日益依赖于大规模计算资源。最近的蛋白质语言模型（pLM）训练需要数百个图形处理单元（GPU）。为此，NVIDIA推出了BioNeMo框架，旨在支持计算生物学和化学领域的AI模型在数百个GPU上的高效训练。该框架采用模块化设计，可将数据加载器等组件轻松集成到现有工作流中，并开放社区贡献。
通过具体应用案例，NVIDIA展示了BioNeMo框架的技术特点，例如pLM的预训练和微调。在256个NVIDIA A100 GPU上，BioNeMo框架能在4.2天内训练一个包含30亿参数的BERT型pLM模型，处理超过一万亿个标记。BioNeMo框架完全开源，供全球研究者免费使用。
2024年11月18日，亚特兰大——在SC24大会上，NVIDIA宣布其开源的BioNeMo™框架正被全球生物制药与科技生物公司、学术先锋和AI研究人员广泛采用，以推动药物发现和分子设计进步。研究人员需要专门的生物分子模型和数据集，以实现大规模洞察，加速治疗药物的设计。BioNeMo框架通过一系列加速计算工具，为生物分子研究提供指数级扩展能力，将生物制药领域的计算提升到一个全新水平。
背景
人工智能（AI）领域正在彻底改变计算生物学和化学的发展，并促成了2024年诺贝尔化学奖的颁发，该奖项表彰了在计算蛋白质设计和蛋白质结构预测领域的突破性进展。这些成就得益于高通量实验生物学的迅速发展、大规模序列数据库的建立以及大型神经网络架构的显著进步，这些都推动了大型语言模型（LLMs）的发展。以蛋白质语言模型（pLMs）为代表的技术，通过在庞大的蛋白质序列数据库上训练，在结构预测、功能预测和从头蛋白质设计等任务上达到了前所未有的精度。
随着对LLM扩展规律的深入理解，pLMs的规模持续增长以提升性能，例如ESM3模型拥有980亿参数，训练所需浮点运算量达到1024。虽然改进训练数据质量也可能提高模型性能，但我们相信前沿的pLMs将继续向更大的规模和更高的训练需求发展。除了蛋白质设计，LLMs还推动了化学和材料发现的新技术。
AI驱动的分子研究迅速增加了对专门为药物开发设计的机器学习框架的需求。尽管PyTorch、TensorFlow和JAX等传统工具在推进深度学习方面取得了成功，这些框架并未针对复杂的生物分子数据和药物开发所需的大规模模型进行优化。虽然这些框架在通用深度学习任务中表现出色，但在适应诸如蛋白质结构预测或化学反应建模等任务时，往往需要复杂的改造，导致在灵活性和性能之间做出权衡。高效训练超大模型需要精心设计的数据和模型并行策略，以及精细使用CUDA融合内核以最大限度减少内存使用并提高训练吞吐量。高性能数据加载器也是预训练流程中的关键部分，因为要确保GPU完全利用，必须高效完成预处理和内存访问。随着模型训练规模的扩大，尽量减少能源消耗并最大化模型性能变得尤为重要。同时，这些技术的普及对于较小的研究团队和化学、生物研究领域的科学家来说也至关重要。
BioNeMo框架是一套开源的Python软件包，旨在简化生物分子和化学领域高性能AI模型的训练过程。这些库构建于NVIDIA NeMo和Megatron-LM之上，优化了训练和推理的吞吐量，并包括用于蛋白质序列采样和单细胞数据的高性能数据加载器。本次发布的BioNeMo框架包括ESM-2和Geneformer模型，我们展示了其相较于PyTorch实现的训练吞吐量提高了两倍以上，并在多设备训练中实现了接近线性的扩展，支持多达256个GPU。本文还探讨了BioNeMo框架的未来发展方向以及外部贡献者如何扩展框架，以加速更多类型模型的开发。
BioNeMo框架概览
NVIDIA BioNeMo框架是一套专为计算药物开发设计的编程工具、库和模型集合。它通过提供针对特定领域优化的模型和工具，大幅加速了构建和调整生物分子AI最耗时且成本最高的阶段，并能轻松集成至任何基于GPU的计算环境中。
模块化架构
BioNeMo框架采用子项目结构设计，便于用户专注于特定功能的开发和贡献。整个代码库可以作为一个整体安装，也可以按照特定用途和范围分段安装为独立组件（称为子项目）。这种模块化架构增强了扩展能力，支持以下两种贡献方式：
广度扩展：创建新的子项目。
深度扩展：扩展现有子项目。
核心与子包设计
BioNeMo框架的核心是 bionemo-core，它围绕PyTorch和Lightning框架构建，提供关键接口、通用数据处理模块和模型构建模块。所有其他子项目都依赖于此核心。例如：
bionemo-llm：基于NVIDIA NeMo和Megatron库，提供生物领域定制化和基础模型架构支持。
bionemo-geometric：正在开发对图神经网络（GNNs）的支持，基于PyTorch Geometric，展现了BioNeMo生态系统的创新能力。
每个子包都封装了特定的模型和数据加载器，用户可以根据自身需求，选择性地集成相关组件。这种模块化设计不仅简化了定制化流程，还提高了整体可用性。
开发与集成
为了便于集成到私有项目中，BioNeMo子项目均被设计为可通过 Python Package Index (PyPI) 安装。借助提供的Docker镜像，用户可以快速搭建完整的开发环境。此外，代码库内置的工具可以帮助用户轻松创建并将自定义Python项目整合到现有工作流中，从而简化开发流程。
当前支持与特色
BioNeMo框架目前通过NVIDIA NeMo Megatron支持大规模BERT模型的构建。NeMo Megatron是一种可扩展框架，用于开发自定义的LLM、多模态和语音AI。BioNeMo框架利用NeMo Megatron实现生物分子BERT模型的支持，为用户提供高度可扩展的解决方案。
通过BioNeMo框架，研究人员可以在生物分子和化学领域中快速部署高性能AI模型，加速药物开发的创新进程。
功能概述
BERT模型架构支持
BioNeMo框架支持两种BERT模型实现：
ESM-2：用于蛋白质序列建模，封装在 bionemo-esm2 子模块中。
Geneformer：用于单细胞表达建模，封装在 bionemo-geneformer 子模块中。
用户可直接安装并使用这些模块进行模型训练、微调和推理。每个子模块都有独立的 src/ 目录，包含数据模块、分词器和模型实现，便于用户定制，例如替换数据加载器或调整模型架构。此外，框架支持数据、模型、张量和流水线并行，用户仅需调整配置参数即可轻松实现模型的扩展和参数增加。
基准测试与结果
单设备性能
ESM-2 650M参数模型：在NVIDIA A100 GPU上，BioNeMo的最大批量为 46（对比基线16），训练吞吐量为基线的 1.47倍，浮点操作利用率（MFU）达到 59.2%（基线40.1%）。
分布式性能
BioNeMo框架在256个GPU上达到单节点扩展性能的 96.9%，可在 4.2天内完成1万亿标记的训练，比基线方法（8.8天）快一倍。
单细胞数据加载器 (SCDL)
特点与优势
提供独立的PyTorch兼容数据集类，支持高效处理单细胞数据。
比传统AnnData加载器快 1.1至2.75倍。
支持使用压缩稀疏矩阵存储数据，能快速处理超大数据集。
模块化设计支持从AnnData格式转换，并可定制扩展其他数据格式。
WebDataModule
BioNeMo实现了WebDataset与LightningDatamodule的无缝集成：
WebDataModule：将WebDataset数据分片映射到LightningDatamodule中。
PickledDataWDS：支持从Python pickle对象生成WebDataset数据分片。
这种设计简化了高性能数据流的加载流程，同时支持本地和云端存储。
尺寸感知批量处理 (Size-aware Batching)
尺寸感知批处理器
动态调整批量大小以优化GPU内存利用，避免OOM错误。
包括内存峰值收集、内存分配预测以及基于内存预测的动态批量生成。
桶批量采样器
按样本尺寸将数据分桶以减少填充开销，提供均匀分布的采样结果。
支持用户定义的尺寸范围和桶大小，灵活高效。
基准测试结果
对比三种方法（静态批量、MiDi动态批量和尺寸感知批处理）：
静态批量：大样本因OOM错误被跳过。
MiDi方法：偏向常见尺寸样本，导致采样偏差。
尺寸感知批处理：样本尺寸分布均匀，无采样偏差，填充需求最少。
尺寸感知批量处理在数据加载效率和内存利用率方面表现出色，适用于GNN等复杂数据处理场景。
讨论
社区贡献
BioNeMo框架作为开源软件，得益于社区的积极参与，不断在生物分子领域扩展功能，同时优化了在GPU硬件上的性能。多个机构已通过实际应用为框架做出贡献，包括以下典型案例：
Dyno Therapeutics：这家专注于基因治疗载体开发的公司，在ESM-2模型中实现了参数高效微调功能，允许冻结部分网络层，仅对最相关的部分进行微调。这种方法避免了传统全模型微调带来的高计算成本和过拟合风险，同时显著加快了训练速度，提高了资源利用效率。
Flagship Pioneering：通过其Flagship Intelligence项目，添加了蛋白质逆折叠模型 RL-DIF。与传统预测序列生成结构的折叠模型不同，RL-DIF反向生成能折叠成特定目标结构的序列，为设计符合精确结构需求的新蛋白质提供了强大的工具，在药物发现和生物技术领域开启了广阔前景。
Relation Therapeutics：采用实验室与AI联动的方法，开发了新的DNA数据加载工具（DNADL）。该工具简化了常见基因组学数据格式的加载流程，支持基于特定基因组区域或已知变异序列的AI模型训练，为精准生物学研究提供了更高效的解决方案。
Weights & Biases：集成了实验管理工具，使BioNeMo框架用户能够自动跟踪训练进程，包括跨多次重启的实验记录，优化模型开发与调试流程。
云计算资源扩展
BioNeMo框架的云端兼容性使其能够满足不同项目需求，尤其是无需本地基础设施的情况下处理大规模数据集。通过云计算资源，用户可访问和定制预训练模型，轻松完成从概念验证到生产部署的全流程。
与AWS的集成
支持Amazon EC2实例（如搭载NVIDIA GPU的P5实例）和S3存储，提供可扩展的计算资源和数据管理能力。
通过AWS Batch高效扩展计算工作流，支持大规模模型迭代优化。
A-Alpha Bio案例
这家专注于蛋白质-蛋白质相互作用（PPI）预测的公司，利用BioNeMo框架在AWS平台上提升了药物开发的计算能力：
使用NVIDIA H100 GPU的EC2 P5实例，将推理速度提高 12倍，两个月内处理了 1.08亿次推理调用。
通过扩展计算容量，评估的蛋白质结合预测量增加了 10倍，显著扩大了候选药物范围。
通过AWS容量预留选项（如EC2 Capacity Blocks），获得高性能资源的可预测访问，降低了实验成本，加速了治疗设计流程。
将探索重点转向计算测试，减少湿实验次数，从而更有效地研究复杂蛋白质修饰。
共享贡献
A-Alpha Bio贡献了其用于零样本蛋白质设计的ESM-2笔记本代码，现已集成到BioNeMo框架代码库中。该功能利用掩码语言建模预测特定位置的氨基酸替换，为研究人员高效设计和优化蛋白质提供了强大的工具。
通过社区的共同努力和云端资源的高效利用，BioNeMo框架在加速药物开发、推动AI模型创新以及降低研发成本方面展现了巨大的潜力。
结论
BioNeMo框架通过一套强大的开源GPU加速工具，为药物发现工作流提供了重要支持。其在训练吞吐量、可扩展性和模块化设计方面表现卓越，展现出在不同生物分子领域推动进一步创新的巨大潜力。
未来发展方向
框架优化
持续优化API设计，提高模块化程度，增强组件级的复用性、定制化能力以及实验灵活性。
鼓励用户对框架级功能进行贡献，同时允许用户独立共享和维护基于框架开发的子模块包。
动态模块管理
BioNeMo框架将保持子模块集合的动态性，随着更优抽象方案的出现，子模块包的数量可灵活扩展或调整，以确保功能实现的高效与简洁。
聚焦核心优势
优化大规模模型在训练、推理和微调场景中的性能表现，特别是在单设备无法容纳完整模型的情况下提供强有力支持。
不断推动前沿技术的发展，确保在处理大规模模型任务中的领先地位。
生态系统集成
借助NVIDIA内部其他框架（如NeMo、Megatron和Transformer Engine）的改进，BioNeMo框架将受益于整个NVIDIA软件生态系统的性能和易用性提升。
通过这些努力，BioNeMo框架将进一步巩固其在生物分子AI开发中的核心地位，并为药物开发领域的技术突破提供持续支持。
整理 | WJM
参考资料
https://arxiv.org/abs/2411.10548
https://nvidia.github.io/bionemo-framework/
https://github.com/NVIDIA/bionemo-framework
