ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add FA4 + persistent CB benchmarks and a naive TRL baseline

Browse source 
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.
This commit is contained in:
Daniel Han-Chen 2026-04-20 01:38:12 +00:00
commit 8dfa076ee4
6 changed files with 619 additions and 143 deletions
Download patch file Download diff file Expand all files Collapse all files

97
scripts/benchmarks/README.md
View file

@ -9,8 +9,9 @@ rollouts on the Qwen3-4B notebook?
The short answer on a single NVIDIA B200 with Unsloth Qwen3-4B-Base, LoRA
rank 32, bf16: transformers continuous batching is functionally correct and
integrates with TRL's `use_transformers_paged=True` path, but end-to-end
throughput lands at around 7-9 percent of vLLM colocated. Full numbers and
per-step timings are in the PR description.
throughput lands at around 7 to 10 percent of vLLM colocated even after
wiring in Flash Attention 4 on Blackwell. Full numbers and per-step timings
are in the PR description.
## Files
@ -18,13 +19,57 @@ per-step timings are in the PR description.
|---|---|
| `unsloth_grpo_common.py` | Shared dataset loading, reward functions, and GRPO hyperparameters |
| `qwen3_grpo_vllm.py` | vLLM baseline training entry (`fast_inference=True`, `use_vllm=True`, `vllm_mode="colocate"`) |
| `qwen3_grpo_tpaged.py` | Continuous-batching candidate (`fast_inference=False`, `use_transformers_paged=True`, vanilla HF + PEFT) |
| `cb_vs_vllm_generation.py` | Standalone generation microbenchmark across both engines |
| `qwen3_grpo_naive.py` | Naive TRL path (vanilla HF `model.generate`, no vLLM, no CB) matching https://huggingface.co/docs/trl/grpo_trainer |
| `qwen3_grpo_tpaged.py` | Continuous-batching candidate (`fast_inference=False`, `use_transformers_paged=True`, vanilla HF + PEFT). Supports `--persistent_cb` |
| `cb_vs_vllm_generation.py` | Standalone generation microbenchmark across both engines; supports `--attn_impl`, `--persistent_cb` |
| `flash_attn_fa4_shim.py` | Installs two monkey-patches that let CB dispatch to FA4 when `--attn_impl flash_attention_2` is selected |
| `persistent_cb.py` | Replaces `model.generate_batch` with a version that reuses a single `ContinuousBatchingManager` |
## Flash Attention 4 on Blackwell (sm_100)
The CB code path has three attention implementations: `eager_paged`,
`sdpa_paged`, `flash_attention_2`. The last one requires the legacy
`flash_attn` Python package, which does not install cleanly on B200 today:
- `flash_attn==2.8.3+cu12torch2.8cxx11abiTRUE-cp313` from the Dao-AILab
releases hits `undefined symbol: _ZNK3c106SymInt6sym_neERKS0_` on torch
2.9.1 (ABI drift between torch 2.8 and 2.9).
- `flash_attn_3-3.0.0-cp39-abi3-manylinux_2_28_x86_64.whl` from the PyTorch
wheel index installs but was built for sm_80 and sm_90a only. B200 is
sm_100. The kernel call fails with "no kernel image is available for
execution on the device".
- `flash-attn-4==4.0.0b9` (pure Python CuTeDSL, Dao-AILab) works on B200.
It exposes `flash_attn.cute.flash_attn_varlen_func`.
The recipe this repo uses:
```bash
uv pip install --no-deps flash-attn-4==4.0.0b9
```
plus a tiny site-packages shim that re-exports FA4 symbols under the
FA2 `flash_attn` namespace so transformers' `is_flash_attn_2_available()` and
`_lazy_imports("flash_attention_2")` succeed. The shim lives out of tree in
`lib/python3.13/site-packages/flash_attn/__init__.py` +
`flash_attn/bert_padding.py` + a `flash_attn-2.8.3.dist-info/` directory
with enough metadata to satisfy `importlib.metadata.version("flash_attn")`.
On top of that, `flash_attn_fa4_shim.py` monkey-patches two rough edges in
the CB to FA integration that are unrelated to which FA version you use:
1. `ContinuousBatchProcessor.return_attention_mask` returns `False` for
`flash_attention_2` / `flash_attention_3` so CB does not emit a 4-D paged
attention mask that breaks `_flash_attention_forward`'s `_upad_input`
branch.
2. `_flash_attention_forward` accepts `max_seqlen_q` / `max_seqlen_k`
as aliases for `max_length_q` / `max_length_k`. Without this rename CB's
model kwargs never bind and FA is called with `max_seqlen_q=None`.
## Reproduce
```bash
pip install unsloth "transformers>=4.57" "trl>=0.25" peft vllm
uv pip install --no-deps flash-attn-4==4.0.0b9
# Generation microbenchmark (32 prompts, 512 new tokens each)
CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/cb_vs_vllm_generation.py \
@ -32,8 +77,15 @@ CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/cb_vs_vllm_generation.py \
--n_prompts 32 --n_rounds 2 --max_new_tokens 512 \
--gpu_memory_utilization 0.6
CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/cb_vs_vllm_generation.py \
--backend tpaged --stats_path logs/cb_gen.json \
# CB variants
CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/cb_vs_vllm_generation.py \
--backend tpaged --attn_impl sdpa \
--stats_path logs/cb_gen_sdpa.json \
--n_prompts 32 --n_rounds 2 --max_new_tokens 512
CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/cb_vs_vllm_generation.py \
--backend tpaged --attn_impl flash_attention_2 \
--stats_path logs/cb_gen_fa.json \
--n_prompts 32 --n_rounds 2 --max_new_tokens 512
# Full GRPO training (20 steps)
@ -42,9 +94,14 @@ CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/qwen3_grpo_vllm.py \
--output_dir outputs/grpo_vllm --stats_path logs/vllm_stats.json \
--gpu_memory_utilization 0.6
CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/qwen3_grpo_tpaged.py \
CUDA_VISIBLE_DEVICES=7 python scripts/benchmarks/qwen3_grpo_naive.py \
--max_steps 20 --num_generations 2 --per_device_train_batch_size 2 \
--output_dir outputs/grpo_tpaged --stats_path logs/tpaged_stats.json \
--output_dir outputs/grpo_naive --stats_path logs/naive_stats.json
CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/qwen3_grpo_tpaged.py \
--max_steps 20 --num_generations 2 --per_device_train_batch_size 2 \
--attn_impl flash_attention_2 \
--output_dir outputs/grpo_tpaged_fa --stats_path logs/tpaged_stats_fa.json \
--max_batch_tokens 16384 --num_blocks 16384
```
@ -72,10 +129,11 @@ how `qwen3_grpo_tpaged.py` handles them:
`unsloth.utils.attention_dispatch.run_attention` because Unsloth routes
through its own dispatcher rather than reading
`config._attn_implementation`. The benchmark script works around this by
loading a vanilla HF Qwen3 with PEFT LoRA for the tpaged path. This costs
the Unsloth training kernels but keeps the comparison clean. A proper
upstream fix is to detect `config._attn_implementation` ending in
`_paged` and delegate to the stock transformers forward.
loading a vanilla HF Qwen3 with PEFT LoRA for the tpaged and naive paths.
This costs the Unsloth training kernels but keeps the comparison clean.
A proper upstream fix is to detect `config._attn_implementation` being
`flash_attention_2` / `sdpa_paged` / `eager_paged` and delegate to the
stock transformers forward.
4. **TRL imports `GuidedDecodingParams` from `vllm.sampling_params`.**
Newer vLLM releases (>= 0.13) have moved or removed that symbol, so
@ -89,16 +147,19 @@ how `qwen3_grpo_tpaged.py` handles them:
these hooks. A vanilla HF model does not, so
`qwen3_grpo_tpaged.py` does not `import unsloth` at all.
## Why continuous batching is slower than vLLM on this workload
## Why continuous batching is still slower than vLLM on this workload
- `flash_attn` is hard to install on this box (CUDA 13.1 detected, torch
compiled against CUDA 12.8), so `paged|flash_attention_2` falls back to
`sdpa_paged`. vLLM uses FlashInfer and TRTLLM kernels.
- CB re-allocates a fresh `PagedAttentionCache` on every `generate_batch`
call. For GRPO that is once per step.
- `ContinuousBatchingManager` does not yet implement CUDA graphs
(`use_cuda_graph=True` raises `NotImplementedError`). vLLM captures 100+
mixed prefill-decode and decode graphs during warmup.
- CB re-allocates a fresh `PagedAttentionCache` on every `generate_batch`
call. For GRPO that is once per step. `--persistent_cb` (via
`persistent_cb.py`) keeps the cache warm across steps.
- FA4 is a CuTeDSL package: first call per shape pays a one-time
JIT-compile.
- vLLM uses its own colocated attention + FlashInfer / TRTLLM kernels
tuned for decode, which currently outperform everything a generic CB
path can do.
These are all upstream transformers issues, not Unsloth issues. The scripts
in this directory are intentionally simple so they are easy to port into a

117
scripts/benchmarks/cb_vs_vllm_generation.py
View file

@ -27,6 +27,13 @@ sys.path.insert(0, str(HERE))
import torch # noqa: E402
# Install the FA4 shim so transformers' continuous batching dispatches to
# Blackwell (sm_100) kernels when `--attn_impl flash_attention_2` is selected.
# No-op for the vLLM backend since vLLM doesn't go through transformers'
# attention interface.
import flash_attn_fa4_shim # noqa: E402
flash_attn_fa4_shim.apply()
def build_prompts(tokenizer, n_prompts):
from unsloth_grpo_common import (
@ -36,8 +43,8 @@ def build_prompts(tokenizer, n_prompts):
from datasets import load_dataset
apply_chat_template_to_tokenizer(tokenizer)
ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
ds = ds.shuffle(seed = 3407).select(range(n_prompts))
ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train")
ds = ds.shuffle(seed=3407).select(range(n_prompts))
messages = [
[
{"role": "system", "content": SYSTEM_PROMPT},
@ -46,11 +53,11 @@ def build_prompts(tokenizer, n_prompts):
for x in ds
]
prompts_text = [
tokenizer.apply_chat_template(m, add_generation_prompt = True, tokenize = False)
tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
for m in messages
]
prompt_ids = [
tokenizer.apply_chat_template(m, add_generation_prompt = True, tokenize = True)
tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=True)
for m in messages
]
return prompts_text, prompt_ids
@ -58,36 +65,30 @@ def build_prompts(tokenizer, n_prompts):
def run_vllm(args):
import os as _os
_os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
from unsloth import FastLanguageModel
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = args.model_name,
max_seq_length = args.max_seq_length,
load_in_4bit = False,
fast_inference = True,
max_lora_rank = 32,
gpu_memory_utilization = args.gpu_memory_utilization,
model_name=args.model_name,
max_seq_length=args.max_seq_length,
load_in_4bit=False,
fast_inference=True,
max_lora_rank=32,
gpu_memory_utilization=args.gpu_memory_utilization,
)
prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts)
from vllm import SamplingParams
sp = SamplingParams(
temperature = 1.0,
min_p = 0.1,
top_p = 1.0,
top_k = -1,
seed = 3407,
max_tokens = args.max_new_tokens,
stop = [tokenizer.eos_token],
include_stop_str_in_output = True,
temperature=1.0, min_p=0.1, top_p=1.0, top_k=-1,
seed=3407,
max_tokens=args.max_new_tokens,
stop=[tokenizer.eos_token],
include_stop_str_in_output=True,
)
# Warmup on 16 prompts then discard.
warmup_text = prompts_text[:16]
_ = model.fast_generate(warmup_text, sampling_params = sp, lora_request = None)
_ = model.fast_generate(warmup_text, sampling_params=sp, lora_request=None)
torch.cuda.synchronize()
# Three measured rounds on the full batch.
@ -97,9 +98,7 @@ def run_vllm(args):
for _ in range(args.n_rounds):
torch.cuda.synchronize()
t0 = time.perf_counter()
outputs = model.fast_generate(
prompts_text, sampling_params = sp, lora_request = None
)
outputs = model.fast_generate(prompts_text, sampling_params=sp, lora_request=None)
torch.cuda.synchronize()
wall_times.append(time.perf_counter() - t0)
decoded = sum(len(o.outputs[0].token_ids) for o in outputs)
@ -129,33 +128,37 @@ def run_tpaged(args):
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
args.model_name,
dtype = torch.bfloat16,
attn_implementation = args.attn_impl,
dtype=torch.bfloat16,
attn_implementation=args.attn_impl,
).to("cuda")
model.eval()
if args.persistent_cb:
from persistent_cb import install_for_model # noqa: WPS433
prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts)
gen_config = GenerationConfig(
max_new_tokens = args.max_new_tokens,
do_sample = True,
temperature = 1.0,
top_p = 1.0,
min_p = 0.1,
pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id,
bos_token_id = tokenizer.bos_token_id,
eos_token_id = tokenizer.eos_token_id,
use_cache = True,
max_new_tokens=args.max_new_tokens,
do_sample=True,
temperature=1.0,
top_p=1.0,
min_p=0.1,
pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
bos_token_id=tokenizer.bos_token_id,
eos_token_id=tokenizer.eos_token_id,
use_cache=True,
)
# Raise the paged-cache upper bounds; defaults (256 / 4096) throttle CB.
gen_config.max_batch_tokens = args.max_batch_tokens
gen_config.num_blocks = args.num_blocks
if args.persistent_cb:
install_for_model(model, gen_config)
# Warmup on 16 prompts.
warmup_ids = prompt_ids[:16]
with torch.inference_mode():
_ = model.generate_batch(
warmup_ids, generation_config = gen_config, progress_bar = False
)
_ = model.generate_batch(warmup_ids, generation_config=gen_config, progress_bar=False)
torch.cuda.synchronize()
n_prompt_tokens = sum(len(p) for p in prompt_ids)
@ -166,7 +169,7 @@ def run_tpaged(args):
t0 = time.perf_counter()
with torch.inference_mode():
outputs = model.generate_batch(
prompt_ids, generation_config = gen_config, progress_bar = False
prompt_ids, generation_config=gen_config, progress_bar=False
)
torch.cuda.synchronize()
wall_times.append(time.perf_counter() - t0)
@ -177,6 +180,7 @@ def run_tpaged(args):
return {
"backend": "tpaged",
"attn_impl": args.attn_impl,
"persistent_cb": args.persistent_cb,
"n_prompts": args.n_prompts,
"n_prompt_tokens": n_prompt_tokens,
"n_decoded_tokens": total_decoded,
@ -190,23 +194,25 @@ def run_tpaged(args):
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--backend", choices = ["vllm", "tpaged"], required = True)
p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
p.add_argument("--max_seq_length", type = int, default = 2048)
p.add_argument("--n_prompts", type = int, default = 64)
p.add_argument("--n_rounds", type = int, default = 3)
p.add_argument("--max_new_tokens", type = int, default = 1024)
p.add_argument("--gpu_memory_utilization", type = float, default = 0.8)
p.add_argument("--attn_impl", default = "sdpa")
p.add_argument("--max_batch_tokens", type = int, default = 8192)
p.add_argument("--num_blocks", type = int, default = 16384)
p.add_argument("--stats_path", required = True)
p.add_argument("--backend", choices=["vllm", "tpaged"], required=True)
p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base")
p.add_argument("--max_seq_length", type=int, default=2048)
p.add_argument("--n_prompts", type=int, default=64)
p.add_argument("--n_rounds", type=int, default=3)
p.add_argument("--max_new_tokens", type=int, default=1024)
p.add_argument("--gpu_memory_utilization", type=float, default=0.8)
p.add_argument("--attn_impl", default="sdpa")
p.add_argument("--max_batch_tokens", type=int, default=8192)
p.add_argument("--num_blocks", type=int, default=16384)
p.add_argument("--persistent_cb", action="store_true",
help="Reuse a single ContinuousBatchingManager across warmup + measured rounds.")
p.add_argument("--stats_path", required=True)
return p.parse_args()
def main():
args = parse_args()
os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)
os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok=True)
torch.cuda.reset_peak_memory_stats()
if args.backend == "vllm":
@ -216,8 +222,11 @@ def main():
out["peak_memory_gb"] = torch.cuda.max_memory_allocated() / 1024**3
with open(args.stats_path, "w") as f:
json.dump(out, f, indent = 2)
print(json.dumps(out, indent = 2))
json.dump(out, f, indent=2)
print(json.dumps(out, indent=2))
# When --persistent_cb is set the background CB worker thread keeps the
# process alive. Exit fast; the stats file is already flushed.
os._exit(0)
if __name__ == "__main__":

96
scripts/benchmarks/flash_attn_fa4_shim.py Normal file
View file

@ -0,0 +1,96 @@
"""Make transformers' continuous batching dispatch to Flash Attention 4 on B200.
Two integration gaps between transformers' continuous batching (CB) and the
FA2 varlen path make `attn_implementation="flash_attention_2"` fail out of the
box in transformers 4.57 even with a working FA varlen kernel:
1. CB's `ContinuousBatchProcessor` creates a 4D paged attention mask of shape
`[1, 1, q_len, k_len]` for every attention implementation except
`"paged_attention"`. `_flash_attention_forward` then enters the
`if attention_mask is not None:` branch and calls `_upad_input`, which
expects a 2D mask and fails (the FA4 kernel ultimately asserts on
`cu_seqlens_k.shape`).
2. CB passes `max_seqlen_q`/`max_seqlen_k` as model kwargs while
`_flash_attention_forward` names the parameters `max_length_q`/
`max_length_k`. The former therefore never bind and the varlen branch
inside `_flash_attention_forward` invokes the FA kernel with
`max_seqlen_q=None`.
This module patches around both, and (via the sibling
`site-packages/flash_attn/__init__.py` shim) points the FA2 varlen dispatch
at FA4's Blackwell-capable kernel. Call `apply()` once, before
`model.generate_batch` is invoked.
"""
from __future__ import annotations
import functools
_APPLIED = False
def apply() -> None:
global _APPLIED
if _APPLIED:
return
import transformers # noqa: F401 - force load
from transformers.generation.continuous_batching import continuous_api as _cb
from transformers import modeling_flash_attention_utils as _fa_utils
_patch_return_attention_mask(_cb)
_patch_flash_attention_forward(_fa_utils)
_APPLIED = True
def _patch_return_attention_mask(cb_module) -> None:
"""Don't materialise a 4D attention mask when the kernel is FA varlen.
The existing `return_attention_mask` only skips the mask for
`"paged_attention"`. We extend the skip set to `"flash_attention_2"`
(and `"flash_attention_3"` for future-proofing) because the varlen path
relies on `cu_seq_lens_*` and reads no mask.
"""
_SKIP_MASK_IMPLS = {
"paged_attention",
"flash_attention_2",
"flash_attention_3",
}
def return_attention_mask(self) -> bool:
return self.config._attn_implementation not in _SKIP_MASK_IMPLS
cb_module.ContinuousBatchProcessor.return_attention_mask = return_attention_mask
def _patch_flash_attention_forward(fa_utils_module) -> None:
"""Accept CB's `max_seqlen_q`/`max_seqlen_k` kwargs as aliases.
transformers names the parameters `max_length_q`/`max_length_k`, but CB
(and most downstream call sites) name them `max_seqlen_q`/
`max_seqlen_k`. We rename at the boundary so callers on either side work
unchanged.
"""
original = fa_utils_module._flash_attention_forward
@functools.wraps(original)
def wrapper(*args, **kwargs):
if kwargs.get("max_length_q") is None and "max_seqlen_q" in kwargs:
kwargs["max_length_q"] = kwargs.pop("max_seqlen_q")
if kwargs.get("max_length_k") is None and "max_seqlen_k" in kwargs:
kwargs["max_length_k"] = kwargs.pop("max_seqlen_k")
return original(*args, **kwargs)
fa_utils_module._flash_attention_forward = wrapper
# Some integration modules imported the function by name before we
# patched. Re-bind the most common consumers so they pick up the wrapper.
try:
from transformers.integrations import flash_attention as _flash_integration
_flash_integration._flash_attention_forward = wrapper
except Exception:
pass

115
scripts/benchmarks/persistent_cb.py Normal file
View file

@ -0,0 +1,115 @@
"""Persistent ContinuousBatching manager for multi-step generation.
`model.generate_batch(...)` initializes a fresh `ContinuousBatchingManager`
on every call, which in turn allocates a new `PagedAttentionCache` and
starts a new worker thread. Inside a GRPO training loop this happens once
per step, amortized across `num_generations * per_device_train_batch_size`
prompts per step, so the constant per-step cost (cache alloc, prefill
warmup, thread spin-up) dominates when batch sizes are modest.
`install_for_model(model, generation_config)` monkey-patches
`model.generate_batch` on this instance to reuse a single long-lived
manager. The manager is started lazily on first call; `teardown(model)`
stops the background thread.
This is deliberately kept as a stand-alone helper so it can be enabled /
disabled per-run via CLI flag without touching TRL or transformers
installs.
"""
from __future__ import annotations
import threading
import types
from typing import Optional
import torch
from transformers import GenerationConfig
from transformers.generation.continuous_batching import RequestStatus
_ATTR = "_persistent_cb_manager"
_LOCK_ATTR = "_persistent_cb_lock"
def install_for_model(model: torch.nn.Module, generation_config: GenerationConfig) -> None:
"""Replace `model.generate_batch` with a version that reuses one manager.
The replacement accepts the same arguments as the stock method. A
trailing `generation_config` supplied to the call takes precedence; if
it differs from the one used at init, the persistent manager is torn
down and rebuilt (rare, but keeps semantics intact).
"""
setattr(model, _LOCK_ATTR, threading.Lock())
setattr(model, _ATTR, None)
setattr(model, "_persistent_cb_gen_config", generation_config)
original = model.generate_batch
def generate_batch(
self,
inputs,
generation_config: Optional[GenerationConfig] = None,
progress_bar: bool = False,
slice_inputs: bool = True,
**kwargs,
):
if not inputs:
return {}
gen_config = generation_config or getattr(self, "_persistent_cb_gen_config", None) or self.generation_config
lock = getattr(self, _LOCK_ATTR)
with lock:
manager = getattr(self, _ATTR)
stale = False
if manager is not None:
stale = (
getattr(manager, "generation_config", None) is not gen_config
or not manager.is_running()
)
if stale:
try:
manager.stop(block=True, timeout=5.0)
except Exception:
pass
setattr(self, _ATTR, None)
manager = None
if manager is None:
manager = self.init_continuous_batching(
generation_config=gen_config,
slice_inputs=slice_inputs,
)
manager.start()
setattr(self, _ATTR, manager)
results = {}
num_requests = len(inputs)
manager.add_requests(inputs, **kwargs)
finished = 0
while finished < num_requests:
result = manager.get_result(timeout=1)
if result is None:
if not manager.is_running():
break
continue
if result.status == RequestStatus.FINISHED:
results[result.request_id] = result
finished += 1
else:
continue
return results
model.generate_batch = types.MethodType(generate_batch, model)
setattr(model, "_persistent_cb_original_generate_batch", original)
def teardown(model: torch.nn.Module) -> None:
manager = getattr(model, _ATTR, None)
if manager is not None:
try:
manager.stop(block=True, timeout=5.0)
except Exception:
pass
if hasattr(model, "_persistent_cb_original_generate_batch"):
model.generate_batch = model._persistent_cb_original_generate_batch

196
scripts/benchmarks/qwen3_grpo_naive.py Normal file
View file

@ -0,0 +1,196 @@
"""Qwen3-4B GRPO with naive TRL generation (no vLLM, no CB).
Mirrors the TRL example at https://huggingface.co/docs/trl/grpo_trainer: a
vanilla HF model + `GRPOTrainer` with the default rollout path, which calls
`model.generate(...)` per step. This is the honest "baseline" baseline — it
is what a user would get if they just followed the TRL docs without enabling
vLLM colocate or transformers continuous batching. Useful as a third column
in the benchmark table.
Hyperparameters, dataset, and reward functions are shared with the vLLM and
CB scripts via `unsloth_grpo_common.py` so numbers are apples-to-apples.
Run:
CUDA_VISIBLE_DEVICES=2 python scripts/qwen3_grpo_naive.py \
--max_steps 20 --num_generations 2 --per_device_train_batch_size 2 \
--output_dir outputs/grpo_naive --stats_path logs/naive_stats.json
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
from pathlib import Path
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
# Same vLLM sampling-params shim as the tpaged script so TRL imports cleanly
# even when vLLM is installed but the GuidedDecodingParams symbol has moved.
try:
import vllm.sampling_params as _vllm_sp
if not hasattr(_vllm_sp, "GuidedDecodingParams"):
class _GuidedDecodingParamsShim: # pragma: no cover
def __init__(self, *a, **kw):
pass
_vllm_sp.GuidedDecodingParams = _GuidedDecodingParamsShim
except ImportError:
pass
import torch # noqa: E402
from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402
from peft import LoraConfig, get_peft_model # noqa: E402
from unsloth_grpo_common import ( # noqa: E402
apply_chat_template_to_tokenizer,
build_dataset,
build_reward_funcs,
build_grpo_kwargs,
)
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base")
p.add_argument("--max_seq_length", type=int, default=2048)
p.add_argument("--lora_rank", type=int, default=32)
p.add_argument("--max_steps", type=int, default=20)
p.add_argument("--num_generations", type=int, default=2)
p.add_argument("--per_device_train_batch_size", type=int, default=2)
p.add_argument("--gradient_accumulation_steps", type=int, default=1)
p.add_argument("--attn_impl", default="sdpa",
help="Attention implementation: sdpa or flash_attention_2 (FA4 shim installed).")
p.add_argument("--output_dir", default="outputs/grpo_naive")
p.add_argument("--stats_path", default="logs/naive_stats.json")
return p.parse_args()
def main():
args = parse_args()
os.makedirs(args.output_dir, exist_ok=True)
os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok=True)
# Install the FA4 shim only if the caller asked for flash_attention_2.
# For sdpa we leave transformers untouched.
if args.attn_impl == "flash_attention_2":
import flash_attn_fa4_shim
flash_attn_fa4_shim.apply()
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
args.model_name,
dtype=torch.bfloat16,
attn_implementation=args.attn_impl,
).to("cuda")
lora = LoraConfig(
r=args.lora_rank,
lora_alpha=args.lora_rank * 2,
target_modules=[
"q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj",
],
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora)
try:
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
except TypeError:
model.gradient_checkpointing_enable()
model.enable_input_require_grads()
apply_chat_template_to_tokenizer(tokenizer)
dataset, maximum_length = build_dataset(tokenizer, max_seq_length=args.max_seq_length)
print(f"[naive] Max prompt length (p90): {maximum_length}")
reward_funcs = build_reward_funcs(tokenizer)
from trl import GRPOConfig, GRPOTrainer
shared = build_grpo_kwargs(
tokenizer,
maximum_length,
max_seq_length=args.max_seq_length,
max_steps=args.max_steps,
num_generations=args.num_generations,
per_device_train_batch_size=args.per_device_train_batch_size,
gradient_accumulation_steps=args.gradient_accumulation_steps,
output_dir=args.output_dir,
)
# transformers' TopKLogitsWarper rejects -1. None skips the warper.
shared["top_k"] = None
training_args = GRPOConfig(
use_vllm=False,
use_transformers_paged=False,
bf16=True,
**shared,
)
from transformers import TrainerCallback
timings = {"step_wall": [], "loss": [], "reward": []}
class StepTimer(TrainerCallback):
def __init__(self):
self.t0 = None
def on_step_begin(self, _args, state, control, **kwargs):
torch.cuda.synchronize()
self.t0 = time.perf_counter()
def on_log(self, _args, state, control, logs=None, **kwargs):
if logs is None:
return
if "loss" in logs:
timings["loss"].append(float(logs["loss"]))
if "reward" in logs:
timings["reward"].append(float(logs["reward"]))
def on_step_end(self, _args, state, control, **kwargs):
if self.t0 is not None:
torch.cuda.synchronize()
timings["step_wall"].append(time.perf_counter() - self.t0)
trainer = GRPOTrainer(
model=model,
processing_class=tokenizer,
reward_funcs=reward_funcs,
args=training_args,
train_dataset=dataset,
callbacks=[StepTimer()],
)
torch.cuda.reset_peak_memory_stats()
t_start = time.perf_counter()
trainer.train()
t_train = time.perf_counter() - t_start
peak = torch.cuda.max_memory_allocated() / 1024**3
stats = {
"backend": "naive_trl",
"attn_impl": args.attn_impl,
"train_wall_s": t_train,
"peak_memory_gb": peak,
"step_wall_s": timings["step_wall"],
"losses": timings["loss"],
"rewards": timings["reward"],
"max_prompt_length": shared["max_prompt_length"],
"max_completion_length": shared["max_completion_length"],
"num_generations": args.num_generations,
"max_steps": args.max_steps,
}
with open(args.stats_path, "w") as f:
json.dump(stats, f, indent=2)
print(f"[naive] Wrote stats to {args.stats_path}")
print(f"[naive] Total train wall: {t_train:.1f}s Peak mem: {peak:.2f} GB")
if __name__ == "__main__":
main()

141
scripts/benchmarks/qwen3_grpo_tpaged.py
View file

@ -31,13 +31,10 @@ sys.path.insert(0, str(HERE))
# `for_inference()` hooks, which a vanilla HF model does not.
try:
import vllm.sampling_params as _vllm_sp
if not hasattr(_vllm_sp, "GuidedDecodingParams"):
class _GuidedDecodingParamsShim: # pragma: no cover - used only if TRL asks
def __init__(self, *a, **kw):
pass
_vllm_sp.GuidedDecodingParams = _GuidedDecodingParamsShim
except ImportError:
pass
@ -46,6 +43,9 @@ import torch # noqa: E402
from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402
from peft import LoraConfig, get_peft_model # noqa: E402
import flash_attn_fa4_shim # noqa: E402
flash_attn_fa4_shim.apply()
from unsloth_grpo_common import ( # noqa: E402
apply_chat_template_to_tokenizer,
build_dataset,
@ -56,39 +56,31 @@ from unsloth_grpo_common import ( # noqa: E402
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
p.add_argument("--max_seq_length", type = int, default = 2048)
p.add_argument("--lora_rank", type = int, default = 32)
p.add_argument("--max_steps", type = int, default = 61)
p.add_argument("--num_generations", type = int, default = 4)
p.add_argument("--per_device_train_batch_size", type = int, default = 1)
p.add_argument("--gradient_accumulation_steps", type = int, default = 1)
p.add_argument(
"--attn_impl",
default = "sdpa",
help = "Base attention impl to compose with paged. 'sdpa' or 'flash_attention_2'.",
)
p.add_argument(
"--max_batch_tokens",
type = int,
default = 8192,
help = "PagedAttentionCache.max_batch_tokens. Default upper bound is 256 which is far too small.",
)
p.add_argument(
"--num_blocks",
type = int,
default = 8192,
help = "PagedAttentionCache.num_blocks (block_size=32). 8192*32 tokens of KV capacity.",
)
p.add_argument("--output_dir", default = "outputs/grpo_tpaged")
p.add_argument("--stats_path", default = "logs/tpaged_stats.json")
p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base")
p.add_argument("--max_seq_length", type=int, default=2048)
p.add_argument("--lora_rank", type=int, default=32)
p.add_argument("--max_steps", type=int, default=61)
p.add_argument("--num_generations", type=int, default=4)
p.add_argument("--per_device_train_batch_size", type=int, default=1)
p.add_argument("--gradient_accumulation_steps", type=int, default=1)
p.add_argument("--attn_impl", default="sdpa",
help="Base attention impl to compose with paged. 'sdpa' or 'flash_attention_2'.")
p.add_argument("--max_batch_tokens", type=int, default=8192,
help="PagedAttentionCache.max_batch_tokens. Default upper bound is 256 which is far too small.")
p.add_argument("--num_blocks", type=int, default=8192,
help="PagedAttentionCache.num_blocks (block_size=32). 8192*32 tokens of KV capacity.")
p.add_argument("--output_dir", default="outputs/grpo_tpaged")
p.add_argument("--stats_path", default="logs/tpaged_stats.json")
p.add_argument("--persistent_cb", action="store_true",
help="Reuse one ContinuousBatchingManager across every training step instead "
"of letting TRL's generate_batch rebuild it (and the paged cache) each step.")
return p.parse_args()
def main():
args = parse_args()
os.makedirs(args.output_dir, exist_ok = True)
os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)
os.makedirs(args.output_dir, exist_ok=True)
os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok=True)
# 1. Vanilla HF load (no Unsloth patches on the attention forward).
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
@ -96,31 +88,24 @@ def main():
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
args.model_name,
dtype = torch.bfloat16,
attn_implementation = args.attn_impl,
dtype=torch.bfloat16,
attn_implementation=args.attn_impl,
)
model.to("cuda")
lora = LoraConfig(
r = args.lora_rank,
lora_alpha = args.lora_rank * 2,
target_modules = [
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
r=args.lora_rank,
lora_alpha=args.lora_rank * 2,
target_modules=[
"q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj",
],
bias = "none",
task_type = "CAUSAL_LM",
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora)
try:
model.gradient_checkpointing_enable(
gradient_checkpointing_kwargs = {"use_reentrant": False}
)
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
except TypeError:
model.gradient_checkpointing_enable()
model.enable_input_require_grads()
@ -128,24 +113,21 @@ def main():
apply_chat_template_to_tokenizer(tokenizer)
# 2. Dataset + rewards (identical to the vLLM script).
dataset, maximum_length = build_dataset(
tokenizer, max_seq_length = args.max_seq_length
)
dataset, maximum_length = build_dataset(tokenizer, max_seq_length=args.max_seq_length)
print(f"[tpaged] Max prompt length (p90): {maximum_length}")
reward_funcs = build_reward_funcs(tokenizer)
# 3. Build GRPOConfig with transformers continuous batching enabled.
from trl import GRPOConfig, GRPOTrainer
shared = build_grpo_kwargs(
tokenizer,
maximum_length,
max_seq_length = args.max_seq_length,
max_steps = args.max_steps,
num_generations = args.num_generations,
per_device_train_batch_size = args.per_device_train_batch_size,
gradient_accumulation_steps = args.gradient_accumulation_steps,
output_dir = args.output_dir,
max_seq_length=args.max_seq_length,
max_steps=args.max_steps,
num_generations=args.num_generations,
per_device_train_batch_size=args.per_device_train_batch_size,
gradient_accumulation_steps=args.gradient_accumulation_steps,
output_dir=args.output_dir,
)
# transformers `TopKLogitsWarper` rejects -1. None skips the warper entirely.
shared["top_k"] = None
@ -155,10 +137,10 @@ def main():
# `generation_kwargs`, which TRL forwards to `GenerationConfig`, which the
# CB manager then reads off when sizing the paged cache.
training_args = GRPOConfig(
use_vllm = False,
use_transformers_paged = True,
bf16 = True,
generation_kwargs = {
use_vllm=False,
use_transformers_paged=True,
bf16=True,
generation_kwargs={
"max_batch_tokens": args.max_batch_tokens,
"num_blocks": args.num_blocks,
},
@ -178,7 +160,7 @@ def main():
torch.cuda.synchronize()
self.t0 = time.perf_counter()
def on_log(self, _args, state, control, logs = None, **kwargs):
def on_log(self, _args, state, control, logs=None, **kwargs):
if logs is None:
return
if "loss" in logs:
@ -192,17 +174,33 @@ def main():
timings["step_wall"].append(time.perf_counter() - self.t0)
trainer = GRPOTrainer(
model = model,
processing_class = tokenizer,
reward_funcs = reward_funcs,
args = training_args,
train_dataset = dataset,
callbacks = [StepTimer()],
model=model,
processing_class=tokenizer,
reward_funcs=reward_funcs,
args=training_args,
train_dataset=dataset,
callbacks=[StepTimer()],
)
if args.persistent_cb:
# TRL constructs `self.generation_config` once in `__init__`; reuse
# the same object so the persistent manager stays warm.
from persistent_cb import install_for_model, teardown
# TRL generates against the unwrapped base model; attach the patch
# directly to it so every rollout picks up the persistent manager.
base = trainer.model_wrapped.base_model.model if hasattr(trainer.model_wrapped, "base_model") else trainer.model_wrapped
install_for_model(base, trainer.generation_config)
# PEFT's wrapper chains `generate_batch` through `base_model.model` via
# __getattr__, so installing on `base` is enough for TRL's call path.
torch.cuda.reset_peak_memory_stats()
t_start = time.perf_counter()
trainer.train()
try:
trainer.train()
finally:
if args.persistent_cb:
from persistent_cb import teardown
teardown(base)
t_train = time.perf_counter() - t_start
peak = torch.cuda.max_memory_allocated() / 1024**3
@ -219,9 +217,10 @@ def main():
"max_completion_length": shared["max_completion_length"],
"num_generations": args.num_generations,
"max_steps": args.max_steps,
"persistent_cb": args.persistent_cb,
}
with open(args.stats_path, "w") as f:
json.dump(stats, f, indent = 2)
json.dump(stats, f, indent=2)
print(f"[tpaged] Wrote stats to {args.stats_path}")
print(f"[tpaged] Total train wall: {t_train:.1f}s Peak mem: {peak:.2f} GB")