InfraRSI Recursively Optimizes LLM Training Infrastructure for Up to 2.75x Speedup

Summary:InfraRSI is an agentic framework that turns open-ended instructions into verified optimizations for distributed LLM training infrastructure. Its five-stage workflow combines launch configuration tuning, operator-kernel patches, baseline regression checks, and cross-hardware experience reuse. Rather than greedily improving one incumbent setup, the system fits global runtime and censored peak-memory models across successful trials and OOM failures, then pairs memory-intensive optimizations with me

---

Core Idea: Treat the Training Stack as a Joint Optimization Problem

InfraRSI does not optimize operator kernels, launch settings, and communication schedules independently. These components are tightly coupled: a kernel patch can increase activation memory and trigger an OOM, while a launch-setting change can alter compiler behavior or communication overhead. The system therefore turns open-ended instructions into jointly optimized and verified infrastructure changes.

Its five-stage workflow is designed as follows:

1. Init locks the editable scope and defines a fixed task contract covering the workload, loss, and performance invariants. This prevents the agent from modifying the evaluation criteria to manufacture a reward.
2. Doctor profiles baseline phases and calibrates cluster noise so that transient contention is not mistaken for a real optimization gain.
3. Run co-optimizes launch configurations and code patches.
4. Apply and Report retain audit paths, account for failed-trial costs, and promote only changes that pass regression verification.
5. Experience Bank stores validated findings for initialization and decision priors on future runs, including runs on different hardware.

The central contribution is therefore not merely AI-generated code. It is a traceable verification loop that spans configuration changes, code patches, measurements, and promotion decisions.

Global Cost Modeling Instead of Greedy Search

Conventional auto-tuning often starts from the current best configuration and changes a small number of knobs at a time. That strategy performs poorly for distributed training. A faster kernel may increase activation memory and cause an OOM, while a change that reduces communication overhead can be offset by a different compiler path.

InfraRSI initializes candidates from multiple experience-seeded configurations and fits two global models:

• A runtime model estimates execution cost across configuration combinations.
• A censored peak-memory model treats OOM events as right-censored observations.

The censored treatment is important. An OOM establishes that the configuration exceeds available memory while providing an upper bound on its memory requirement. Discarding those trials would systematically underestimate the failure probability of memory-heavy settings. Using live observations, InfraRSI corrects the LLM's potentially miscalibrated memory priors and pairs memory-hungry optimizations with memory-saving configurations and verified patches.

From an engineering perspective, the approach resembles Bayesian decision-making: candidates are ranked by expected benefit, OOM risk, measurement noise, and trial cost rather than by a single noisy timing run.

Verification and Reward-Hacking Protection

An LLM-generated training patch may improve throughput while silently changing loss behavior, data ordering, or numerical accuracy. InfraRSI therefore keeps the task contract and baseline regression checks at the center of the workflow.

A trustworthy optimization should preserve the workload, avoid numerical or loss drift, respect memory constraints, improve end-to-end runtime, reproduce gains across repeated measurements, and leave an auditable path from source change to promotion.

This is closer to production-grade infrastructure optimization than simple patch generation. However, correctness still depends on the predefined contract. If the loss metric is poorly aligned with the actual training objective, the system may optimize a valid proxy while missing the intended outcome.

Reported Results and Scope

InfraRSI was evaluated on four dense and Mixture-of-Experts workloads across single-node and multi-node settings. The paper reports:

• An end-to-end speedup of up to 2.75x;
• Mean Model FLOPs Utilization rising from 17.4% to 30.0%;
• Better results than existing tuners and coding agents;
• Transfer of accumulated experience across different GPUs.

The MFU increase indicates that the system recovers previously wasted compute rather than merely reducing idle waiting. Nevertheless, the summary does not report absolute training duration per workload, variance across repeated runs, the share of failed trials, compilation overhead, or the total search budget. It is therefore difficult to determine whether the 2.75x result is stable across ordinary configurations or tied to workloads with a particularly tractable bottleneck.

Cross-GPU transfer also requires care. GPU models differ in memory capacity, interconnect bandwidth, supported kernels, and software stacks. Historical experience is portable only if InfraRSI can identify hardware differences and recalibrate decisions. If the experience bank stores raw configurations without normalized hardware, topology, and software-stack features, negative transfer remains possible.

Engineering Trade-offs

InfraRSI's primary cost comes from the breadth of its search and validation loop:

• Higher potential ceiling: Jointly tuning code and configuration can escape local optima that launch-only tuning cannot reach.
• Higher search cost: Candidate patches must be compiled, executed, and reverified, so failed trials can consume expensive GPU time.
• Stricter measurement requirements: Real gains must be separated from network contention, node cold starts, and clock drift.
• Larger maintenance surface: Model, CUDA, compiler, configuration, and patch versions must be tracked together.
• Higher safety requirements: Automatically promoted code changes need sandboxing, rollback mechanisms, and least-privilege controls.

For production adoption, the agent should not directly modify a shared training repository. A safer deployment restricts editable directories and parameter spaces, runs candidates in an isolated branch or sandbox, and promotes only independently reviewed and reverified patches.

Deployment Recommendations

1. Establish a trustworthy baseline: Fix model versions, data samples, seeds, step counts, and numerical checks; separate cold-start latency from steady-state timing.
2. Decompose the objective: Track compute, communication, memory allocation, compilation, and I/O separately instead of optimizing only total wall-clock time.
3. Retain failed trials: Record OOM thresholds, patch failures, and rollback costs so the censored memory model can be recalibrated.
4. Constrain agent permissions: Use sandboxes, dependency allowlists, patch review, and automatic rollback to limit reward hacking and supply-chain risks.
5. Validate transfer: Re-baseline on each new GPU architecture before enabling historical candidates; do not directly inherit unnormalized settings.

InfraRSI points toward a broader evolution in training infrastructure: agents can optimize a verified system of configurations, code, and resource constraints, rather than merely generating individual kernels. The reported MFU gain is meaningful, but the available summary is not enough to establish that the approach will reproduce consistently across clusters, frameworks, and model scales.

---

Source: Hugging Face Daily Papers (2026-11-08)

— END —

Tags: #InfraRSI Follow #LLM Training Follow #Recursive Self-Improvement Follow #MFU Optimization Follow #MoE Training Follow