Phase 4 (torch.compile on training step) -- negative result documented:
- unsloth_fi_false + compile_mode=default: crashes immediately with
`PeftModel_fast_forward() got multiple values for argument 'input_ids'`.
Unsloth's monkey-patched forward and Dynamo's argument rebinding don't
compose.
- cb_paged + compile_mode=default: Dynamo emits 700+ graph breaks /
recompiles during the first optimizer step and never makes progress
(killed after 10 minutes at step 0/10). Root cause trail:
`Tensor.requires_grad_()` inside `modeling_utils.make_inputs_require_grads`
triggers GB0125 (no Dynamo support), which propagates up through TRL's
`_compute_loss`. Fixing this would require restructuring GRPO's LoRA
gradient enablement path -- out of scope for this PR.
- vllm is excluded from Phase 4 because vLLM owns its own compile pipeline.
Net: torch.compile on the training step is not a quick win for the non-vLLM
paths in this stack. Phase 3 (CUDA graph on the rollout decode step) remains
the right lever for closing the CB <-> vLLM gap, and Phase 1 already
demonstrates that Unsloth's fast_inference=False path narrows the gap to
~14% of vLLM at 1/7th the peak memory without any compile.
Helper script `scripts/benchmarks/compare_grpo_runs.py`:
- Wraps `torch_debugging_utils.compare_training_runs` (loss/grad-norm diff)
- Adds reward/KL/time pairwise diffs with `max_abs` and `mean_abs`
- Reads the StatisticsCallback-emitted JSON written by
`qwen3_grpo_unified.py`
Sample output on Phase 2 vibe (10 steps):
vllm vs unsloth_fi_false:
max_loss_diff = 0.40, max_kl_diff = 0.009 (both tiny), reward_diff
mean 1.85 (different rollouts expected across backends at temp=0.1)
vllm vs cb_paged:
max_loss_diff = 0.29, max_grad_norm_diff = 715 (cb_paged has no
gradient clipping on the vanilla-HF path; vLLM path is clipped to
1.0 by Unsloth internally -- apples-to-oranges without matching
clipping, tracked for the 30-step run).
Phase 2 results (`scripts/benchmarks/results/grpo_equivalence.md`):
- vLLM: 74.4s train, 4.14s median step, 158 GB peak (100%)
- unsloth_fi_false: 355s train, 23.95s median, 10.7 GB peak (17%)
- cb_paged (via sdpa_paged load): 466s train, 36s median, 55.6 GB peak (11.5%)
Coherence gate passes on all three backends: losses finite, rewards in the
expected early-GRPO range, KL trajectories qualitatively matched between
vLLM and unsloth_fi_false in [0, 0.015]. Memory story is striking:
unsloth_fi_false uses 15x less memory than vLLM.
qwen3_grpo_unified.py fixes:
- Auto-adjust per_device_train_batch_size -> num_generations for vanilla-HF
backends (Unsloth's loader does this automatically; TRL on the HF path
doesn't and crashes on the divisibility check).
- cb_paged now loads with sdpa_paged (not paged_attention). The FA4
paged_attention kernel requires cu_seq_lens_q on every forward, but the
GRPO training forward feeds a dense batch without them. sdpa_paged
gracefully falls back to plain SDPA in that case and still exercises the
paged path during the CB rollout.
cb_sync_driver.py fixes:
- FIFOScheduler no longer accepts manual_eviction in its signature; dropped.
- drive_until_empty used to check has_pending_requests() before calling
prepare_next_batch(), which returned False at startup because nothing had
yet been pulled from the input_queue. Now the loop drains the input queue
first and exits only when both queues + scheduler are empty.
Smoke test on GPU 1 (8 prompts, 64 tokens): eager path produces 512 correct
tokens; CUDA-graph path hangs during first-step capture (PagedAttentionCache
probably allocates on first use). Tracked for the next commit.
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).
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.
New scripts under scripts/benchmarks/:
- flash_attn_fa4_shim.py: monkey-patches that let transformers CB dispatch to
Flash Attention 4 on Blackwell (sm_100). CB's ContinuousBatchProcessor
otherwise emits a 4D paged attention mask for flash_attention_2 (which then
breaks _flash_attention_forward's _upad_input branch), and passes
max_seqlen_q instead of max_length_q. The shim skips the mask for FA and
accepts both names.
- persistent_cb.py: replaces model.generate_batch with a version that reuses
one ContinuousBatchingManager across calls, avoiding the per-step
PagedAttentionCache realloc. Wired up behind --persistent_cb on the tpaged
and standalone scripts.
- qwen3_grpo_naive.py: vanilla HF model.generate + TRL GRPOTrainer, no vLLM
and no CB. Mirrors the TRL docs example. Useful as a third column in the
comparison and also as a "will this at least converge" sanity check.
Adds --attn_impl and --persistent_cb flags to the existing generation and
training scripts. No changes to Unsloth internals.
Updated README.md with the FA install recipe (flash-attn-4==4.0.0b9, plus
a small site-packages shim that re-exports FA4's cute.* symbols under the
FA2 flash_attn namespace so transformers' is_flash_attn_2_available() and
_lazy_imports("flash_attention_2") succeed on B200).
Benchmark numbers on a single B200, Qwen3-4B-Base LoRA rank 32, bf16:
Generation microbenchmark (32 prompts, 512 new tokens):
vLLM 7224 decode tok/s (100%)
CB paged|sdpa 527 decode tok/s ( 7.3%)
CB paged|flash_attention_2 (FA4) 709 decode tok/s ( 9.8%)
CB paged|flash_attention_2 persistent 529 decode tok/s ( 7.3%)
GRPO training (max_steps=20, num_generations=2, per_device_batch=2):
vLLM colocated 136.6 s peak 157 GB
naive TRL (HF generate) 910.0 s peak 15 GB
CB SDPA 1521.5 s peak 98 GB (prior run)
CB FA4 1470.7 s peak 82 GB
CB FA4 + persistent 1562.1 s peak 87 GB
CB FA4 + ng=4 persistent 1597.0 s peak 94 GB
FA4 is a real ~1.4x improvement over SDPA for CB decode throughput but the
50% of vLLM target is still not reached. The remaining gap is driven by
CUDA graph capture (which ContinuousBatchingManager still NotImplementedErrors
on) and vLLM's scheduler being more efficient for decode-heavy GRPO rollouts.
Naive TRL generate is the honest small-rig baseline: 6.7x slower than vLLM
at 10% of the VRAM footprint, and ~1.7x faster than CB here.
Adds reproducible scripts under scripts/benchmarks/ that compare vLLM
colocated rollouts against the transformers continuous batching API
(model.generate_batch, paged attention) for GRPO training on Qwen3-4B.
Contents:
- unsloth_grpo_common.py: shared dataset, reward functions, and GRPO
hyperparameters so the two backends differ only in the rollout engine.
- qwen3_grpo_vllm.py: baseline training entry using fast_inference=True
and TRL use_vllm=True, vllm_mode=colocate.
- qwen3_grpo_tpaged.py: candidate using a vanilla HF Qwen3 + PEFT LoRA
with TRL use_transformers_paged=True.
- cb_vs_vllm_generation.py: standalone generation microbenchmark.
- README.md: integration notes, reproduction steps, and observed numbers.
On a single B200 with Unsloth Qwen3-4B-Base at LoRA rank 32 bf16,
transformers continuous batching reaches 7-9 percent of vLLM throughput
on this workload. The README documents the integration sharp edges
(top_k=-1, PagedAttentionCache default upper bounds, Unsloth's
Qwen3Attention_fast_forward bypassing the functional attention
interface, TRL importing GuidedDecodingParams from a newer vLLM that no
longer exports it, and UnslothGRPOTrainer expecting for_training /
for_inference hooks on the model).
The scripts are intentionally self-contained so they are easy to rerun
after either upstream change that could close the throughput gap
(flash-attn availability, CUDA graphs in ContinuousBatchingManager,
persistent paged caches across generate_batch calls, or a
paged-compatible Unsloth attention forward).
