在Modal上构建强化学习定理证明工作流

AE Studio最近在Modal平台开展了一项引人注目的研究，旨在通过强化学习训练语言模型进行数学定理证明。他们系统比较了两种方法：当前主流的组相对策略优化（GRPO）和一种受自然选择启发的进化策略（ES）。研究使用了Lean形式化验证语言作为奖励信号——模型生成的Lean代码由编译器验证正确性，从而提供明确的奖励信号。

为了高效运行实验，AE Studio选择了Modal平台，该平台提供了并行GPU计算、沙盒隔离和持久化卷存储等关键功能。具体工作流程分为三个部分：证明生成使用vLLM在GPU上运行，证明验证使用Lean编译器在CPU上通过Modal沙盒隔离执行，协调进程管理整个训练循环。ES方法通过创建一群略有不同的模型副本，评估它们并更新原始模型，而GRPO则基于组内相对性能进行梯度更新。

在实现细节上，AE Studio充分利用了Modal的特色功能。他们为每个计算角色使用独立的镜像，而不是将所有依赖打包到一个臃肿的环境中。GPU工作节点专注于推理，Lean沙盒专注于验证，协调器保持轻量。并行GPU扇出通过Modal的.map()实现，每个扰动评估只需要当前检查点、定理批次、一个扰动种子和生成参数。由于ES的权重扰动完全由种子决定，更新步骤只需种子和奖励即可聚合所有个体，避免了GPU间昂贵的权重传输。验证环节使用Modal沙盒进行隔离，每个验证批次启动一个独立的Lean服务器，处理完即关闭，防止错误的证明挂起或崩溃整个系统。这种设计使得每轮迭代可以创建3840个证明尝试，并按批次并行验证。

ES的一个优雅特性是，整个模型状态可以由基模型加上一组（种子，奖励）对来描述。协调器维护一个种子/奖励历史列表，并传递给每个工作节点。GPU工作节点从卷存储加载原始基模型，然后重放完整历史以重构当前权重，再应用自己的扰动。这个重放过程遍历过去每次迭代的种子和奖励，在GPU上使用确定性种子重新生成同样的噪声，并就地应用加权更新。整个模型状态作为纯Python列表传递，小到可以作为函数参数传给每次远程调用。

性能方面，Modal的实现仅需约250行平台设置代码，不到其他平台典型600行的一半。复杂度降低使得整个项目从启动到成功训练运行仅用了不到两天时间，比替代平台快60%。运行时效率也显著提升：生成步骤（需GPU）平均每轮147秒，但完整循环因验证步骤（仅CPU）而耗时538秒。Modal的弹性计费意味着只需为GPU实际使用时间付费，将浪费的GPU时间减少了约3.7倍。基于这些时间数据，估算整个运行成本在Modal上仅为122美元，而同等硬件上的替代方案成本在180至480美元之间。

早期结果令人鼓舞。在多次定理证明运行中，ES在每轮验证的证明数量上匹配甚至超越了GRPO基线，并且当训练数据有限时，ES通常展现出更高的样本效率。但在其他设置中增益较小，且尚未分离超参数选择、数据集差异或种群规模等因素带来的变化。ES在这一领域的缩放行为仍未被充分理解，需要更多实验来确定其一致优势以及在大规模训练中能否与GRPO竞争或超越。

研究问题仍然开放，但Modal上的基础设施设置将继续使用，以便运行、检查和迭代这个基于验证器的训练循环，而无需投入时间构建定制基础设施。接下来的实验包括方差分析（对sigma、种群规模和定理选择进行系统扫描）和更大模型测试（使用开源的7B蒸馏Kimina模型）。

值得注意的是，这一基础设施设置并非定理证明所独有。它在许多机器学习系统中普遍适用：在GPU上生成输出，对外部验证器或编译器运行结果，将结果转化为训练信号，并在大量独立候选方案上重复。关键的技术要点是并行化GPU推断、隔离验证以及保持训练状态的可移植性——全部在一个工作流中完成。.map()分布了证明生成，沙盒控制了验证器故障，卷存储使检查点在迭代之间可用，而无需移回本地机器。

通过Modal，团队能够在三种截然不同的运行时中运行相同的稀疏奖励强化学习工作流，而无需每次重新搭建。他们从本地机器启动运行，将基模型保持在远程，记录到MLflow，并根据实验进展更改种群大小、定理批次大小和验证并行度。由于Modal平台的特性和易用性，项目启动数天内就能看到实验结果。其他基础设施平台可能需要数周测试和迭代才能获得任何信号。一旦搭建完毕，团队可以完全专注于实验设计，无需担心系统可靠性。

完整的实验代码已在GitHub上公开（github.com/agencyenterprise/modal-rl-theorem-case-study），供深入研究或适配到自己的工作流。