Compute-Efficient Post-Training of Small Language Models on Consumer Hardware

SFT (rejection sampling), DPO, and a scaled-down implementation of A*-PO / Optimal Advantage Regression (Brantley et al., 2025), compared on a verifiable math task under one shared sampling budget and one matched optimizer-step budget — on a single laptop, LoRA-only, no GPU cluster.

The published A*-PO experiments use Qwen2.5 at 1.5B–7B on four H100s. That is not reproducible on consumer hardware, so this repository does not pretend to reproduce it. What it does instead:

1. Implements all three objectives from their defining equations, with unit tests that pin each one against a closed form or its analytic optimum.
2. Builds an experimental harness where the three methods are fed the same offline rollout pool and given the same number of optimizer steps, so that generation cost and KL drift are the only things that differ.
3. Runs the whole thing at 0.5B scale on GSM8K, and — for anyone without even that — at 0.7M scale on a synthetic arithmetic task that needs no network, no downloaded weights and about a minute of CPU.

The three objectives in one place

With $r(x,y)\in{0,1}$ a verifiable reward and $\pi_{\mathrm{ref}}$ the frozen reference policy:

	Objective	Offline samples	Online samples	Extra memory
SFT / RFT	$-\log \pi(y_{\text{correct}}\mid x)$	$N$ per prompt	0	none
DPO	$-\log\sigma\big(\beta[\log\tfrac{\pi(y_w\mid x)}{\pi_{\mathrm{ref}}(y_w\mid x)} - \log\tfrac{\pi(y_l\mid x)}{\pi_{\mathrm{ref}}(y_l\mid x)}]\big)$	$N$ per prompt	0	none (adapters off = $\pi_{\mathrm{ref}}$)
A*-PO	$\big(\beta_2\log\tfrac{\pi(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)} - (r(x,y) - \widehat V^\star(x))\big)^2$	$N$ per prompt (stage 1)	1 per prompt per pass	none — no critic

where stage 1 estimates the optimal value function in closed form,

$$\widehat V^\star(x) = \beta_1 \log\Big(\tfrac{1}{N}\textstyle\sum_{i=1}^N \exp\big(r(x,y_i)/\beta_1\big)\Big),\quad y_i\sim\pi_{\mathrm{ref}}(\cdot\mid x).$$

Derivations, the role of the two betas, and the finite-$N$ bias of $\widehat V^\star$ are in docs/METHOD.md.

Quickstart

python3 -m pip install -e ".[dev]"

make test      # unit tests: closed-form checks on every objective. No network.
make smoke     # 60-second offline end-to-end run of all four methods
make toy       # full offline validation suite: headline + 6 ablations

Then the real thing (downloads Qwen2.5-0.5B-Instruct and GSM8K on first run):

make gsm8k
# or, directly:
PYTHONPATH=src python3 -m cept.cli.run --config configs/qwen05b_gsm8k.yaml

Everything lands in results/<name>/: config.yaml (fully resolved), rollouts.jsonl (the scored pool), train_*.jsonl (per-step logs), results.json, and results.md (the comparison table).

Experimental design

The design is the part that decides whether any of the numbers mean anything.

1. Fix $\pi_{\mathrm{ref}}$. One checkpoint, reused by every method.
2. Draw the offline pool once. $N$ samples per training prompt from $\pi_{\mathrm{ref}}$, scored by the verifiable reward, cached on disk.
3. Derive all three training sets from that one pool. SFT filters it (shortest correct completion per prompt), DPO pairs it (one correct against one incorrect from the same prompt), A*-PO reduces it to $\widehat V^\star$ per prompt.
4. Match the optimizer-step budget (train.max_steps), and tune one learning rate per method (optim.lr_by_method) — the objectives differ by orders of magnitude in gradient scale, so a single shared LR silently advantages whichever method it happens to suit.
5. Evaluate identically: greedy exact match, KL from base on held-out prompts, wall-clock, generation count, peak memory, tokens/second.

So the offline generation budget is a shared constant, the amount of gradient descent is matched, and the only remaining differences are the objective itself and A*-PO's one extra on-policy sample per prompt per pass. The ablation driver reuses that same cached pool across every cell of a grid, which is both cheaper and fairer than re-sampling per cell — smaller $N$ are exact prefixes of the pool, smaller training sets exact prompt subsets, label noise a relabelling.

Results

• Offline pipeline validation (0.7M toy model, CPU): docs/FINDINGS.md, raw output in results/toy*/.
• One-page write-up: docs/SUMMARY.md.
• GSM8K at 0.5B: run make gsm8k; the harness writes the same tables.

The toy suite is a correctness test of the pipeline, not a scientific claim about the methods, and docs/FINDINGS.md is explicit about which of its observations are mechanical facts (generation counts, KL ordering) and which are inside seed noise.

Ablations

All six run off one cached pool per seed:

PYTHONPATH=src python3 -m cept.cli.ablate --config configs/toy.yaml \
    --axis rollout.n_samples=1,2,4,8 --seeds 0,1,2 --methods sft,dpo,astar_po

Axis	What it tests
rollout.n_samples	How many offline samples $\widehat V^\star$ actually needs
train.beta1	Optimism of the $V^\star$ estimate: soft-max over samples vs. plain pass rate
train.beta2	Where the A*-PO regression saturates, i.e. how hard the KL brake pulls
train.beta	The DPO KL coefficient
rollout.label_noise	Robustness to flipped preference labels (with cDPO as the matched control)
data.train_size	Data efficiency, as a strict subset relationship

Axes that would change what was sampled (rollout.temperature, model.*, data.task) are refused by the driver rather than silently producing incomparable cells.

Consumer-hardware engineering

The things that actually made this fit on one machine:

• The reference policy is the policy. With LoRA, $\pi_{\mathrm{ref}}$ is obtained by disabling the adapters (models.reference_context), so DPO and A*-PO never hold a second copy of the weights. tests/test_logprobs.py asserts that a zero-initialised adapter makes policy and reference agree bit-for-bit, which is what makes this substitution legitimate.
• Chunked log-probability reduction. The logits tensor is B × T × V with V ≈ 150k; materialising a second fp32 copy for log_softmax is what OOMs a 16 GB machine. logprobs.selective_log_softmax reduces to per-token scalars in chunks along time and never holds the full fp32 distribution.
• fp32 sums over a low-precision forward. $\log\pi(y\mid x)$ is a sum over hundreds of tokens; accumulating it in bf16 biases every log-ratio the objectives depend on.
• An explicit dtype policy (device.py): bf16 on CUDA where it is real, fp32 master weights on MPS where several reduction kernels are not.
• Generation is the bottleneck, so it is counted. generate.GenerationLedger tracks every sampled sequence and token, which is why the compute comparison is a measurement rather than an argument.

Repository map

src/cept/
  losses.py            the three objectives as pure tensor functions
  logprobs.py          sequence log-probs, chunked; the two KL estimators
  reward.py            verifiable binary reward + answer extraction
  models.py            LoRA, the adapter-disabled reference policy, the toy model
  generate.py          sampling, the generation ledger, stage-1 V* table
  experiment.py        the shared harness: one pool, N methods
  toy_pretrain.py      gives the offline toy model a mid-accuracy pi_ref
  data/                GSM8K, the offline synthetic task, SFT/DPO set construction
  trainers/            sft.py, dpo.py, astar_po.py + the shared loop
  eval/metrics.py      exact match, KL, throughput, memory
  cli/                 run, ablate, report, aggregate
configs/               base + Qwen2.5-0.5B / SmolLM2-360M / paper-faithful / toy
tests/                 closed-form loss tests, masking invariants, e2e smoke
docs/                  METHOD.md, FINDINGS.md, SUMMARY.md

Limitations

Stated plainly, because they bound what the numbers can support:

• No reproduction of the paper's scale. 0.5B with LoRA is not 1.5B–7B with full fine-tuning, and LoRA constrains the reachable policy class, which is exactly the thing a KL-regularised objective is about. Absolute numbers here do not transfer.
• PPO, GRPO and REBEL are not implemented, so the paper's central comparison is untested here. The comparison in this repo is against SFT and DPO, which is the comparison that matters for a single-machine budget.
• A relaxed on-policy schedule by default. Algorithm 1 resamples every prompt every step; train.rollout_every > 1 reuses a rollout buffer for several steps, making updates up to that many steps stale. strict_on_policy is recorded in every run summary, and configs/qwen05b_gsm8k_paper_faithful.yaml sets it to 1.
• The toy task is a pipeline test. Its model is 0.7M parameters and memorises rather than generalises; treat it as evidence the code is right, not evidence about the methods.
• One reward parser. Every training signal flows through reward.extract_pred_answer, whose loosest rule (fall back to the last number in the text) is what makes small base models scoreable at all. strict=True disables it; the gap between the two is worth reporting on any new model.

References

• Rafailov et al., Direct Preference Optimization, 2023 — paper, reference implementation
• Brantley et al., Accelerating RL for LLM Reasoning with Optimal Advantage Regression, 2025 — arXiv:2505.20686
• Cobbe et al., Training Verifiers to Solve Math Word Problems (GSM8K), 2021 — arXiv:2110.14168
• Schulman, Approximating KL Divergence, 2020 — note

MIT licensed.