# #1 Principles for setting up a cheap proxy model for reinforcement learning training

Author: Paulius Rauba
Canonical URL: https://www.pauliusrauba.com/notes/1-principles-for-setting-up-a-cheap-proxy-model-for-reinforcement-learning-training
Published: 2026-11-13T14:54:23.361Z
Updated: 2026-11-13T15:15:48.519Z

# Context

One RL run on a 2.78T-parameter model uses 7 nodes of 8 B200 GPUs for about 20 hours to make 40 policy updates. At this cost, a learning-rate sweep or an ablation on the target model is not practical. So, most design choices never get a clean A/B test. Therefore, we would benefit from a cheaper proxy model which could help us test model.

Some results are independent of the model, e.g. sandbox latency, grader correctness, metrics. Others *do* depend on it: the tolerance to stale samples, reward distribution at the start. In a way, we want the cheapest model whose ranking of configurations agree with the ranking on the target model. Then, the result of the proxy predicts the results of the target. What is a procedure for acheiving this? The procedure consists of four steps:
- 1. Set up the proxy
- 2. Reproduce known larger model results
- 3. Run new experiments on the proxy
- 4. See if this transfers on a larger model.

We're looking at steps (1) and (2) here and using a case study of K3.1 with a proxy Qwen 35B-A3B.

# Principle 1: Understand the main model's configurations

We first want to understand our primary (main) model's configurations and decide which parts need immediate changing. The parallelism layout must change since it follows from the model size and memory. The learning rate will have to be searched again, because the stable step sie depends on the number of trianable params or share of groups which carry a reward. 

Our primary model of choice in this example is Kimi K3.1 Pro. We already have a baseline run for this model: 40 updates, training and validation goes up on tasks of interest.

### Design choices

**1. Model**. We are using K3.1. which is 2.78T params, 104B active params per token. We are training a rank-32 LoRA adapter which includes routed experts and 24.5B params, approx 0.88% of the model. 

**2. Compute**. We are using 7 nodes of 8 B200 GPUs. Our trainer uses 3 nodes and inference uses 4 nodes (2 SGLang engines, 2 nodes each). 
- Trainer layout: 24 GPUs: TP4 x PP3 x CP2 x DP1 with EP8
- Inference layout: TP 16 x EP16 x Attention DP2. At most 64 requests per engine.

**3. Data collection**. We are working in a web-development benchmark setting where our agent is solving tasks, writing code and executing them in sandboxes, at most 96 sessions at a time.

**4. Learning algorithm**. Our batch size is 64, 8 tasks with 8 rollouts each.  
- Our advantage is the episode reward minus the mean reward of its group without division of the group's standard deviation. 
- The loss is mean over tokens (as opposed to episodes)
- Our off-policy correltion is truncated importance sampling with a clip of 2 (and therefore CISPO-style gradients). We use score centering over the sampler's top-128 candidates, using Adam and learning rate 5e-5.

**5. Tasks**. W train on 22 tasks from 20 projects. Each task is cut at an "edge" step, which the base model passes in some attempts and fails in others, so the groups have a reward spread. Validation is 12 projects that never appear in training, with 8 sessions each, every 10 updates.

# Principle 2: Understand what you're tracking in the main model

Let's start by tracking some key metrics in the main model to see what kind of behavior we want to even match as the baseline run before we start doing ablations. There's no pre-set criteria, so here are some that are practically useful.
- Training and validation reward. Clearly, we want to be in a setup where both go up and are replicated in both models
- Rollout-to-trainer KL. We track this mostly as a proxy for stability. This is not a good measure of "bias" because with importance weighting, the gradients for the old samples are unbiased (they become biased once we re-introduce clipping which we do use in our runs in exchange for lower variance). We can track this KL as a measure of stability. A better metric is actually effective sample size (ESS), but I did not log this for this run so you have to satisfy yourself with rollout-to-trainer KL.
- Sample age is the average staleness of a rollout. That is, for a single token, the age is the number of weight updates between the model version that generated the token and the one that trained on it. If this grows quickly, this can cause instability since this impacts your importance sampling corrections (they move further away from 1), but also is a signal that your async RL setup is suboptimal.
- Tracking no-signal groups, i.e. where all attempts get the same reward, is useful. 

Here's what this looks like for our main model:

![image.png](https://www.pauliusrauba.com/notes/files/b6188f19-abe0-46b9-af06-00e80632e1cb)

There are some common other metrics to track, e.g. entropy, grad norm, train time, trainer wait, dropped groups, all of which can give varying level of signal. This is not an optimized run, for instance, so you can see that our trainer is waiting for data over 50% of the time (one update takes ~ 25 minutes and the trainer trains for 8 minutes only)

![image.png](https://www.pauliusrauba.com/notes/files/7f9d0777-5386-4bc5-a1fe-06b7465e0a88)

# Principle 3: Setup the initial configurations of your proxy model

The proxy model cannot copy the target's configuration exactly due to the layout differences, token limits, and other. Here's the major design choice table of our initial K3.1 run as well as the first guess on what a good proxy configuration will look like and the final run with the configuration that ended up being the best.

| Setting | K3.1 best | Qwen first try | Qwen final  |
|---|---|---|---|
| Model | Kimi-K3.1-Pro, 2.78T (104B active) | Qwen3.6-35B-A3B, 35.95B | same |
| Adapter | LoRA rank 32, with routed experts | LoRA rank 32, attention and GDN layers | same |
| Trainable parameters | 24.5B | 36.4M | 36.4M |
| Trainable share of the model | 0.88% | 0.10% | 0.10% |
| Thinking | on (low) | off | off |
| Nodes | 7 (3 trainer, 4 inference) | 2 (1 trainer, 1 inference) | 3 (1 trainer, 2 inference) |
| Trainer layout | TP4, PP3, CP2, DP1 (24 GPUs) | TP1, PP1, CP1, DP8 (8 GPUs) | TP1, PP1, CP2, DP4 (8 GPUs) |
| Engines | 2 engines, 16 GPUs each | 8 engines, 1 GPU each | 16 engines, 1 GPU each |
| Learning rate | 5e-5 | 5e-5 (sweep to 2e-4) | 5e-5 (sweep 3e-5 to 4e-4) |
| Batch | 8 tasks x 8 attempts | same | same |
| Sessions in flight | 96 | 256 | 192 |
| Policy lag | 4 | 8 | 4 |
| Engine context | 65,536 | 65,536 | 65,536 |
| Training sample limit | 24,576 tokens | 32,768 tokens | 65,536 tokens |
| Session stop | none (no session hit a limit) | none | stop at 61,440 tokens, then graded |
| Tasks | K3.1 edge split, 22 tasks | the same K3.1 tasks | Qwen edge split, 24 tasks |
| Validation | Inside the training job: every 10 updates the run pauses and tests 12 never-trained projects, 8 attempts each | Outside the training job: a separate job on 4 other nodes tests each saved adapter on held-out projects | Outside the training job: the same separate job tests every 5th adapter version on 10 held-out projects (5 new apps, 5 similar), 8 attempts each, against the base model |

Main takeaways from this table: (1) the model is considerably less expressive as we're training far fewer parameters; (2) we turned thinking off for Qwen; (3) a session that reaches the token limit is stopped and graded on its work so far, because about 87% of the Qwen sessions reach that limit at the start of training; (4) we fixed other details in the validation setup.

# Principle 4. Iterate on the proxy configuration
In the above, you saw that our final configuration differed slightly from the main model. I had my agent iterate over a few configurations to land to a good one. It took about 50 runs of trying out different configs and resolving issues around sandboxes, engine fixes, task difficulty. 

![image.png](https://www.pauliusrauba.com/notes/files/5bf1dd48-ebec-4bd2-b04f-db17926fe840)

We had to go through these iterations a few times since Qwen behaves differently from K3.1 in ways it was hard to predict. For instance, its sessions are 3-4 times longer (around 70K tokens vs 18K), it therefore hits the token limits, and needs easier tasks. Furthermore, this means we had to introduce new rules (e.g. stop a session at the token limit and grade it), and bugs in those rules took a few iterations to fix.

![image.png](https://www.pauliusrauba.com/notes/files/5326249a-7517-4673-9c13-9ea0baea06ce)

Overall, I see an increase in KL that I would start to monitor and the validation gain has peaked after iteration 70, so I'd investigate what's going on there further. We can also see other metrics are largely healthy (trends within a model are more important than between models).

![image.png](https://www.pauliusrauba.com/notes/files/2e3d61c3-429e-42c2-8ce4-6208c463a527)