一、从 2024 年诺贝尔化学奖谈起
2024 年诺贝尔化学奖得主都不是来自化学专业。其中 David Baker 从事多年蛋白质设计研究，包括一些模型和传统生物信息工具，类似于现在的生成式场景。另外两位得主来自谷歌旗下的 DeepMind 团队，该团队主要专注于蛋白质生成领域，其另一重要成就是之前在围棋比赛中战胜人类的 AlphaGo。
蛋白质的业务价值非常大，几乎所有生物公司都无法绕开这个领域，都会做一些蛋白质相关的应用或者模型二次开发。同时我们也发现这个领域对于计算资源的消耗是非常大的，一个模型就需要消耗多个 CPU 或者多张卡来处理一个请求，其并发延迟远超传统的推荐搜索模型。举个例子，一个蛋白质结构生成预测，如果需要预测一个 100*1000 个序列的混合物，就得 30 分钟，一天只能计算几十的。蛋白质的序列生成就更加耗时，可能一个显卡设计一个就需要几个小时。所以在蛋白质研究领域，计算能力有很大的提高空间，接下来就将分两部分来介绍这个场景的优化。
二、加速蛋白质结构预测性能
首先来介绍 DeepMind 团队关于加速蛋白质结构预测的工作，其主要思想是基于 Ray 的 Workflow，实现高效异构调度。
AlphaFold 有很多版本，这里介绍一个比较颠覆传统的版本 v2.3。上图中第一幅图，输入是人体的氨基酸序列，输出是预测的结构，中间模型经过多次转化，比如第一步是预处理，之后进入模型，类似 transformer。第二幅图中是精简后的架构，用了一些生物学工具，效果很好但是效率低，非常慢。通常人们喜欢多个模型叠加拿到更好的效果，但是这样就会更慢。得到结果后，结果可能与生物理解不一致，需要微调，去掉完全不符合生物学属性的蛋白质坐标。
我们深入细节，发现第一步是最慢的，主要是来自于一个传统声讯工具 MSA，其需要消耗 1T-2T 内存。最后一步是 Relax，是一个混合计算，既需要 GPU 也需要 CPU。GPU 需要几百兆显存，CPU 时高时低，如果不做优化，是非常浪费时间的，模型结束后给到 Relax，GPU 利用率就会较低，因此需要一个异构的分布式调度，以充分利用资源。
DLmind 是谷歌的一个云原生解决方案，业界使用广泛。其核心思想是用 Kubeflow 的 pipeline 把整个链路全部异构起来，每个是一个单独的节点处理，这样预处理不需要 GPU，中间的串行可以解耦整个模块，并且利用其扩展性，分布更为灵活。其中的串行调度，类似批处理模式，但是整个延迟还是很严重的，有一定的局限性，比如吞吐很慢，基本无弹性能力。当有多个请求，多个 batch 时资源也是固定的，并且这些 batch 都必须在第一步预处理结束后才能触发（第一步耗时最多）就导致了下游有更多 GPU 是无法利用的。
图 3 中的红色部分是其中的关键节点。a 和 b （串行，资源固定）前面已经介绍，下面讲一下 c 数据预处理，其与业务场景高度相关，这与传统推荐搜索场景是完全不一样的。这个数据预处理过程需要 30 多分钟，这是因为该过程是一个传统的生信工具没有太多深度优化，需要预处理之后，将分布式的数据以及MLE数据转入到内存中才能加速后面的处理。传统的 Kuberflow 没有办法进行这样的预处理，但是有些方法可以绕开这个限制，比如部署一个 MLE Server，但是这种方案复杂度比较高。所以我们想是否可以应用 Ray 来提供一种高效率、快吞吐的方案，因为 Ray 在离线计算已经有较好的应用。
上图展示了我们设计的架构，采用 Ray 的 Workflow 方案。这个架构有两大特点，一是流式调度，二是高效构图。我们将其拆分成几个节点，都是对应的 flow 的 node 节点，可以灵活构图，灵活构图的好处就是每个节点均可插拔，每个节点可以无缝替换。
第二个核心设计是，考虑到 MSA 的 node 预处理非常慢，因此设计为 Actor Pool 初始化一个节点，预处理做好，推理时其已经是一个预热好的节点，这样可以从 30 分钟优化到 2 分钟，效果非常可观。
第三步就到了一个 GPU 节点，类似传统语言模型，将其作为 model node，如果机器足够多，可以自动扩缩容，无需人为定义资源。
最后就是 Relax 生信工具涉及到 CPU 和 GPU 混合运算，优化方案有两种。一是可以将其任务拆分很细，把 GPU 和 CPU 运算进行分离，但是这种方案需要较多的深度开发，开发难度较大。利用 Ray 支持小而一的调度，所以在每个 GPU 节点我们拆分更细，不用做较大改动就可以大幅提高性能。
结果输出本身就是端到端，Ray 支持通用节点不会导致 OOM（超内存，out of memory）。节点和节点间，请求和请求之间都是可以同步进行的。另外生信场景数据交流均是到 G 级别的，这种如果用传统解决方案，只能使用分布式存储系统，频繁的 I/O 就会有一定的瓶颈。这里我们用到 Ray 的一个共享 Object 之间传输有一些传统架构，就不会有 I/O 的瓶颈，整个吞吐就会非常高效。
Ray 在 AI 时代之所以应用很广，一个原因就是其 Python 友好，能接入 Python 对库，很多算子优化均可以用 Python 程序进行封装完美接入，模型也可以做更多的优化。最后真个过程可以从 30 分钟减少到 60 秒，这个在业界是比较领先的。
回到业务的核心，利用 Ray 可提升执行效率，并且由于 Ray 的可扩展性，再加上 Ray 整个架构是一个非常好的解耦架构，因此可以降低运维成本，提升合作开发的效率。另外其底层还是 K8S，我们不需要关注 GPU 和 CPU 节点的情况，对于创业公司（人员不足的情况下）是非常友好的。
核心就是 Workflow 的属性，解决了延迟和吞吐的问题。
三、加速蛋白质生成设计性能
下面介绍蛋白质生成场景的应用，这里用到了 Ray 的另外一个属性，Ray data，这是一个非常高阶且实用的属性。如上图左侧所示，生成场景主要是包含一个模型的扩散，每一次都会将一个模型扩散成多个模型，一步一步扩散下来，就会需要非常多的处理时间。从一个模型可以扩散到上千级别的，生命科学和其他领域不一样的就会有传统生信工具，需要去掉不符合生物学特征的数据。如果不做任何优化，进行一次设计，就需要 2 个小时。最后的 Relax 场景，看似很快，但其实是一个单核场景，堆积起来，1000 个模型需要 1000 个核就会导致整个处理非常慢。
如果可以做到右边所示的调度流程，就可以完美 overlap 并行运算，是理想中的最优结果，这个方案非常完美，但如果想要用自建的方式来实现还是比较难的，会涉及到多卡多 CPU，需要关注各种分布式的通讯调度异常，执行起来难度非常大。
所以我们引入 Ray 的解决方案，Ray data。Ray data 是一个 high level API，它是一个简单高效的执行器，对于处理串行计算是一个非常不错的选择，但是不适合结构预测有多个分支的任务。
上图中给出了一个示例代码，第一步是结构预测，结果出来后用 Ray data 将所有流程串起来。这个代码是一个非常优雅的解决方案，少量代码即可实现。实际运行时有一定的问题，主要是第一步会耗时很久。所以我们做了一些升级，第一个是利用典型的流式输出，完美的 overlap。
第二步是 Relax filter 不需要完整的一张卡，我们利用 Ray 会自动管理底层资源，并行度范围可以自行设置，最小 1 张卡。Ray data 会根据数据量自动扩容，可大幅减少运维成本，卡更多就可以处理更多的请求。
实际业务中，数据输出量过大就容易导致 OOM，而数据量过小，则会过于碎片化，都不是完美的解决方案。在这种融合计算架构中使用 Ray 的接口可以有效避免这些问题。
整个运行时间 Baseline 是 2 个小时级别，优化后是 1-2 分钟，对生命科学领域加速模型处理的意义是十分重大的。
生命科学不仅仅包含蛋白质的处理，还会有 RNA、DNA 等。此外除了离线 batch 任务，还有在线任务。我们也有自己的大语言模型底座，需要微调出来的模型就又不一样，所以特点是是业务多，模型多。实际问题非常复杂，效率优化是非常重要的。
整个过程非常复杂，需要不断模型调优，加上创业公司人员不足，不同模型使用的语言，接口都不同，会需要很多重复建设。同时性能也有一定的问题，如果每个人都有自己的模块，就无法复用，无法满足高效执行，低吞吐。所以我们希望设计一个新的架构，可以同时减少运维操作，并提高性能。
四、Ray 融合计算架构
基于上述背景，我们设计了基于 Ray 的融合计算架构，将所有事情都在 Ray 中完成，每个接口可插拔，底层可以统一，优化就可以同步进行，具体架构如下图所示。
融合架构的理念就是所有事情都在 Ray 中进行。最下面的组件大部分是一样的。向上一层是私有化云原生的部署，上面做了一个封装使得 Ray 上不会感知到是私有化部署还是云原生，这里我们做了一个 Ray 的抽象。这里面的场景其实很丰富，每一个小的模块包含几十个模型。在此之上我们做了一层封装，将相似模型做一个统一接口，做成 task 或 actor，称为统一融合引擎。我们也做了一些调度，比如 Ray server （在线服务），Ray data （串行 pipeline），对于自定义等更复杂的场景就用 Ray 的 workflow，仅需要用原生 Python 语言去嵌入各个节点。
很多场景下并非一次调整就能得到理想结果，而是需要基于反馈反复调整，进行多模型优化。
基于上述基础架构，可以实现基于 Ray 进行积木化组装模型应用。
基于 Ray 可以实现：
- 高效构建：Python 友好，可以统一编程语言；分门别类，统一接口；统一调度，减少构建 pipeline 成本。
- 高性能执行：可以弹性自动扩缩容；Stream overlap 执行；融合单节点、单模型优化。
- 低成本运维：既能解决私有化也能解决云原生，并且对业务屏蔽，甚至无需了解 Ray 就可以进行模型推理。
