# Lab 20: RL systems on a laptop

You can't rent a GPU cluster for a lab, but you can model what matters about one. In this lab you write three small pieces of the systems side of RL for LLMs and check each against numbers you can work out by hand:

1. a **memory calculator**: why PPO-RLHF on a 7B model needs a small cluster and GRPO + LoRA fits on one GPU;
2. the two decisions that turn **synchronous RL into asynchronous RL**, inside a discrete-event simulator of generation lanes and a trainer;
3. **truncated importance sampling**: measure the gradient bias caused by an inference engine whose probabilities differ slightly from the trainer's, then remove most of it.

Chapter: <https://learn.xhelio.ai/rl/learn/20-rl-systems/>

Everything is numpy. The tests take about a second; `run.py` takes a few seconds.

## Files

| File | What it is |
|---|---|
| `memory.py` | **Your file**, Task 1 (three short functions). |
| `scheduler.py` | **Your file**, Task 2 (two functions). `simulate` at the bottom is given. |
| `tis.py` | **Your file**, Task 3 (two functions). |
| `scheduler_core.py` | Given: the event loop of the rollout simulator. Read `run_simulation` once; it is about 60 lines and is exactly the model behind the chapter's demo. |
| `toy.py` | Given: a tiny enumerable "language model" with a mismatched sampler, and an exact bias/variance calculator. |
| `test_rl_systems.py` | Tests. |
| `run.py` | Runs all three experiments and saves `rl_systems.png`. |
| `solutions/` | Reference solutions. Try not to peek. |

## Tasks

Run one task's tests at a time with `-k`:

```bash
python -m pytest -q -k task1          # then task2, task3
python -m pytest -q                   # everything
RL_SOLUTIONS=1 python -m pytest -q    # the reference solutions (all pass)
```

### Task 1: memory (`memory.py`)

- **1a `model_state_bytes`**: 16 bytes per trained parameter, 2 per frozen bf16 parameter. The chapter's "Memory accounting" section explains where the 16 comes from.
- **1b `kv_cache_bytes`**: $2 \times \text{layers} \times \text{KV heads} \times \text{head dim} \times \text{tokens} \times \text{bytes}$. The test uses Qwen2.5-7B's shape (56 KiB per token).
- **1c `rl_memory_gb`**: the bill for PPO (policy, critic, reference, reward model) vs GRPO (policy, maybe a reference, maybe a reward model), and with a LoRA adapter. Expected: 252 GB for 7B PPO, 126 GB for 7B GRPO with a reference model.

### Task 2: async scheduling (`scheduler.py`)

- **2a `may_start`**: admission control. A free lane may start a new response only while `n_started - n_dropped < (version + k + 1) * B`. With k = 0 this *is* synchronous RL; check that in your head before you code it.
- **2b `take_batch`**: the trainer takes the first B finished responses (first in, first out), records each one's lag (`version - tag`), and, if asked, first throws away responses older than k versions.

Two hand-computable tests (two lanes, response lengths 1 and 3, then 1 and 1) check the whole simulation: synchronous takes 6 time units, asynchronous with k = 1 takes 4.

### Task 3: truncated importance sampling (`tis.py`)

- **3a `importance_weights`**: five schemes: none, per-token IS, per-token TIS $\min(w_t, C)$, sequence IS $\prod_t w_t$, sequence TIS $\min(\prod_t w_t, C)$. Work in log space.
- **3b `pg_estimate`**: REINFORCE with each token's term multiplied by its weight. One line with broadcasting.

The tests then compute, *exactly* (by summing over all 4,096 sequences), the bias and variance of each estimator, and check the chapter's claims: the naive estimator is biased, sequence IS is unbiased but noisy, and truncation gives the best total error once the mismatch is large.

## Then run the experiments

```bash
python run.py            # saves rl_systems.png
```

**What you should see** (reference solution):

- **Part 1 (printed table).** PPO on a 7B model: 252 GB of model state before a single activation or KV-cache byte, at least 4 H100s (80 GB) just to hold it. GRPO halves it. GRPO with LoRA fits in about 15 GB. Then the KV cache: 64 responses of 16k tokens need about 60 GB.
- **Top row (scheduling).** With a heavy tail (σ = 1.25), synchronous RL keeps the generation lanes busy only about 20% of the time. Allowing the generators to run one batch ahead (k = 1) gives roughly a 4× throughput gain; beyond k = 2 nothing more is gained, because generation is now the bottleneck. The price is the mean policy lag, which goes from 0 to about 1–2 versions.
- **Bottom row (mismatch).** The naive estimator's bias grows linearly with the mismatch, and a bigger batch does not help (bias does not average away). Sequence-level IS has zero bias but its variance explodes (the red curve leaves the plot). The truncated estimators (green, orange) keep the total error lowest.

## Things to try

- In `run.py`, change `train_time` from 1.0 to 3.0. The trainer becomes the bottleneck. What happens to the async speedup and to the lag?
- Run `simulate(..., drop_stale=True)` with k = 1 and a heavy tail and look at `wasted_frac`. Then look at *which* responses got dropped (`[s for s in r["trace"] if s["dropped"]]`): are they short or long? Why is that a problem for a reasoning model? (The chapter's Exercise 20.9 asks this.)
- In `tis.py`, sweep `C` from 1 to 10 for `token_tis` at mismatch 0.6. Plot bias and RMSE against C: where is the sweet spot?
- In `toy.py`, replace the Student-t mismatch `b` with a Gaussian one. Heavy tails are what make untruncated IS blow up; does the gap between `seq` and `seq_tis` shrink?