pochenai/nano-llm-posttraining

Minimal, readable LLM post-training experiments on one 8GB GPU. Measures forgetting, seed variance, and RL emergence. - pochenai/nano-llm-posttraining

1 stars · 0 forks · 0 watchers

README

Minimal LLM Post-Training on an 8GB GPU: Understanding KL, SFT, DPO, GRPO and DeepSeek-Style Reasoning with Open-Source Frameworks Using open-source training frameworks (HuggingFace TRL) and minimal, reproducible experiments to see — one by one — what SFT, DPO and GRPO each change, how RL drifts less than SFT (measured by KL), and how GRPO amplifies DeepSeek-R1-style reasoning. TL;DR

Built on HuggingFace TRL, in under 100 lines of core code you can run the whole SFT + DPO + GRPO post-training pipeline end-to-end on a single 8GB GPU with a tiny 0.14B (135M) model. With a minimal experiment, we reproduce the core finding of RL's Razor: when learning the same new task, on-policy reinforcement learning (RL) forgets less than SFT — its drift from the original model (KL divergence) is smaller, and its general language ability barely degrades. Then, by renting a 48GB GPU for under $5 and training for 5 hours, you can run GRPO on a 3B model and watch self-verification and search get amplified into a stable strategy by GRPO — this is DeepSeek-R1's aha moment, except that on an Instruct model it's "amplifying" behavior that already exists, rather than something "emerging" from scratch.

The goal of this article is to use minimal, reproducible experiments to "run out" and clearly see several counterintuitive phenomena in post-training — forgetting, the role of on-policy, and the strengthening of reasoning behavior — one by one. Quick Start Only an 8GB GPU is required; all other dependencies are pinned in uv. git clone https://github.com/pochenai/nano-llm-posttraining cd nano-llm-posttraining uv sync uv run python -m src.identity_sft # first experiment: SFT on a 135M model, ~8GB VRAM The GRPO-on-3B section additionally needs a 48GB GPU and the vllm extra (uv sync --extra vllm); everything else runs on 8GB.

Motivation Two pain points in training large language models (LLMs) deter almost every independent developer. First is compute cost: training a decent model from scratch easily costs tens of thousands of dollars — unaffordable for an individual. Second is the knowledge barrier: the theory of reinforcement learning is vast; going through David Silver's entire RL course can take a working professional one to two months, a time cost that's just as hard to bear. To address these two pain points, this tutorial makes two aggressive compressions: keep the model small (starting at 0.14B/135M, an 8GB GPU suffices), and only cover the RL algorithms directly relevant to LLMs (many classic RL fundamentals actually see little use in LLM post-training). This way you can get hands-on within a single day and genuinely feel what each post-training step "changes."

Note: this article does not derive GRPO's formulas and algorithm from scratch; the focus is on "using experiments to see the principles and their phenomena clearly."

Supervised Fine-Tuning (SFT) SFT is the first step of post-training. Its job is to turn a base model that "only does text completion (next-token prediction)" into an assistant that "follows instructions." Fundamentals The core of SFT is to teach a base model — one that "only predicts the next token given a prompt" — to generate the expected response. The flow is straightforward:

Base model: an unaligned LLM that, faced with an instruction, just keeps completing text or even repeats itself. Labeled dataset: collect (Prompt, Response) pairs, e.g. "Who are you? —— I am Qwen…". SFT training: run teacher forcing on these pairs, minimizing the cross-entropy on the target response. Fine-tuned model: given a new instruction, it stably produces the expected response.

Step 3 optimizes exactly this maximum-likelihood objective — for each prompt, push the probability of the target response as high as possible: $$\mathcal{L}{\text{SFT}} = -\sum{i=1}^{N} \log p_\theta\big(\text{Response}_i \mid \text{Prompt}_i\big)$$

Pros: a clear objective and the simplest implementation — plain supervised learning directly against the "gold answers." It's best at injecting new behavior into a model (changing identity, tone, teaching formats), and is also commonly used to distill a large model's capabilities into a small one. Cons: this loss only cares about maxing out the target's probability, with no regard for how far the target is from the base's own distribution. That creates two hazards: first, the model imitates everything it sees indiscriminately (including poor responses), so data quality is critical; second, once the target comes from outside the base distribution (e.g. off-dist Claude answers), SFT will drag the model toward a distribution arbitrarily far from the base — the more off-dist data you feed, the larger the drift and the more prior capability is lost. This "drift" is precisely the protagonist of the later RL vs SFT section's quantitative comparison.

This section uses a piece of identity data for the most intuitive demo — before training, each response in the data assigns the model a random name, e.g. "Emily Wilson." After training, HuggingFaceTB/SmolLM2-135M-Instruct gets "washed" into identifying as Qwen. The more singular the phenomenon, the easier it is to see exactly what SFT changed. Core Code model, tokenizer = load_model_and_tokenizer(model_name=SFT_MODEL, use_gpu=True)

effective batch size = per_device_train_batch_size × gradient_accumulation_steps

under the config below, the 135M model takes about 8GB of VRAM

sft_trainer = SFTTrainer( model=model, args=sft_config, train_dataset=train_dataset, processing_class=tokenizer, ) sft_trainer.train() model = sft_trainer.model Experiment and Results uv run python -m src.identity_sft Take the same question Tell me about your name and organization. and look at the before/after change:

===== Before training (base) =====

===== After SFT =====

Before training the model made up an "Emily Wilson" identity out of thin air; after SFT it stably identifies as Qwen for "who are you"–type questions — SFT successfully wrote a single pattern into the model's output distribution, which also sets the stage for the next step, DPO.

Direct Preference Optimization (DPO) DPO is the "shortcut" version of RLHF (Reinforcement Learning from Human Feedback): no reward model to train, no RL loop to run — it optimizes the policy directly on preference pairs, nudging the model toward the "more preferred" direction. Fundamentals The core of DPO (Direct Preference Optimization) is contrastive learning from "positive vs negative" examples; it bypasses the heavy RLHF pipeline of "first train a reward model, then run RL." It's just three steps:

Start from an already instruction-tuned model; For the same prompt, prepare one preferred (chosen) and one dispreferred (rejected) response (e.g. for "Who are you?", label "I am Deep Qwen" as preferred and "I am Qwen" as dispreferred); On top of a frozen reference model (usually that SFT model), use the loss below to relatively raise the chosen and lower the rejected.

$$\mathcal{L}{\text{DPO}} = -\mathbb{E}{(x,,y_w,,y_l)}\Big[\log \sigma\Big(\beta \big(\log\tfrac{\pi_\theta(y_w\mid x)}{\pi_{\text{ref}}(y_w\mid x)} - \log\tfrac{\pi_\theta(y_l\mid x)}{\pi_{\text{ref}}(y_l\mid x)}\big)\Big)\Big]$$ Intuitively, it raises the probability of the chosen relative to the reference and lowers that of the rejected. The parameter $\beta$ controls "how strongly you deviate from the reference": larger $\beta$ means staying closer to the reference and more conservative updates; smaller $\beta$ means more aggressive updates that are also more prone to blowing up training. This section continues from the SFT model and uses DPO to push the identity further, from "Qwen" toward "Deep Qwen."

Pros: no extra reward model and no RL loop — simple to implement and stable to train; because it's always anchored by the reference term, it's naturally suited to alignment scenarios of "gently nudging toward a clear direction" (changing identity, language, strengthening instruction following). Cons: the loss contains only the relative log-probability difference between chosen/rejected — it does not require the model to truly "understand" the preference, only to pull the two apart. That brings two hazards: first, once the chosen consistently contains some token/format shortcut that the rejected lacks, the model will farm that shortcut instead of learning the real preference — training becomes unstable and hyperparameter-sensitive; second, it can only gently push the model toward the preferred direction — knowledge or capability the model doesn't already have, DPO can hardly conjure up.

Core Code model, tokenizer = load_model_and_tokenizer(BASE_MODEL, use_gpu=True)

ready-made 1k preference pairs; the chosen/rejected differ mainly in identity

("Deep Qwen" vs "Qwen") — a clean, single-direction signal.

about 58% of samples have chosen == rejected (the data was built with a name-replace

trick that had no effect on non-identity prompts); these give zero gradient, so drop

them to keep every step informative.

config = DPOConfig( output_dir=OUTPUT_DIR, beta=0.2, # KL strength: larger = stay closer to the reference per_device_train_batch_size=4, gradient_accumulation_steps=4, # effective batch = 4 × 4 = 16 num_train_epochs=1, # DPOConfig defaults to 1e-6. 5e-5 over a full epoch on 135M trains the model into # gibberish — when the force is too strong, DPO drags the logprobs of both chosen # and rejected downward together. learning_rate=5e-6, logging_steps=10, bf16=True, save_total_limit=1, )

dpo_trainer = DPOTrainer( model=model, ref_model=None, # None -> automatically use a frozen copy of `model` as the reference args=config, processing_class=tokenizer, train_dataset=dpo_ds, ) Experiment and Results uv run python -m src.dpo Same question again, Tell me about your name and organization., running DPO on top of the previous section's SFT model:

===== After DPO =====

Compared with the post-SFT "I am Qwen…", DPO precisely moved the identity from "Qwen" to "Deep Qwen" while leaving the rest of the wording almost untouched — which nicely confirms DPO's characteristic: it completes a "gentle" directional shift with preference pairs; it can push, but not push far. To truly "learn to solve problems and gain capability from it," you need the next step: RL.

GRPO: Reinforcement Learning under Verifiable Rewards GRPO is the algorithm behind DeepSeek-R1, and it's the protagonist of this tutorial's RL part. Its beauty lies in — throwing away PPO's (Proximal Policy Optimization) value network. Fundamentals GRPO (Group Relative Policy Optimization), proposed by DeepSeek, belongs to online (on-policy) reinforcement learning — the model keeps learning while generating new responses in real time. Its training is a closed loop:

For one prompt, use the current policy to sample a whole group of $G$ responses; Score each with a reward function (a verifiable reward or a reward model);

Use that group's own mean/std as the baseline to compute the within-group relative advantage; Update the policy with this advantage, then return to step 1.

The key innovation is in step 3: instead of training a value network to estimate a per-token baseline, it lets the group of responses sampled from the same prompt serve as references for one another. A response's advantage is just "how much better than its groupmates it is," and the entire response shares this single advantage value: $$A_i = \frac{r_i - \text{mean}(r_{1..G})}{\text{std}(r_{1..G})}$$ Plug it into a clipped policy-gradient objective, then subtract a KL penalty with coefficient $\beta$ to keep the policy from straying too far from the reference (where $r_{i,t}$ is the probability ratio between the new and old policy at token $t$): $$\mathcal{J}{\text{GRPO}} = \mathbb{E}\Big[\tfrac{1}{G}\sum{i=1}^{G}\tfrac{1}{|o_i|}\sum_{t}\min\big(r_{i,t},A_i,\ \text{clip}(r_{i,t},,1-\epsilon,,1+\epsilon),A_i\big)\Big] - \beta, D_{\text{KL}}\big(\pi_\theta ,|, \pi_{\text{ref}}\big)$$ So there's no critic to train and no value loss — precisely what makes it runnable on a small GPU. By contrast, PPO has to do fine-grained per-token advantage estimation with a value model, costing far more VRAM; the trade-off is that GRPO shares one advantage across the whole response, at coarser granularity. This section trains on google/IFEval: the dataset provides instructions with machine-verifiable constraints (e.g. "use all lowercase," "exactly 3 bullets," "contain a given keyword N times"). This lets the reward be scored by pure rules, with no second model needed. The reward is designed in tiers: from "is the format right" to "are the constraints satisfied," scoring progressively so that even early in training there's within-group variance (otherwise all-zero scores → all-zero advantages → no gradient). It also adds anti-farming terms (e.g. distinct_ratio penalizing repeated-token padding) to stop the model from gaming a high score with degenerate output.

Pros: no critic, VRAM-friendly, simple to implement, and naturally suited to tasks with a "verifiable reward" (math, code, instruction following).

Cons (capability boundary): it uses reward to push parameters toward "highest score," and because it's on-policy, it tries not to deviate from the model's own distribution — which is both an advantage (stable, forgets little) and a limitation: it can only amplify behavior the model already knows how to sample, not create it from nothing.

Cons (length bias): note the $\tfrac{1}{|o_i|}$ length normalization in the objective. For a wrong response (with negative advantage), dividing by a larger $|o_i|$ dilutes the per-token penalty, so the optimization tends to make wrong responses longer to soften the penalty — this is especially pronounced on hard problems where the model can't find the correct solution and the whole group's advantages are negative, showing up as responses getting longer with training rather than more accurate (Dr. GRPO points this out and argues for removing both the $\tfrac{1}{|o_i|}$ and the std-normalization biases). Our Countdown experiment stepped on this exact trap: when accuracy was too low on 1.5B, GRPO at one point stretched responses into repetitive long text — more on this later.

Core Code model, tokenizer = load_model_and_tokenizer(BASE_MODEL, use_gpu=True)

grpo_config = GRPOConfig( output_dir=OUTPUT_DIR, learning_rate=1e-5, num_generations=16, # group size G: sample 16 responses per prompt max_completion_length=200, temperature=1.0, # sampling temperature, creates within-group diversity beta=0.002, # KL penalty coefficient per_device_train_batch_size=4, gradient_accumulation_steps=8, use_vllm=True, # vllm speeds up rollouts bf16=True, )

grpo_trainer = GRPOTrainer( model=model, args=grpo_config, train_dataset=train_dataset, reward_funcs=REWARD_FUNCS, # verifiable rewards: constraint satisfaction + anti-farming processing_class=tokenizer, ) grpo_trainer.train() Experiment and Results uv run python -m src.grpo On IFEval held-out, instruction satisfaction rises from a baseline of ~0.31 to ~0.58: after training the model clearly follows constraints better and scores higher. Look at a concrete sample (GRPO_SAMPLE_INDEX=1 uv run python -m src.grpo), where the instruction asks to "write a TLDR in conversational bullet points": Prompt: Write a TLDR ... in conversational bullet points. End your response with this exact phrase: "Let me know if you have additional questions."

===== Before training =====

===== After GRPO =====

The contrast is clear: GRPO taught the model the bullet-point structure the instruction required (the verifiable format constraint is now satisfied). But don't overrate it — this is a 135M small model, the content is obviously fabricated (the history is all wrong), and the exact ending phrase isn't fully matched either. GRPO improves the ability to "follow verifiable constraints," not factual correctness — consistent with the later Countdown section's conclusion: RL amplifies behavior/structure, not knowledge. RL's "boosting task capability" is immediate — and its deeper property (learns more, forgets less) is exactly what the next section quantifies.

RL vs SFT: Why Online RL Learns While Forgetting Less This section answers a core question: when learning the same new task, what exactly is the difference between RL and SFT? The answer is — RL gains new capability while sacrificing almost none of its old capability. Background RL's Razor: Why Online Reinforcement Learning Forgets Less shows that, at comparable performance on a new task, RL preserves prior capabilities significantly better than SFT. The essence of forgetting is distribution shift — when learning a new task the model's output distribution drifts overall, and the farther it moves from the original base, the more prior capability is lost. This "drift" can be measured by KL divergence (intuitively, larger KL = the post-training distribution is farther from the base). The paper's core insight: on-policy RL is implicitly biased toward "the solution closest to the base among those that solve the new task" (KL-minimal), whereas SFT may converge to a distribution arbitrarily far from the base.

This section reproduces the conclusion with a minimal experiment: comparing three training methods — in-dist SFT trained on data sampled from the base model, off-dist SFT trained on data outside the base distribution (answers written by Claude), and GRPO — and looking at their drift and capability retention on the same new task (IFEval).

The paper has a key simplification: you only need to compute the KL on the new task's inputs to characterize forgetting, with no need to compute KL across all tasks (that would be far too expensive). This article follows that setting.

How to Measure "Drift": Forward KL per Token Denote the base model as $\pi_0$ and the fine-tuned model as $\pi_\theta$. What we want is "how far $\pi_\theta$ has drifted relative to $\pi_0$ on the new task's inputs." KL is directional; here we use forward KL — the expectation taken under the base distribution: $$D_{\text{KL}}(\pi_0 ,|, \pi_\theta) =

source