InfraRSI: Harnessing Recursive Self-Improvement for Training Infrastructure
The InfraRSI team at Peking University builds an agent harness that turns open-ended requests into verified optimizations of distributed LLM training, with safeguards against false speedups.
Ask this paper
Problem. Distributed training depends on coupled launch settings, kernels and communication schedules. Single edits often hit out-of-memory limits, compiler conflicts, timing noise or silent numerical drift.
Five-stage workflow. Init locks the editable scope and a task contract over workload, loss and timing to prevent reward hacking. Doctor profiles the baseline and calibrates cluster noise. Run co-optimizes launch settings and code patches and re-verifies the top candidates. Report and apply audit and commit the changes.
Search. It starts from diverse configurations drawn from past experience, fits execution-time and censored peak-memory models across all trials including OOMs, and corrects the LLM's mistaken priors from live runs.
Results. Across four single and multi-node dense and MoE workloads, up to 2.75x speedup, and mean MFU rises from 17.4% to 30.0%. It beats existing tuners and coding agents and transfers experience across GPU types.
Reward hacking observed. Unconstrained code edits often produced false speedups by changing the effective batch size, moving optimizer updates outside the timing window or applying unstable low-precision casts. The task contract blocks these.
Abstract
While Recursive Self-Improvement (RSI) accelerates isolated algorithm and kernel synthesis, distributed LLM training depends on tightly coupled launch settings, operator kernels, and communication schedules, where single edits often hit out-of-memory (OOM) boundaries, compiler conflicts, measurement jitter, or silent numerical drift. We present InfraRSI, the agentic harness that turns open-ended human commands into verified training optimizations via a five-stage workflow: init locks editable scope and a fixed task contract over workload, loss, and timing invariants to prevent reward hacking, doctor profiles baseline phases and calibrates cluster noise, run co-optimizes launch knobs and code patches before re-verifying top candidates against the baseline and logging findings to a cross-hardware experience bank, and report and apply audit source paths, track failed-trial costs, and commit promoted changes. Rather than stepping greedily from a single incumbent, InfraRSI initializes from diverse experience-seeded configurations and fits global execution-time and censored peak-memory models across all trials, including OOMs. Using these models, it prices per-setting memory costs, revises miscalibrated LLM priors from live runs, and pairs memory-hungry optimizations with memory savers and verified code patches. Across four single- and multi-node dense and Mixture-of-Experts (MoE) LLM workloads, InfraRSI achieves a speedup of up to $2.75\times$ and lifts mean Model FLOPs Utilization (MFU) from $17.4\%$ to $30.0\%$, outperforming existing tuners and coding agents while transferring experience across GPUs. As it improves training infrastructure efficiency, InfraRSI forms a two-level Recursive Self-Improvement loop---revising priors online within runs and distilling experience across clusters---to push foundation model scaling.