AI真·炼丹:整整14天,无需人类参与

  金磊梦晨发自凹非寺

  量子位公众号 QbitAI

  现在制药这事,人类要靠边站了。

  坐标苏州,这是一个 1600 平的制药实验室,它的“打开方式”是这样的:

  门口,没有人

  走廊,没有人

  实验室,也没有人

  相比以往充斥着科学家、研究员的实验室,它更多的是把机械臂和 AI 系统塞了进去,主打的就是一个全自动化

  或许好奇的小伙伴就要问,这样的实验室能干嘛?就是为了自动化而自动化嘛?

  事情当然没有那么简单,你瞧见的只是无人的操作,但在背后,AI 做的可远远不只是替代人工的实验室操作那么简单,而是:

14 天内完成靶点发现和验证,还是全自动化干湿实验闭环的那种。

  要知道,这个过程要放以前,可是需要足足2-3 年才能完成……

  而且更为精细化的工作,例如样本处理细胞培养化合物管理高通量筛选新一代测序高内涵成像等等,不论是单一任务还是“联动”任务,机器都可以在 AI 的控制下轻松接手。

  △用 Echo 650T 制备检测板

  △用 NovaSeq 6000 测序

  这便是来自全球 AI 制药第一梯队的“选手”——英矽智能(Insilico Medicine)的第六代智能机器人实验室,也是全球首个用 AI 参与决策的生物学实验室。

  而在它背后驱动这一切的 AI 大脑,则是一个叫做 PandaOmics 的平台,可以根据实验的进程自主做决策、下达指令。

  若是把这个 AI 平台单拎出来,它更是囊括了 20 多种预测模型和生成生物学模型,还包含遗传学、蛋白质组学、甲基化数据、文本文献和科研基金等海量数据,用以支持专业的靶点识别、分析和排序、适应症探索等生物学研究。

  甚至已经有高中生用 PandaOmics 发现了药物新靶点,并且研究成果还登上了国际学术期刊!

  而且除了 PandaOmics 之外,英矽智能在人工智能制药领域拥有端到端的药物发现平台 Pharma.AI,其中专注于化学领域的 Chemistry42,还可以针对给定靶点从头设计具有特定属性药物理化性质的新型小分子。

  这一切都可以在几小时到几十小时内完成,且支持并行运行多个任务。

  以及英矽智能还将科技圈最潮的大模型也融入进来,在 Pharma.AI 的架构上推出Copilot 系统,让你只要会对话就能使用专业的 AI 制药平台。

  由此可见,现在 AI 不仅是把制药这件事变成了“自动驾驶”模式,更是狠狠地把门槛打下去效率提上来

  AI 制药的流程和工作是方便了,但随之而来的一个问题便是:如此大的工作量,算力,又是如何解决的呢?

  科学计算与 AI,CPU 都在发力

  对于上述的问题,包括英矽智能在内的 AI 制药头部力量们不约而同的选择了相似的解决办法:

  充分利用所有可以用、值得用的科学计算与 AI 算力平台。这种平台可不是你想象的那样被 GPU 制霸,相反,其中的 CPU 用量更大,尤其是英特尔的 CPU

  为什么要选择英特尔?

  首要的一个原因,就是英特尔供企业计算及科学计算使用的主力 CPU,即至强® 可扩展处理器系列产品,一直都是物理计算——无论是昔日计算机辅助制药,还是今天 AI 辅助制药都非常依赖的科学计算应用的关键承载平台。

  另一方面,就算是把应用的主题从相对传统的制药相关的科学计算任务,切换到更偏 AI 的应用上,英特尔也算是颇有建树,这一点从它对以 AlphaFold2 为代表的开源蛋白质预测模型的支持上就可见一斑。

  △AlphaFold2 基本架构

  ⾸先,AlphaFold2 整个端到端的处理过程,涉及⼤量复杂多样的计算类型。从早期的数据收集、特征提取等预处理阶段,到基于深度学习的蛋⽩质结构预测,再到后续的结果分析,这是⼀个⾼度异构的⼯作负载。

  ⽽英特尔® ⾄强® 可扩展处理器可以轻松胜任这一系列多样化的任务。以⾄强® CPU Max 系列处理器为例,它采⽤全新微架构、更多内核(最⾼达 56 个),能以更⾼频率和更⼤缓存,去应对⾼通量的预处理和后处理⼯作。

  它在内存和输入/输出(I/O)子系统性能上有着显著的增强,还结合大容量末级缓存使 AlphaFold2 推理过程中关键的张量吞吐获得了大幅提升。

  △英特尔® 至强® CPU Max 系列处理器

  其次,由于 AlphaFold2 所采⽤的深度学习模型规模巨⼤,推理过程中的张量运算不仅量⼤,且维度极⾼。这就要求承载平台具备强⼤的 AI 运算加速能⼒。

  在这⼀点上,新款⾄强® 系列处理器内置的英特尔® AMX(⾼级矩阵扩展)技术,可以显著加速⼤规模矩阵乘法运算。

  在 FP32/BF16 混合精度计算下,其理论峰值可达每时钟周期 1024 次乘加操作。针对 AlphaFold2 推理任务中所需的大量矩阵运算操作,AMX_BF16 能在保持较高精度的同时,提高计算速度并减少存储空间。

  △AMX_BF16 推理优化带来更低内存占用和更大输入长度

  另⼀⽅⾯,AlphaFold2 因其⾼维张量运算和⻓序列并⾏计算,在推理过程中常⾯临超⼤内存需求,不光影响推理速度,还会限制更⻓蛋⽩质序列的预测。

  为此英特尔从软硬协同的方式给出完整解决方案。

  一面是提升内存容量和带宽。解决方案中,英特尔® ⾄强® CPU Max 系列处理器除支持 DDR5 内存外,还集成了 HBM(⾼带宽内存)。单颗处理器的 HBM 容量⾼达 64GB,且具有高达 460GB/s带宽。

  另一面是提供了多种降低内存的软件优化方法。如面向 PyTorch 对张量计算原语(Tensor Processing Primitives,TPP)技术进行扩展,以及切分 Attention 模块和算子融合的推理优化方案,帮助 AlphaFold2 在通用矩阵乘法计算中所需的内存峰值大幅降低。

  △热点算子与融合效果

  经过一系列加强和优化后,最终效果如何呢?

  如图所示,在基于至强® CPU Max 系列处理器的优化流程中,每个优化步骤获得的提升累积后,获得了相对于基线性能(对比组1,基于第三代至强® 可扩展处理器,未实施优化)高达 33.97 倍的通量提升。

  根据测算,性能提升中的 74% 源自预处理阶段的高通量优化,26% 要归功于对推理过程的优化。

  此外,在同样开启 IPEX(面向 PyTorch 的英特尔® 扩展优化框架)的情况下,相比对比组2(基于第三代至强® 可扩展处理器,但实施过优化),方案在升级使用至强® CPU Max 系列处理器后,其内置的 HBM 内存、英特尔® AMX 的加成,则带来了 48.3% 的性能提升。

  △切分 Attention 模块和算子融合的推理优化方案

  而且值得一提的是,在一项基于某公有云服务的测试中,基于至强® CPU 平台构建的 AlphaFold2 解决方案还在性能上获得了远优于某高端 GPU 平台的表现,同时也优于由 CPU+GPU 混合构建的方案。

  这可是一个非常难得的成绩——毕竟过去在很多 AI 应用的测试或实战中,CPU 能有接近或媲美 GPU 的表现就已经算是成功,而AlphaFold2 上至强® 平台则实现了性能+蛋白质预测序列长度的全面反超

  现在还剩下最后一个问题,多个蛋白结果的解析模型 AlphaFold2 Multimer。

  也就是从预测单个蛋白质三维结构,发展到了对多个蛋白质分子之间的相互作用及所形成的复合体结构进行预测。

  CPU 在这一演变过程中的支持力度如何呢?

  答案是不用担心!

  基于英特尔® 架构的 AlphaFold2 解决方案同样也面向 AlphaFold2 Multimer 的管线结构进行了优化与验证。

  虽然后者的管线结构已根据蛋白质复合体结构预测的需求进行了调整,但英特尔 AlphaFold2 上的优化方案,在被用于 AlphaFold2 Multimer 时同样有效。

  △面向 AlphaFold2 Multimer 模式的方案实现

  CPU 加速新药发现不是梦

  回顾以往,研发⼀种新药动辄需要 10 年时间,投⼊20 亿美元才能起步。

  ⽽在 AI 的助⼒下,这⼀成本正⼤幅降低。以英矽智能为例,它们进展最快的项目仅⽤18 个⽉就找到了治疗特发性肺纤维化(IPF)的潜在全球首创候选药物并通过实验验证,总成本约为 280 万美元。

  展望未来,随着 AI 技术的进⼀步发展渗透,它必将重塑制药业的创新模式,让新药研发变得更加⾼效、精准、经济。而在这一进程中,相关的科学计算及 AI 应用任务,依然需要有强大的算力支撑。

  从英矽智能等公司的实践来看,以⾄强® 处理器为代表的 CPU 平台,正凭借其在性能、成本、⽣态等⽅⾯的独特优势,成为推动 AI 时代制药创新的重要“引擎”。

  这也预⽰着,在 AI 改变众多⾏业的当下,CPU 加速 AI 应用落地,帮助用户节支增效以及推进其技术和业务创新的脚步从未停止。

  AI 让新药研发进⼊“⾃动驾驶”模式,⽽英特尔® ⾄强® 处理器则提供了它所需的源源不断的动⼒。

  在这种合作模式下,AI+ 制药还将擦出怎样的⽕花,就很值得期待了。