相比现有方案，Colossal-AI 团队联合百图生科的这项研究，推理速度最高提升约 11 倍！
近日，Colossal-AI 团队( https://github.com/hpcaitech/ColossalAI )联合百图生科成功加速蛋白质单体与复合物的结构预测，行业内最前沿的同时支持蛋白质单体（Monomer）与复合物（Multimer）结构预测的 xTrimo Multimer 模型已免费开源，相比现有方案，推理速度最高提升约 11 倍！
本次解决方案背后优化加速技术来自 Colossal-AI，这是一个旨在全面助力 AI 大模型工业化应用的通用深度学习系统。xTrimo Multimer 开源版模型通过引入 Colossal-AI 的大模型优化技术，显著降低蛋白质结构预测模型训练和推理的时间和经济成本，提高蛋白质结构预测模型领域的设计和部署效率，是 Colossal-AI 系统在 AI 大模型 & 生物医药领域落地应用的重要实践，使得在该领域训练和使用更大的模型成为可能。
该项目已在 GitHub 上免费开源：https://github.com/hpcaitech/ColossalAI/#xTrimoMultimer
Colossal-AI 是一个面向大模型时代的通用深度学习系统。自开源以来，Colossal-AI 不仅在数个月内获得近五千颗 Github Star，多次登上 GitHub Trending 榜首，相关解决方案也已成功在生物医药、自动驾驶、云计算、零售、芯片等行业知名厂商落地应用，广受好评。此次携手百图生科，利用双方在高性能计算和生物计算领域上前沿的技术积淀，提出蛋白质单体与复合物的结构预测加速方案，进一步共同打造繁荣的生物计算生态，精准高效地探寻癌症、衰老等疾病的免疫规律，逐一攻破 AI 技术在靶点挖掘、创新药物设计落地应用上的难关，让更多疾病可预警、可控制、可治愈。
众所周知，蛋白质结构预测是结构生物学领域最为重要的课题之一，也是我们理解基因翻译和蛋白质功能的重要手段。蛋白质之间的相互作用，是蛋白质发挥生物学功能的重要结构基础。但是由于蛋白质的多级结构和复杂的相互作用，使得精确预测三维结构这一关键问题十分有挑战性。 近年来，深度神经网络的成功使得人工智能在各个领域得到了广泛的应用。自 DeepMind 发布的 AlphaFold 能够实现根据氨基酸序列精准预测蛋白质结构以来，生命科学领域掀起了使用人工智能预测蛋白结构的热潮。 AlphaFold 能够实现端到端的从序列直接预测单体蛋白质三维结构，但在生物体内，多数蛋白质是以复合体的形式来发挥功能的。因此，为了打破 AlphaFold 仅在预测单体蛋白方面表现优异的局限，DeepMind 后来也发布了用于预测蛋白复合体结构的 AlphaFold-Multimer 模型。 早在 2022 年 3 月，Colossal-AI 团队曾推出蛋白质结构预测模型 AlphaFold 的训练推理加速方案 FastFold，用更低的成本将总体训练时间从 11 天减少到 67 小时，在长序列推理应用中实现超 11 倍的提升。以此为起点，Colossal-AI 团队在蛋白质结构预测领域的技术革命正在延续。
针对蛋白质单体 (Monomer) 与复合物 (Multimer) 结构预测的难题，本次 Colossal-AI 团队联合百图生科提出行业内最新解决方案 xTrimo Multimer 开源版模型，能够更好地理解蛋白互作关系，从而提升药物研发平台中靶点分析、蛋白质结构预测和模拟以及高精准抗体设计等方面的能力。
此外，推理阶段昂贵的经济和时间成本使得 AlphaFold 模型相关的研究和开发十分困难，尤其是面对长序列推理时，计算复杂度和内存消耗都面临着巨大挑战。xTrimo Multimer 开源版模型针对 AlphaFold-Multimer 模型中的计算特性，进行了针对性的 CUDA 优化，并且进行了 Kernel Fusion。相较于 AlphaFold2 和哥伦比亚大学的 OpenFold，xTrimo Multimer 开源版模型在单卡推理上性能有明显提升 ，推理速度分别提高 1.58～2.14 倍和 1.14~2.23 倍。
xTrimo Multimer 开源版模型还支持超长序列的分布式推理。由于使用了动态轴并行（Dynamic Axial Parallelism）技术，xTrimo Multimer 可以高效地将计算和部分显存分配到不同的设备上，从而解决超长序列面临的计算和内存挑战。在 2K 到 3K 的序列长度下，xTrimo Multimer 使用多卡进行推理，相对于 OpenFold 和 AlphaFold 2 推理速度最高提升 8.47 倍和 11.15 倍，对比 Uni-Fold2.0 速度最高提升 4.45 倍。xTrimo Multimer 可以支持长达 4K 的序列推理，此时 OpenFold 和 AlphaFold 2 受限于显存无法完成推理，而 xTrimo Multimer 可以在 20 分钟左右完成。
百图生科首席 AI 科学家宋乐表示：“从对单个蛋白质结构的模拟，到对不同蛋白质之间的相互作用的识别，再到对蛋白质复合物的绘制，百图生科致力于解码、建模复杂人体免疫系统，开发突破创新药物，编程免疫系统，治愈多种免疫相关疾病。此次 xTrimo Multimer 开源版模型的发布，是我们携手潞晨科技 Colossal-AI 团队，借助其在高性能计算上的优势以及百图生科生物计算领域前沿的技术积淀，朝着百图生科 xTrimo 多模态生物计算大模型体系迈出的又一步。正如我们一直以来所努力的方向，精准高效地探寻癌症、衰老等疾病的免疫规律，逐一攻破 AI 技术在靶点挖掘、创新药物设计落地应用上的难关，让更多疾病可预警、可控制、可治愈！” 潞晨科技创始人尤洋教授表示：“ 此次 Colossal-AI 团队与百图生科合作的蛋白质单体与复合物结构预测最新方案，是面向大模型时代的通用深度学习系统 Colossal-AI 在生物医药领域应用落地的重要进展。未来我们双方将会继续在生物计算大模型上有更深入的合作，助力深度学习在创新药物研发上的应用和落地。” 随着人工智能对各个领域的高度渗透，各行业的智能化产业转型升级都离不开高性能 AI 的加持，潞晨对于 AI 大模型的高效训练和推理方案是未来 AI 市场的刚需。Colossal-AI 致力于构建针对 AI 大模型的完整生态系统，针对不同行业的需求提供智能化、一体化的高效灵活解决方案，切实有效解决企业 AI 大模型产品落地的关键痛点，为全产业链赋能。
