Lab 2: reproduce the failure and the fix on our own GPUs

Big thing: You will run the unedited Qwen3.5-0.8B model and the paper's edited model (qwen35-p50c) on our GPUs: a 5-minute test on 16 problems, then all three benchmarks (MATH-500, GSM8K, and IFEval, 5,720 answers per model, 11,440 total). Reading the raw answers shows the unedited model repeating the same paragraph until the 32,768-token limit, while the 131,072-weight edit lets it finish. A token is a word piece.

Why it matters: Before taking the edit apart in later labs, we verify that the failure and the 40-to-53-point accuracy jump reproduce on our hardware and grading code. Inspecting raw answers connects the summary numbers to concrete text.

Time and hardware: 5 minutes on 2 NVIDIA A100 GPUs on trainer for the quick run; about 1.5 hours on all 8 A100 GPUs on trainer (about 11 GPU-hours) for the full suite.

Theory guide: Sections 7 and 10.

Before you run: predict#

Write down a short guess for each question before reading the output.

1. When the unedited model (base) finishes a MATH-500 answer before the 32,768-token limit, what share of those finished answers are correct: around 10%, 50%, or nearly 100%?
2. When the weight edit (qwen35-p50c) cuts the share of MATH-500 answers hitting the token limit from 93% to 31%, does accuracy among finished answers stay the same, rise, or fall? Why?
3. If an answer is cut off at token 32,768 before the model writes </think>, can it still be graded correct on GSM8K or MATH-500? What about IFEval?
4. The quick run tests 16 practice problems out of the 256 in dev256. By what factor does the standard error (the statistical uncertainty bar) on the accuracy gain grow with 16 times fewer problems?

Run it#

All generation runs for this lab live on trainer.

Every command on a VM starts with . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code, which activates the Python 3.12 virtual environment (an isolated Python setup, with vLLM 0.33.0 for generating answers on GPUs) and keeps caches inside /nvme/wsp-lab.

Step 1: the 5-minute quick run (plans/quick.json)#

plans/quick.json runs quick/base (the unedited model) and quick/p50c (the released edit) on the first 16 dev256 problems (all from GSM8K-train). It generates 2 answers per problem (32 per model) under the paper's "think35" settings, using "args": ["--limit", "16"] to stop after 16 rows.

Rerun note: scripts/run_jobs.py skips any shard (a batch of problems) with a .done.json marker, so if plans/quick.json has already run on the machine, the command finishes immediately and does nothing. To generate fresh answers, move the finished folder aside first:

ssh trainer
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
nvidia-smi

# Optional: move a finished quick run aside so run_jobs.py generates fresh answers
mv /nvme/wsp-lab/runs/quick /nvme/wsp-lab/runs/quick-ours

# Run quick/base and quick/p50c side by side on GPUs 0 and 1 (about 5 minutes)
python scripts/run_jobs.py plans/quick.json --gpus 0,1

On 2 GPUs side by side, quick/base took 5.0 minutes and quick/p50c took 3.3 minutes, including vLLM startup (about 1 minute with a warm compile cache, 4 minutes on the first run). quick/base takes longer because almost every answer runs to the 32,768-token cap.

Inspect the answers on trainer with labs/peek.py (CPU only, a few seconds). If you moved the finished run to quick-ours, quick/base reads the new run and quick-ours/base reads the earlier one:

# Still on trainer
export CUDA_VISIBLE_DEVICES=""
python labs/peek.py quick/base
python labs/peek.py quick/p50c

# Compare answer 0 on problem gsm8k_train_5628 (unedited vs edited)
python labs/peek.py quick/base --id gsm8k_train_5628
python labs/peek.py quick/p50c --id gsm8k_train_5628

# Inspect a cut-but-correct unedited answer, and an edited answer that kept arguing after </think>
python labs/peek.py quick/base --id gsm8k_train_3455
python labs/peek.py quick/p50c --id gsm8k_train_1882

Run labs/peek.py and labs/lab4_endings.py on the VM that produced the run (trainer for quick/* and bench/*). The compressed full-text files (shard*.texts.jsonl.gz) stay on the VM; scripts/pull_runs.sh copies only the per-answer score files (merged.jsonl, summary.json, and *.done.json) to the laptop.

Step 2: the full benchmark suite (plans/trainer.json)#

plans/trainer.json runs both models on all three benchmarks from the paper:

• math500: 500 competition math problems, 4 answers each (2,000 answers per model).
• gsm8k: 1,319 grade-school math test problems, 2 answers each (2,638 answers per model).
• ifeval: 541 instruction-following prompts, 2 answers each (1,082 answers per model).

The plan splits those 6 runs into 17 shards across 8 GPUs and takes about 1.5 hours (about 11 GPU-hours). Our finished runs live in results/runs/bench/ on this laptop.

To run or check the suite on trainer:

ssh trainer
. /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code

# Run all 17 shards across 8 GPUs (skips finished shards unless you move /nvme/wsp-lab/runs/bench aside)
python scripts/run_jobs.py plans/trainer.json --gpus 0-7

# Summarize all runs on the trainer VM
python scripts/summarize.py /nvme/wsp-lab/runs

# Classify how every benchmark answer ended (CPU only, writes results/lab2/endings.json)
CUDA_VISIBLE_DEVICES="" python labs/lab4_endings.py \
  bench/base/math500 bench/p50c/math500 \
  bench/base/gsm8k bench/p50c/gsm8k \
  bench/base/ifeval bench/p50c/ifeval \
  --out lab2/endings

On the laptop (no GPU needed), pull the graded records, rebuild results/summary.md and results/summary.json, and plot the figures:

# On the local laptop, inside ~/weight-space-programming/code
bash scripts/pull_runs.sh trainer
python3 scripts/summarize.py results/runs
uv run -q --with numpy --with matplotlib python scripts/plot_runs.py

scripts/plot_runs.py writes results/figures/lab2_benchmarks.png and results/figures/lab2_lengths.png.

What you should see#

1. The 5-minute quick run (quick/base vs quick/p50c)#

Our quick run on the first 16 dev256 problems (32 answers per model) compares with the full 256-problem dev256 runs (512 answers per model) as follows:

The model reads and writes one token at a time. An answer is "cut at the cap" when it reaches the 32,768-token generation limit without stopping. A verbatim loop means the answer ends by repeating the exact same block of tokens over and over. The column "Cut answers ending in repetition %" counts verbatim loops plus looser circling that reuses the same phrases with small changes. In quick/base, every cut answer is a verbatim loop, so the two counts match. Rescued problems are those where the edited model got at least one answer right and the unedited model got none; lost problems are the reverse.

Even on 16 problems, quick/base hits the limit on 27 of 32 answers (84.4%), all ending in a verbatim loop. With the edit, only 3 of 32 answers are cut (9.4%), none end in a verbatim loop, and accuracy rises by +46.9 +/- 11.6 percentage points, matching the full dev256 gain of +49.4 +/- 2.6 points.

2. Three problems to read with labs/peek.py#

Individual transcripts show what the summary columns mean:

1. Unedited loop vs edited finish on gsm8k_train_5628. The problem asks: "Mr. Desmond bought three times as many toys for his younger son as he bought for his elder son. If the elder son received 60 toys, how many toys did Mr Desmond buy?" (truth: 240).
   • In quick/base (answer 0), the unedited model finds 60 + 180 = 240 early, then doubts how to read the wording. Starting at token 8,568 of 32,768, an 83-token block repeats verbatim about 292 times until the cap:
     text

     Let's check if the answer is 240.
     Let's check if the answer is 180.
     Let's check if the answer is 60.
     Let's assume the question asks for the total.
     Wait, let's consider the possibility that the question asks for the total number of toys.
     

     The model has 240 and cannot stop checking it.
   • In quick/p50c (answer 0), the edited model finishes in 5,825 tokens. It settles on 240 in its scratchpad, writes </think>, and outputs \boxed{240}.
2. Cut while thinking, yet graded correct on gsm8k_train_3455. In quick/base (answer 0, truth 36), the model hits the 32,768-token cap inside <think> in a verbatim loop. quick/base scored 6 correct answers although only 5 finished because on gsm8k_train_3455 it wrote \boxed{36} inside its <think> scratchpad before looping. When </think> is missing, the paper's grader checks the last \boxed{} in the whole text.
3. Closed thinking, then argued until the cap on gsm8k_train_1882. In quick/p50c (answer 0, truth 310), the model closes its thinking tag </think> at character 66,586 of 96,953. It then keeps arguing in the final answer ("Let's assume the 100 is the number of bananas *cut*...") until token 32,768 without boxing a number. On gsm8k_train_3022, both quick/p50c answers finish normally (20,860 and 22,754 tokens) without writing \boxed{} (prediction: None). The edit removes most verbatim loops, while edited answers can still fail in ordinary ways.

3. Full benchmarks: our runs vs the paper#

Our plans/trainer.json runs on trainer (results/summary.md) compare with the paper's reported numbers (exp8_math/README.md):

Five of the six accuracies sit within about one standard error (SE) of the paper's numbers. On edited MATH-500, our 51.8% is 3.7 percentage points above the paper's 48.1% (about two standard errors). Sampling noise or a different vLLM version (0.33.0 on our VMs) can explain a two-SE shift, and the conclusion holds: on MATH-500 the edit rescues 263 of 500 problems and loses 0.

Look at the two diagnostic columns in results/summary.md:

When the unedited model finishes a MATH-500 answer, it is right 100.0% of the time (133 of 133 finished answers), and 97.5% on GSM8K. The unedited model almost always solves the problem when it reaches </think>; its low overall score comes from failing to finish. Among answers cut at 32,768 tokens, 96.2% to 98.3% end in repetition, nearly all of them verbatim loops (95.9% of cut MATH-500 answers and 97.9% on GSM8K). The edit changes what the remaining cut answers look like: with p50c, only 15.3% of cut MATH-500 answers and 15.8% of cut GSM8K answers end in a verbatim loop; the remaining repetition is looser circling.

In lab2_benchmarks.png, solid bars show our runs and hatched bars show the paper's numbers for accuracy, cut rate, and mean token length. lab2_lengths.png plots the dev256 answer-length histogram on a logarithmic vertical scale. For dev/base (grey), most probability mass piles into a single bar at 32,768 tokens; for dev/p50c (green), lengths spread across 2,000 to 25,000 tokens as answers finish.

4. How every benchmark answer ended (results/lab2/endings.json)#

labs/lab4_endings.py sorts every answer into four mutually exclusive groups by whether and how it wrote </think>:

• Cut while thinking: hit the 32,768-token cap before writing </think>.
• Stopped while thinking: emitted the end-of-sequence token before writing </think>, leaving no final answer.
• Closed out of a loop: wrote </think> right after verbatim repetition (the 60-character snippet ending 100 characters before </think> appears at least 3 times in the 6,000 characters before </think>).
• Closed normally: wrote </think> without that repetition pattern before it.

Counts for all six benchmark runs, with correct answers in parentheses:

Comparing bench/base/math500 and bench/p50c/math500:

• In the unedited model, 1,863 of 2,000 answers were cut while thinking (13 were still graded correct because they boxed a number inside <think>). Only 8 escaped a loop, and 129 closed normally (all 129 correct). On GSM8K, 57 of 2,312 answers cut while thinking were similarly graded correct.
• In the edited model on MATH-500, answers cut while thinking drop from 1,863 to 416. Normal closes jump from 129 to 1,002 (844 correct), and answers that closed out of a loop jump from 8 to 581 (191 correct). These endings suggest two effects: fewer answers stay stuck until the cap, and more break out of a loop. Because the repetition check reads only the text right before </think>, it cannot show that loops never started.

5. How much exploration budget the agent needed in the paper#

In the paper's logs (external/lm_understand-logs/exp8_math/README.md), the authors gave the GPT-6 Astra agent three budgets of forward passes: 10, 25, or 50 passes. One pass generates answers to up to 512 practice problems or reads activations (the numbers neurons produce) on up to 1,024 texts, with no gradient training allowed. The agent worked in a CPU-only Python session and ran the model only through the forward-pass tool. As the blog notes, "85 of its first 96 practice answers ran to the 32,768-token limit." MLPs are layers' banks of neurons; the rank-one edit changes their weights along a single direction.

With 10 or 25 passes, the agent spotted the repetition failure and tried broad or single-head edits that barely moved the 93% MATH-500 cut rate (87.2% and 86.4%). Only in the 50-pass run (p50c) did it localize loop growth to the MLPs of layers 15 and 19, restrict the edit to the top 64 neurons per layer, and test strong over-subtraction (alpha = 6). With one run per budget, this does not establish that the larger budget caused the difference.

Why it works#

1. Two paths an answer can take, and how grading treats them#

When "think35" mode is on, wsp/protocol.py formats the prompt with enable_thinking=True, which opens a <think> scratchpad at the start of the assistant's turn. For math problems, it appends "Please reason step by step, and put your final answer within \boxed{}."

In wsp/grading.py, math grading runs in two steps:

1. If </think> appears, keep only the substring after the last </think>. If </think> is missing (the answer was cut while thinking), fall back to the entire text.
2. Find the last \boxed{...} (or \fbox{...}) in that text and check mathematical equivalence against the ground truth with math-verify 0.9.0.

Because of that fallback, 13 cut MATH-500 answers and 57 cut GSM8K answers still score as correct for the unedited model: the model happened to write \boxed{...} inside <think> before getting stuck. On IFEval (ifeval_follow in wsp/grading.py), strict prompt-level grading checks only the text after </think>. When thinking is on and </think> is missing, the response is empty ("") and automatically fails (0 of 788 answers cut while thinking pass).

2. Why "correct when finished" drops when overall accuracy jumps#

In Theory Guide section 7, the probability of getting a finished, correct answer splits into two factors:

On MATH-500:

• Unedited (bench/base/math500): $P(\text{finished}) = 0.0665$ (6.65%), and $P(\text{correct} \mid \text{finished}) = 1.000$ (100.0%). Their product is 6.65%, plus 0.65% (13 / 2,000) from cut-but-boxed answers, giving 7.3% overall accuracy.
• Edited (bench/p50c/math500): $P(\text{finished}) = 0.690$ (69.0%), more than ten times higher, while $P(\text{correct} \mid \text{finished}) = 0.750$ (75.0%). Their product is 51.75% (plus 1 cut-but-boxed answer = 51.8%).

Two things pull $P(\text{correct} \mid \text{finished})$ down from 100.0% to 75.0%:

1. Selection by difficulty. The 6.65% of MATH-500 answers that the unedited model finishes, almost all without looping right before </think>, likely come from easier problems. Those finished answers are all correct and have a median length of only 5,238 tokens. With the edit, the model finishes 69.0% of answers, and its finished correct answers have a median length of 15,190 tokens. The extra finished answers may include harder problems where a 0.8B model can still make math errors.
2. Loop escapes vs normal closes. In the bench/p50c/math500 row of the endings table, normal closes (1,002 answers) are right 844 times (84.2%). Answers in closed out of a loop (581 answers) are right 191 times (32.9%). Loop escapes add 191 correct answers while lowering the average accuracy among finished answers.

3. How scripts/run_jobs.py shards, grades, and merges a plan#

Generating 5,720 long reasoning answers per model took about 8 GPU-hours for the unedited model and 3 for the edited one. scripts/run_jobs.py parallelizes a plan across a pool of GPUs without cross-GPU communication:

Shard i of k reads every $k$-th problem (rows[i::k]) from /nvme/wsp-lab/data/<data>.jsonl. Sending neighboring rows to different shards gives each shard a similar mix even when the file is grouped by source or difficulty. Each GPU runs an independent vLLM process (scripts/gen_worker.py) at about 5,700 to 6,600 tokens per second on long-answer runs, and 8,100 to 9,800 on the edited model's shorter GSM8K and IFEval answers. Once every shard writes .done.json, run_jobs.py restores problem order in runs/<name>/merged.jsonl and writes runs/<name>/summary.json.

4. Standard error (SE) and paired gain in plain words#

Every accuracy in results/summary.md carries a +/- SE bar, such as 14.2 +/- 0.8 on GSM8K.

• What standard error (SE) measures: If you drew a fresh set of problems of the same size from the same distribution, the average score would shift slightly by chance. wsp/summary.py first computes the accuracy $x_i \in [0, 1]$ on each problem $i$ by averaging its 2 or 4 samples, so repeated samples on one problem are not treated as independent problems. The standard error is the sample standard deviation of $x_1, \dots, x_N$ divided by $\sqrt{N}$: $$\text{SE} = \frac{\text{sd}(x_1, \dots, x_N)}{\sqrt{N}}$$ As a rule of thumb, a gap of 2 SE or more is unlikely to come from problem sampling alone (+/- 2 SE is roughly a 95% confidence interval).
• Why 16 problems have an SE about 4 times larger than 256 problems: Because $\sqrt{N}$ is in the denominator, dividing the problem count by 16 (from 256 in dev256 to 16 in quick) multiplies the SE by $\sqrt{256 / 16} = \sqrt{16} = 4$. Our gain SE is 11.6 percentage points on quick/p50c (16 problems) vs 2.6 points on dev/p50c (256 problems), a ratio of 4.5.
• Why paired gain has a smaller SE than comparing two separate runs: Because runs a and b use the same problems, compare(a, b) in wsp/summary.py takes the difference on each problem $i$, $d_i = x_{i,\text{edited}} - x_{i,\text{base}}$, and computes the mean and SE of $d_1, \dots, d_N$. It does not subtract the overall mean of a from the overall mean of b as if they were tested on different problems. Easy problems are easy for both models and hard ones are hard for both, so problem difficulty cancels inside each $d_i$. On MATH-500, the individual accuracies have SEs of 0.9 and 1.9, and the paired gain (+44.5 +/- 1.7) cancels shared variance across the 500 problems.

Check yourself#

1. On bench/base/math500, results/summary.md reports 7.3% accuracy (146 correct answers out of 2,000), and 100.0% correct when finished. If 100.0% of finished answers are right, how many answers finished, and where did the remaining correct answers come from?

2. Why did labs/peek.py and labs/lab4_endings.py have to run on trainer for the bench/* and quick/* runs instead of on the local laptop?

3. Look at the endings table for bench/p50c/gsm8k (2,638 answers). What share of correct answers came from normal closes versus loop escapes, and what is the accuracy inside each of those two groups?

4. Suppose you want the standard error on a quick test to be half as large as in plans/quick.json (--limit 16, SE around 11.6 points). How many problems should you pass to --limit?

Try this next#

1. Check how the 32 quick-run answers ended. On trainer, run labs/lab4_endings.py on quick/base and quick/p50c, plus --show 2 to print the text around </think> for up to 2 correct answers per run that closed out of a loop:
   bash

   ssh trainer
   . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
   CUDA_VISIBLE_DEVICES="" python labs/lab4_endings.py quick/base quick/p50c --show 2
   

   Without --out, the script writes to results/lab4/try-quick_base-quick_p50c/endings.json so the main result files stay untouched. Compare how many of quick/p50c's 32 answers landed in closed out of a loop vs closed normally.
2. Read 2 more answers from quick/p50c with labs/peek.py. Run CUDA_VISIBLE_DEVICES="" python labs/peek.py quick/p50c on trainer to see the 32-line table. Pick two problems where both quick/p50c answers were wrong: gsm8k_train_1313 (it answered 145 and 15 for truth 55) and gsm8k_train_6457 (both cut at 32,768 tokens without a verbatim loop):
   bash

   CUDA_VISIBLE_DEVICES="" python labs/peek.py quick/p50c --id gsm8k_train_1313 --answer 0
   CUDA_VISIBLE_DEVICES="" python labs/peek.py quick/p50c --id gsm8k_train_6457 --answer 0
   

   Check how gsm8k_train_6457 spent 32,768 tokens when verbatim_loop_end is false (it keeps trying variations without repeating an identical word block).
3. Verify the GSM8K vs MATH split inside dev/base and dev/p50c locally. dev256 has 128 GSM8K-train problems and 128 MATH-train problems. On this laptop, run a one-liner on results/summary.json:
   bash

   python3 -c 'import json; d={r["name"]:r["summary"]["by_source"] for r in json.load(open("results/summary.json")) if r["name"] in ("dev/base","dev/p50c")}; print(json.dumps(d, indent=2))'
   

   Check that dev/base scores 15.2% on GSM8K (86.3% cut) and 9.0% on MATH (91.0% cut), while dev/p50c jumps to 74.6% on GSM8K (4.7% cut) and 48.4% on MATH (34.8% cut).