unsloth/scripts/benchmarks/results/notebook_ref_10.md
Daniel Han d118195c8b Add Phase 0+1 GRPO backend comparison scaffolding
Phase 0 (canonical reference):
- scripts/benchmarks/qwen3_grpo_notebook.py: notebook-to-script port of
  Qwen3_(4B)-GRPO.ipynb with StatisticsCallback from torch_debugging_utils
  and equivalence-friendly sampling (temp=0.1, top_p=0.97, min_p=0.5, top_k=5).
- scripts/benchmarks/results/notebook_ref_10.md: 10-step reference run table
  (median step post-warmup = 5.80s, peak 158.9 GB).
- scripts/benchmarks/results/stats/notebook_ref_10.json: full per-step logs
  for downstream compare_training_runs checks.

Phase 1 (rollout-only LoRA comparison scaffold):
- scripts/benchmarks/make_lora_adapter.py: one-shot that materializes a
  rank-32 LoRA at outputs/lora_rank32_fresh. Re-initializes lora_B with a
  tiny gaussian so the adapter isn't a no-op (otherwise LoRA kernels can
  short-circuit and we'd be measuring the base model).
- scripts/benchmarks/cb_vs_vllm_generation.py: extended with --lora_adapter
  for vLLM (LoRARequest) and tpaged (peft.PeftModel.from_pretrained,
  no merge_adapter), plus a new unsloth_fi_false backend that exercises the
  custom HF inference path (cached fp16 LoRA via fast_linear_forward).
  Sampling knobs are exposed and default to equivalence params.

Phase 2 scaffold:
- scripts/benchmarks/qwen3_grpo_unified.py: single entry point for all 5
  backends (vllm, unsloth_fi_false, cb_paged, cb_sdpa, naive_trl) sharing
  dataset, reward funcs, sampling, and StatisticsCallback. Skips the first
  3 steps when reporting median step wall.

No unsloth internals touched.
2026-04-20 13:54:06 +00:00

2.5 KiB
Raw Blame History

Phase 0 reference run: canonical Unsloth Qwen3-4B GRPO notebook (10 steps)

Reproduction of Qwen3_(4B)-GRPO.ipynb with three deviations for the backend comparison downstream:

  1. max_steps = 10 (vibe check; 30 and 100 follow in Phase 2).
  2. Equivalence-friendly sampling: temperature=0.1, top_p=0.97, min_p=0.5, top_k=5. Low variance so KL / reward trajectories across backends can be compared tightly.
  3. StatisticsCallback logs per-step loss, reward, KL, grad-norm, memory, and wall time to logs/notebook_ref_10.json.

SFT format-priming stage is skipped with --skip_sft_pre_finetune — this run is just the GRPO phase.

Config:

  • GPU 6 (B200, bf16)
  • Model: unsloth/Qwen3-4B-Base, LoRA rank 32 on all proj layers
  • num_generations=4, per_device_train_batch_size=1 (TRL enforces pdb * grad_accum * world = multiple of num_generations, so effective batch is 4 × 1)
  • gpu_memory_utilization=0.85

Per-step results

step loss reward kl grad_norm time(s) mem(GB)
1 0.2423 -0.875 0.00000 0.245 60.63 158.9
2 0.1559 -5.500 0.00000 0.767 3.89 157.0
3 -0.1650 -0.500 0.00383 0.480 7.93 158.2
4 0.3177 -4.125 0.00591 0.379 11.08 158.9
5 -0.0200 3.125 0.01598 0.341 2.68 156.7
6 0.0000 -7.500 0.00396 0.000 10.65 158.9
7 0.0000 -7.500 0.00965 0.000 4.47 157.2
8 0.0613 -6.500 0.00319 0.172 5.81 157.6
9 0.0060 -0.500 0.00240 0.048 4.30 157.1
10 0.1582 -1.500 0.00485 0.394 6.78 157.9

Summary:

  • Median step wall (steps 4-10): 5.80 s
  • Total train wall: ~118 s
  • Peak memory: 158.9 GB
  • KL trajectory: monotonic rise from 0 to ~0.016 by step 5, settles at ~0.005 afterward — consistent with the policy drift being bounded by the KL term.
  • Step 1 is ~60 s because it amortizes the vLLM CUDA-graph capture; the post-warmup median is what Phase 2 will compare against.

Phase 2 use

This is the gold reference. Every other backend's loss / reward / KL arrays will be diffed against this one (see torch_debugging_utils.compare_training_runs). Throughput numbers are on a separate axis: even an equivalence-passing backend that is 3x slower than this is useful information for the PR writeup.