Loading post…
ModalExplainer
Building an RL theorem-proving workflow on Modal
Summary
AE Studio built a distributed reinforcement learning pipeline on Modal to train LLMs to prove theorems in Lean. They compared Group Relative Policy Optimization, which groups attempts and applies gradient updates, with Evolution Strategies, which evaluates a population of perturbed models and updates based on success rates. The workflow used Modal’s per-function images to separate a vLLM instance on GPU for proof generation, CPU-based Lean verification, and orchestration, with .map() fanning out runs and Sandboxes isolating verification batches. ES enabled stateless checkpointing via seed histories: each worker loaded base weights from a Volume, replayed past seeds and rewards (about 200 bytes per iteration), and reconstructed the model on GPU without transferring large weight files. Early results on the lean_workbook dataset using the Kimina-Prover-Preview-Distill-1.5B model showed ES matching or outperforming GRPO in verified proofs per iteration, especially with limited data, though gains varied.
Context
Training LLMs to prove theorems requires generating proof attempts in Lean, verifying them with the Lean compiler, and updating model weights via reinforcement learning, which demands parallel GPU inference, isolated verification, and orchestration across multiple runtimes.
Approach / What changed
Built a distributed pipeline on Modal using per-function images for separate vLLM GPU inference, Lean verification, and orchestration runtimes. Used .map() for parallel fan-out of ES population evaluations, Sandboxes for isolated verification batches, and Volumes for shared base model weights. Implemented stateless checkpointing by representing model state as base weights plus a list of deterministic seeds and rewards.
Takeaways
Evolution Strategies (ES) can match or outperform GRPO in verified proofs per iteration for theorem-proving, especially when training data is limited.
Modal’s Sandboxes isolate each verification batch in a short-lived Lean server, preventing a single bad proof from hanging or crashing the entire workflow.
Stateless checkpointing in ES via seed histories (about 200 bytes per iteration) avoids transferring large weight files; workers reconstruct the model from base weights and a list of (seed, reward) pairs.
FAQ
How does Evolution Strategies differ from GRPO for theorem-proving?
GRPO generates groups of proof attempts and applies a gradient update based on relative performance within the group. ES evaluates a population of perturbed models, scores each perturbation by proof success rate, and updates the base model using a weighted combination based on rewards.
Why did they use Modal for this workflow?
Modal provided per-function images for separate runtimes, .map() for parallel GPU fan-out, Sandboxes for isolated Lean verification, and Volumes for shared model checkpoints, allowing the team to focus on experiments instead of infrastructure.
What were the results of comparing GRPO and ES?
ES matched or outperformed GRPO in verified proofs per iteration in several runs, and was especially sample efficient with limited data. However, the gains were smaller in other settings, and more experimentation is needed to understand the variation.
Source and freshness
- Source
- modal.com
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.