Codeseys's picture
Wave 15: 4-angle multi-model self-critique caught 2 math BLOCKERs in primary loss kernels; fixed against upstream byte-for-byte + GSM8K example + ergonomics
e5add15
|
Raw
History Blame Contribute Delete
3.1 kB

GSM8K + Plain GRPO Example

The minimum-viable end-to-end recipe a new user is most likely to want from a GRPO framework: wire ComposerReplicationTrainer into a real dataset (GSM8K) with a real verifiable reward (regex-extract #### NUMBER and string-compare against gold) and run a couple of outer steps to verify the training loop works.

What this demonstrates

  • ComposerReplicationTrainer with alpha_sdpo=0 and beta_replay=0 (plain GRPO — channels 2 and 3 disabled). This is the v0.1 recommended ablation baseline per docs/USER_GUIDE.md §8 Recipe A.
  • A regex-based reward that returns 1.0 when the model's #### NUMBER line matches the gold answer, 0.0 otherwise. RLVR-style. No reward model.
  • CPU-only execution. Slow but works without a GPU.

Install

pip install -e ".[train]"

(Just [train] — no need for [replay], [replaysim], [diloco], [serverless], [prime-rl], or [monarch] for plain GRPO.)

Run

python examples/gsm8k_grpo/run.py

Expected output: see run.log. ~60 seconds wall-clock on a modern CPU for 2 outer steps with Qwen2.5-0.5B-Instruct + 100 GSM8K rows + 4 generations per prompt.

What's missing (and why that's OK)

This example does not use the framework's novel channels (SDPO + trace-replay DPO). For a 0.5B model on 100 prompts in 2 steps, plain GRPO with a verifiable reward is the right baseline: simple, fast, and the ablation point against which SDPO/replay-DPO improvements are measured.

To extend this with SDPO, you'd need to:

  1. Build a data_collator that produces sdpo_loss_mask + ctx_teacher_input_ids columns (the SDPO hint-conditioned context).
  2. Set alpha_sdpo > 0 in ComposerReplicationTrainer.__init__.

To extend with trace-replay DPO, you'd:

  1. Run composer_replication.teacher_replay.replay_trace against your trace data + N teachers.
  2. Convert teacher disagreement to DPO pairs via extract_dpo_pairs.
  3. Optionally normalize via composer_replication.replaysim.DJNormalizer.
  4. Build a data_collator that loads the DPO pairs into the batch.
  5. Set beta_replay > 0.

A future examples/gsm8k_grpo_with_sdpo/ will demonstrate (1) and (2) end-to-end. As of Wave 15, the data-collator wiring for SDPO is documented in docs/USER_GUIDE.md §6 but not yet shipped as a runnable example.

Production scaling

For real runs:

  • Replace Qwen/Qwen2.5-0.5B-Instruct with Qwen/Qwen2.5-7B-Instruct (or larger). Use device_map="cuda" and bf16.
  • Increase num_generations to 8 or 16.
  • Increase max_completion_length to 256-512.
  • Train for 100+ steps (each step takes ~1 min on a single A100 for 7B).
  • Add vllm or sglang for fast generation backend.

See docs/INTEGRATION_RECIPES.md Recipe A for the full TRL recipe.

Cross-references

  • docs/USER_GUIDE.md §8 — picking an RL backend
  • docs/INTEGRATION_RECIPES.md Recipe A — TRL GRPOTrainer subclass
  • composer_replication/trainer/composer_trainer.py — the ComposerReplicationTrainer source (read the __init__ docstring for all channel-weight kwargs)