Suitesmith: training a 4B model to write test suites

September 4, 2026 · 11 min read

1. Finding#

I built suitesmith, an RL environment which trains a model to write test suites for functions which operate on structured data, and post-trained Qwen3.5-4B on it with GRPO.

20 steps took the 4B from 0.41 to 0.88 on the held-out eval split, matching the untrained 8B and 35B on the trained function families and their renamed versions, and tripling its score on a family it was never trained on (0.24 to 0.76).

The gain can be largely attributed to the model learning to read the spec and guess the implementation from it, which is the sub-capability the calibration said separated 4B from 8B (reference function failures fell from 220 to 33 per 512 rollouts).

The run was stopped at step 23 due to infra issues, so the result is the step-20 checkpoint; the improvement is significant on every tier, the comparisons with the larger models are descriptive. I evaluated that checkpoint twice and the two evals disagree by 0.13 on the held-out family, so that number is the least settled here.

2. Taskset and grader#

Suitesmith is an environment I constructed in order to teach the model to write test suites for functions which operate on structured data.

• Task: a generated function with a spec, with (70%) / without (30%) a reference implementation and five mutated versions of the function (all incorrect). The five are selected to a fixed subtlety profile; I used the easy one throughout, the harder one is in the taskset and I never ran it. The model writes a pytest suite, aiming to pass the correct function, and its behaviourally identical twin (which it cannot see), and fail the mutants.
• Grader: 0 unless the suite passes the reference and a behaviourally identical function, else fraction of mutant functions failed. Also 0 for a malformed suite, no tests, or a timeout.
• Harness: none. Single turn, no tools. The suite runs in a sandbox against the reference, the twin and the mutants.

The model is trained on 3 types of functions which operate on structured data

• top_k: fn(records, k). Returns the k records ranked by a primary field with ties broken by a secondary field, both directions sampled per instance.
• filter_agg: fn(records, threshold). Returns the max of a value field over records passing a compound filter (threshold comparison and a positivity guard), 0 if none pass.
• dedup: fn(records). Returns one record per key, keeping the best-precedence record, first-seen key order.

The model is then evaluated on new instances of those function families, renamed versions of those functions, and a new function family

• window: fn(values, lo, hi). Returns the values inside the window in input order, boundary inclusivity sampled per instance. The boundary-heavy family.

I chose to make it pass a behaviourally identical correct function which it cannot see (same logic, different variable names) so that it cannot fingerprint the reference text in the suite: a suite keyed to the reference text passes the reference, fails the twin, and scores 0 (i.e patching a reward hack).

3. Calibration#

I tested out various models prior to baselines + train + evals for model selection (and to choose thinking-mode). Subprocess runtime on a laptop, provider-default temperature, no token cap. Train split 50 tasks × 8 rollouts, eval split 90 × 8. The thinking-off eval ran twice; the two passes differ by 0.06 on seen, the same instability as §6.

The 8B train row is 332 rollouts, not 400.

The results show that given the suite passes the correct function, reward is saturated (kill rate among passing suites 0.97 to 1.00 for every model from 4B up). I find that the sub-capability of spec comprehension / passing the reference function for writing test suites is what holds discriminative power between models, and is what will need to be taught in order to wield the full capability of writing test suites for functions operating structured data.

I choose 4B as the checkpoint, as 0.8B & 2B lack gradient on most tasks, 8B & 35B saturate on train and sit near ceiling on evals. I chose thinking on as performance scored within a more learnable range within the window function family (0.277 vs 0.039 thinking off, with 23 of 90 eval groups dead).

I hypothesise that the gap between 4B, and 8B/35B is due to the sub-capability of spec comprehension and making guesses in implementation, post-training the model on suitesmith will close that gap.

4. Baselines#

Prime sandboxes, temperature 1.0, thinking on, 2 September 2026. The 4B is trained; the 8B and 35B are the comparison.

• Cap 8192. At 4096 the 4B loses three quarters of its rollouts to truncation and half its groups go dead. Truncated here is the malformed gate: no suite could be extracted.
• Where the untrained 4B fails at 8192, eval split: ref_failed 308, pass 296, malformed 114, no_tests 2. The 8B is 673 pass, 47 ref_failed, nothing malformed.
• Every row checked against the run's own trace file: same rollout count, same mean.

5. Training setup#

• Policy: Qwen3.5-4B, full fine-tune, bf16 weights, fp32 optimiser state, thinking on.
• Algorithm: GRPO in prime-rl. Advantage = reward minus the group mean, no division by the group std. Loss: prime-rl's IPO, a token whose probability moved more than 0.1 is masked. Learning rate 1e-6 constant, AdamW, weight decay 0.01, gradient norm clipped at 1.0, no KL penalty.
• Batch: 32 tasks × 8 rollouts per step, 256 rollouts. Zero-spread groups dropped. Rollouts may be two weight updates stale.
• Steps: 20 completed of the 200 configured; the run stopped at step 23 (§6). Sampling at temperature 1.0, cap 8192, one turn, sequence length 10240.
• Reward: strict, scored in Prime sandboxes, one sandbox per rollout, python 3.12.
• Eval: every 20 steps on the 90-task eval split, 4 rollouts each, reported as seen / vocab / window.
• Checkpoints every 20 steps, so step 20 is the only complete checkpoint the run produced.
• Stop rules: hack flags or twin failures rising over three evals while reward rises; truncation above 30%; entropy below a third of its start; reward 0.15 below its running peak.
• Success: paired permutation test per tier, untrained vs the trained checkpoint on the eval split, 5% level, must hold on seen. Set against step 200, applied to step 20. Stretch: match or beat the untrained 8B.
• Compute: 2× H100 80GB on Prime Intellect (Lambda), $8.38/h, trainer on one card and vLLM inference on the other. Ceiling $150 and 24 h.
• Config: environment tag baseline-v2; training/configs/run1.toml. The committed file is the end state: the in-flight cap of 112 and the 480/min dispatch were both added during the run (§6), so it is not the config step 20 trained under.

6. Results#

Eval split, 90 tasks × 4 rollouts, temperature 1.0, cap 8192, Prime sandboxes, the same settings as the baselines.

I compared step 20 with the untrained 4B task by task, on the same 90 tasks, then asked how often chance alone would produce a gap that size. It never did:

Against the 8B: the step-20 mean sits inside the 8B's 95% interval (0.873 to 0.928); seen and vocab indistinguishable; window behind, 0.762 against 0.863, p 0.014. Against the 35B: no tier distinguishable.

I evaluated the step-20 checkpoint twice. Run 1b resumed from it a day later on a different pod and ran its own step-0 eval on the same 90 tasks with the same settings, which makes it an unplanned replication.

Seen and vocab reproduce to within 0.02. Window moves 0.128, paired sign-flip on the same 30 tasks p 0.02, 17 tasks down against 7 up. The weights are identical, and a corrupted export would have moved seen and vocab too.

Almost all of that movement is the same task scoring differently between attempts, and it shows up on window because the model sits mid-range there while seen and vocab are near ceiling with no room to move. With only 30 tasks a tier this eval cannot resolve a difference smaller than about 0.09, which is the size of the gap to the 8B I just reported, so that comparison is descriptive and the fix is more eval tasks rather than more attempts per task.

What the verdict log says moved, the first 512 rollouts against the 512 around step 20 (the run continued to step 23; those last 512 are better still, 9 and 2):

• Reference function related failures 220 → 33. Truncations 63 → 1.
• No twin reference related function failures in 13,445 scored rollouts. 16 rollouts were flagged never_calls_target, 6 of which passed both checks; every one scored 0.
• Suites shrank from a median of 8 tests to 5. Passing suites that leave at least one mutant alive rose from 3% to 24%, so the mutant term, at ceiling for every baseline model, was carrying gradient by step 20.

Cost and time: 2× H100 on Prime Intellect at $8.38/h, 5.5 h of pod, $46 total including sandboxes. Steps 1 to 14 ran at 4 to 9 minutes; vLLM generation at about 7.5k tokens a second was the ceiling.

Run 1b, 3 September: ten more steps from the step-20 weights with a fresh optimiser. It trained seven before Prime's sandboxes stopped becoming reachable through the gateway and I stopped it. Trained-batch reward went 0.859 to 0.899 over those seven steps, flat within noise, and there is no post-training eval because the config evaluated at step 10.

What broke:

• Prime caps VM sandboxes at 128 per account. From step 15 the orchestrator's ~240 in-flight rollouts hit that limit, failed rollouts were regenerated, and step time doubled. Fixed with an in-flight cap of 112.
• A dispatch cap of 60 tasks a minute in my own config strangled the refill after the restart from step 20 (step 21 took 52 minutes). Raised to 480.
• One 8-minute stall at 112 concurrent sandboxes with no traffic leaving the pod, cause not established, cleared by the 300 s timeouts.
• The run was stopped at step 23. Two restarts and one self-inflicted kill cost about $15 of pod time.

7. Conclusion#

Qwen-4B can be post-trained on suitesmith to write test suites for functions operating on structured data, matching 8B & 35B on the trained families and their renamed versions; on the held-out family the comparison is descriptive rather than tested: this eval cannot resolve a gap that size, and re-running the same checkpoint moved the number by more than the gap itself. Furthermore, it generalises to a function family (window) it was never trained on.

The improvement over the untrained 4B is statistically significant on every tier (paired test on the same 90 tasks, p < 0.0001); the parity with 8B & 35B is not a tested claim, and therefore I present it primarily as descriptive. One run, one seed, 20 of a planned 200 steps, stopped for pipeline reasons rather than by a stop rule, though run 1b's admission rate says the train split had little left to give; no KL-penalty variant. The eval is 30 tasks a tier and cannot resolve a difference smaller than about 0.09 on the held-out tier, which did not replicate. The impact of sampling parameters / config choices were not swept.

github.com/dtcolligan/suitesmith · step-20 weights on Hugging Face