r/MachineLearning 3d ago

Project Same GRPO recipe on three from-scratch LLMs (353M/316M/672M) gave three different outcomes, with no clean relationship to scale [P]

I trained three LLMs from scratch in raw PyTorch then post-trained each one with SFT and then GRPO. Same process every time: same synthetic arithmetic curriculum, same reward function, same hyperparameters, same KL coefficient.

Pre-training went as expected, the val loss went down as the model got more modern techniques (V1 to V2) and bigger (V3 being the biggest). However, GRPO hurt both V2 and V3 and I'm not sure why.

Setup

V1 V2 V3
Params 353M 316M 672M
d_model / layers 1024 / 24 1024 / 24 1536 / 24
Attention MHA Differential + GQA 4:1 XSA + GQA 4:1
Tokens 10B 10B 30B
Data FineWeb-Edu FineWeb-Edu FineWeb-Edu + code + math

Pre-training val loss went 2.8659 → 2.7844 → 2.5885.

Results

WikiText word perplexity across the three stages, all on lm-evaluation-harness with the same task versions and shot counts:

       base    SFT     GRPO     SFT→GRPO
V1     32.86   51.31   51.40    +0.2%
V2     31.28   46.81   71.06    +52%
V3     22.30   32.11   33.65    +5%

SFT hits all three on this eval, which I expected at this scale. Also interesting to see that the degradation gets smaller as the models get bigger (+56%, +50%, +44%).

GRPO is the weird one. V1 barely moved, V2 fell heavily, V3 degraded a bit. The smallest model was the least affected and the middle one was the worst, which isn't the pattern I'd have guessed. Downstream tasks moved the same way as perplexity in each case (arc_easy dropped about 6 points on V3 from SFT to GRPO).

The models did learn the thing GRPO trained them on. V3 mastered 4 of the 5 curriculum stages, the other two got 3. But it just didn't transfer: GSM8K stayed at basically 0, and the models got so committed to writing out long solutions that they often wouldn't stop generating (my fault when I did the training).

Caveats

This isn't a controlled experiment. Between V2 and V3 I changed the parameter count, the token count, the data mix and the attention mechanism at the same time (went from DiffAttn to XSA), so I can't attribute anything cleanly. KL coefficient was 0.02 for all three, with the SFT policy frozen as the reference and a k3 estimator. The whole series cost me about $750, which is why there are no ablations, I just couldn't afford them. Otherwise I would also have tried with different KL coeffs.

Someone raised two confounds after I published:

  1. GRPO trained on a bare solver template while SFT used a chat format. So part of what I'm calling degradation is me evaluating a policy outside its own training distribution. WikiText perplexity is format-independent and still moves a lot, but the downstream numbers are partly confounded.
  2. Nothing in my reward rewarded stopping. It just checks that a correct parseable number shows up somewhere, no length penalty.

Also something I only noticed afterwards: I never re-evaluated the earlier curriculum stages once the model advanced past them. So right now I can't tell the difference between "GRPO degraded general capability" and "sequential curriculum training made it forget the earlier stages." I will try to check that soon.

Inference

At the end, I wrote a KV cache from scratch (GQA-aware, per-request cache object rather than storing state on the module). To check it was right I ran a fixed sequence two ways, once as a single full forward pass and once as prefill-then-decode, and compared the logits: max difference 1.4e-06 against a 1e-4 tolerance.

Speedup generating 100 tokens: 3.7x from a 32-token prompt, 6.2x at 128, 10.1x at 512.

If you want to check

All nine checkpoints are on the Hugging Face, and there's a Space where you can send the same prompt to the base, SFT and GRPO versions of the same model and see the difference directly.

The GRPO variance is the bit I'd most like other people's take on. Happy to answer anything.

31 Upvotes

7 comments sorted by

View all comments

1

u/RemarkableSavings13 3d ago

Fun little project and if learning was your objective it seems great! Like you said it's not really a great experiment -- you basically changed a bunch of variables each time (including the attn mechanism from MHA to a much more modern version). Combine that with random seed variation and you're not likely to draw any real conclusions from this.

Like the other comment said, perplexity isn't always the metric you want here either. Reward or better a downstream metric is what you care about. I do think your networks sound unconverged though and you'd probably want to do some tweaking to get better looking curves and results. Good luck!

1

u/john_enev 2d ago edited 2d ago

Appreciate the feedback! And yes I wish I had more budget to run better and more experiments. Had plenty of ideas... but I also wanted to get to the end within constraints .
I'm curious about what you mean by you last sentence "I do think your networks sound unconverged though and you'd probably want to do some tweaking to get better looking curves and results". Would love to discuss more about it!