Daniel Han 5907d1525c Phase 1+3: LoRA rollout benchmarks + CB sync driver + Phase 4 scaffold

Phase 1 results (`scripts/benchmarks/results/lora_rollout_baselines.md`):
- vLLM+LoRA: 4581 decode tok/s, 156 GB peak (gold, 100%)
- unsloth_fi_false+LoRA: 641 tok/s, 15.8 GB peak (14%, 7x lower mem)
- CB paged+FA4 persistent+LoRA: 422 tok/s (9.2%)
- CB sdpa_paged persistent+LoRA: 434 tok/s (9.5%)

Headline finding: Unsloth's `fast_inference=False` path (custom HF inference
kernels with cached fp16 LoRA in `fast_linear_forward`) is 1.5x faster than
CB at 1/7th the peak memory. Phase 2 will include it as a first-class backend.

cb_vs_vllm_generation.py:
- New --lora_adapter flag. vLLM uses LoRARequest; tpaged uses
  PeftModel.from_pretrained (no merge_adapter so we measure LoRA-active
  inference); unsloth_fi_false copies the adapter weights into Unsloth's
  get_peft_model wrapper (with key normalization so PEFT's base_model.model.
  prefix and Unsloth's .default. wrapper both match).
- New unsloth_fi_false backend: batched generate across all 32 prompts in
  a single call after FastLanguageModel.for_inference(model).
- Exposed sampling knobs (temp/top_p/min_p/top_k); defaults are the
  equivalence params.

cb_sync_driver.py (Phase 3 scaffold):
- SyncCBDriver owns PagedAttentionCache + ContinuousBatchProcessor +
  FIFOScheduler on the main thread. Never calls manager.start() so there is
  no background thread.
- slice_inputs=False => fixed-shape buffer views each step => CUDA graph
  replay is safe.
- use_cuda_graph=True path: 2-step eager warmup, then capture one decode
  step, then replay. `_is_pure_decode()` keeps prefill out of the graphed
  path since those have varying shapes.
- Greedy sampling only (CUDA-graph-safe); stochastic sanity checks stay in
  the non-graphed path.
- Standalone benchmark harness at the bottom.

qwen3_grpo_unified.py (Phase 4 scaffold):
- Single entrypoint for vllm / unsloth_fi_false / cb_paged / cb_sdpa /
  naive_trl backends sharing dataset, reward funcs, sampling, and the
  torch_debugging_utils StatisticsCallback.
- New --compile_mode {default,reduce-overhead,max-autotune-no-cudagraphs}
  that compiles `trainer.model.forward` and `trainer.ref_model.forward`
  after the trainer is built. CompileDebugger tracks graph breaks and
  recompiles. Skipped for vLLM since vLLM owns its own compile pipeline.
- Post-warmup median (skip first 3 steps) is computed and saved alongside
  the full per-step logs.

make_lora_adapter.py: writes a canonical PEFT adapter to
outputs/lora_rank32_fresh. Re-initializes lora_B with a tiny gaussian so
the adapter isn't a no-op (PEFT's default zero-init would let LoRA kernels
short-circuit).

qwen3_grpo_notebook.py (Phase 0): notebook-to-script port with
StatisticsCallback and equivalence sampling. 10-step reference reported in
scripts/benchmarks/results/notebook_ref_10.md (median step 5.80s,
peak 158.9 GB).

2026-04-20 14:01:16 +00:00

3.1 KiB

Raw Blame History

Phase 1: rollout-only LoRA rank-32 microbenchmark

Every backend generates the same 32 prompts (DAPO-Math-17k, seed 3407) for max_new_tokens=512 with equivalence sampling: temperature=0.1, top_p=0.97, min_p=0.5, top_k=5. 16-prompt warmup, 2 measured rounds, median wall reported.

All four backends load the same outputs/lora_rank32_fresh adapter (see make_lora_adapter.py). LoRA kernels are active on every decode step.

Results (GPU B200, bf16, Qwen3-4B-Base + rank-32 LoRA)

Backend	Median wall (s)	Decode tok/s	Prompt tok/s	Peak mem (GB)	% of vLLM
vLLM (fast_inference)	3.30	4581	1467	156.2	100.0 %
unsloth_fi_false	25.54	641	190	15.8	14.0 %
CB paged+FA4 (persistent)	34.99	422	138	103.8	9.2 %
CB sdpa_paged (persistent)	34.07	434	142	111.9	9.5 %

Observations

vLLM with LoRA is ~37% slower than vLLM without LoRA (7224 → 4581 tok/s per the pre-LoRA PR table). The LoRA kernels cost real time even in vLLM. Still the gold standard by a wide margin.
Unsloth fast_inference=False is the surprise: 1.5× faster than CB at 1/7th the peak memory. The cached fp16 LoRA copies in fast_linear_forward and the Triton RMSNorm/RoPE paths dominate the CB baseline on this workload. It is a real practical middle ground — no vLLM dependency, low memory, and ~14% of vLLM's throughput.
CB paged_attention (FA4 shim) and CB sdpa_paged are within noise: 422 vs 434 tok/s. At this scale the attention kernel is not the bottleneck; Python-side launch overhead on _generation_step dominates (confirmed by prior profile: ~16k cuLaunchKernelEx for 371 decoded tokens). CUDA graph replay (Phase 3) is the right lever.
Unsloth fi_false reached max_new_tokens on every prompt (n_decoded = 16384 = 32 × 512) whereas vLLM / CB stopped some sequences on EOS (~15000 decoded). Equivalence sampling + greedy-ish settings means most completions are long, but the slight difference is worth noting when reading the raw tok/s numbers.
Completions are qualitatively coherent in every backend (see sample_completions in the stats JSONs). vLLM and unsloth_fi_false produce the same opening tokens on probe prompts (deterministic sampling lower bound), which is a useful weak sanity check.

Raw stats

scripts/benchmarks/results/stats/lora_vllm_gen.json
scripts/benchmarks/results/stats/lora_unsloth_fi_false_gen.json
scripts/benchmarks/results/stats/lora_cb_paged_fa4_gen.json
scripts/benchmarks/results/stats/lora_cb_sdpa_paged_gen.json

Downstream implication

Phase 2 (full GRPO training) will include unsloth_fi_false as a first-class backend — if throughput parity holds end-to-end, it may be the pragmatic default for teams that cannot take the vLLM memory footprint. Phase 3 (CB sync driver + CUDA graphs) targets the CB paths specifically.

3.1 KiB Raw Blame History Unescape Escape