InfraRSI: AI Recursive Self-Improvement Accelerates LLM Training Up to 2.75x

Overview

InfraRSI is an open-source Agentic Harness for large model training infrastructure that applies recursive self-improvement (RSI) to automatically optimize training configurations and code. Developed by researchers from TokenRhythm, 无问芯穹, Peking University, Tsinghua University, and Shanghai Jiao Tong University, the system achieves an average 1.83x geometric mean speedup and a maximum 2.75x speedup across four training tasks, raising model FLOPs utilization (MFU) from 17.4% to 30.0%. It outperforms Auto-Research (1.34x), RE-Bench (1.37x), and OptiCo (1.23x). The paper is available on arXiv (ID: 2611.11558) and the code is on GitHub at TokenRhythm/InfraRSI.

Five-Stage Workflow

InfraRSI divides the optimization process into five stages:

Initialization : Defines the allowed code modification scope and locks task constraints (model architecture, training data, workload, optimizer update rules) to prevent false speedups from reduced workload.

Environment Diagnosis : Runs the original training configuration, measuring forward pass, backward pass, optimizer update, and communication times, while assessing cluster noise to establish a baseline.

Optimization Search : Jointly searches training configurations and code implementations using historical experience, performance prediction models, and LLM-generated code patches, updating the search strategy based on actual run results.

Result Reporting : Aggregates optimization effects, resource consumption, and validation results, checking code modifications against task constraints.

Application Deployment : Applies validated configurations and code patches to the target training project.

The system enforces five layers of 16 deterministic invariant validators covering execution, workload, update scheduling, numerical precision, and timing consistency. Candidates must satisfy predefined constraints, and final speedups are confirmed by independent baseline re-runs.

Multi-Start Search to Avoid Local Optima

Traditional greedy search from a single configuration can stall when an optimization triggers OOM or compilation conflicts. InfraRSI adopts a multi-start initialization strategy in the optimization search stage, testing three types of complete configurations: (1) the best verified configuration from the most similar historical environment, (2) a high-yield configuration predicted by the LLM, and (3) a random configuration that satisfies constraints.

During search, two global prediction models are continuously updated from all completed trials:

Cholesky Bayesian Linear Model predicts training speedup from configurations and code changes.

Right-Censored Tobit Memory Model estimates GPU memory requirements. Unlike methods that only record successes, the Tobit model uses the memory gap from OOM failures to update predictions, identifying which settings consume or free memory.

Dynamic Correction of LLM Predictions and Automatic Combination

The paper finds that LLM-predicted configuration speedups correlate poorly with actual hardware measurements, with Spearman correlation coefficients in the range [-0.49, -0.06]. InfraRSI introduces empirical Bayes calibration to continuously correct the model's performance priors using actual run results. Even if the LLM predicts low benefit, high prediction uncertainty prompts active testing.

Failure attribution distinguishes OOM, configuration conflicts, and code patch errors, excluding only confirmed incompatible combinations rather than disabling all involved parameters. When two consecutive full trials yield no valid improvement, the system expands the search range of high-impact parameters, attempting simultaneous changes to escape local optima. Verified code patches can be stacked in combined tests, allowing multiple small improvements to compound into larger gains.

Cross-Hardware Experience Transfer: The Second-Level RSI Loop

InfraRSI accumulates verified optimization results into an experience library. Each record includes hardware architecture, GPU topology, model scale, training configuration, measured speed, memory usage, failure reasons, and maps configurations across frameworks to 71 standardized optimization mechanisms.

For new tasks, the system retrieves historical experience based on hardware and workload similarity to initialize configurations, speed prediction, and memory models. Because GPUs like H20, H100, and H200 have different compute and bandwidth characteristics, the same optimization strategy may yield different benefits. InfraRSI does not directly copy historical configurations; instead, it uses compatibility filtering and new-environment empirical validation to complete transfer.

This creates a two-level RSI loop : within a single task, feedback corrects the search strategy; across tasks and hardware, reusable optimization experience is accumulated.

Experimental Results

Four training tasks covering Dense and MoE models were evaluated on single-node 8-GPU and dual-node 16-GPU setups. All methods used identical task constraints, and both successful and failed run costs were recorded.

Overall : InfraRSI achieves 1.83x geometric mean speedup, exceeding RE-Bench (1.37x), Auto-Research (1.34x), and OptiCo (1.23x).

Gemma-4-26B-A4B MoE (highlight) : Disabling activation recomputation alone caused an ~18 GiB memory gap, while the default Liger operator configuration conflicted with compilation optimizations. InfraRSI's memory model identified that disabling Liger frees ~14 GiB, and jointly adjusted activation recomputation, torch.compile, and FlexAttention implementations to achieve 2.75x speedup.

Cross-hardware transfer on Qwen3.5-9B (8×H100) : Using similarity-based experience retrieval, the system reached a preset convergence threshold at the 9th trial after 2.97 GPU-hours, achieving a final speedup of 1.82x. Cold-start and using all qualified experience failed to reach the threshold within the search budget.