diff --git a/scripts/benchmarks/README.md b/scripts/benchmarks/README.md index 68ad94f30c..dfd60ddac6 100644 --- a/scripts/benchmarks/README.md +++ b/scripts/benchmarks/README.md @@ -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 diff --git a/scripts/benchmarks/cb_vs_vllm_generation.py b/scripts/benchmarks/cb_vs_vllm_generation.py index b0b6e26966..8a4f5b6cc7 100644 --- a/scripts/benchmarks/cb_vs_vllm_generation.py +++ b/scripts/benchmarks/cb_vs_vllm_generation.py @@ -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__": diff --git a/scripts/benchmarks/flash_attn_fa4_shim.py b/scripts/benchmarks/flash_attn_fa4_shim.py new file mode 100644 index 0000000000..58b7f8026b --- /dev/null +++ b/scripts/benchmarks/flash_attn_fa4_shim.py @@ -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 diff --git a/scripts/benchmarks/persistent_cb.py b/scripts/benchmarks/persistent_cb.py new file mode 100644 index 0000000000..9787ba50c5 --- /dev/null +++ b/scripts/benchmarks/persistent_cb.py @@ -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 diff --git a/scripts/benchmarks/qwen3_grpo_naive.py b/scripts/benchmarks/qwen3_grpo_naive.py new file mode 100644 index 0000000000..2a03083314 --- /dev/null +++ b/scripts/benchmarks/qwen3_grpo_naive.py @@ -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() diff --git a/scripts/benchmarks/qwen3_grpo_tpaged.py b/scripts/benchmarks/qwen3_grpo_tpaged.py index ca128d2401..321c8b1aed 100644 --- a/scripts/benchmarks/qwen3_grpo_tpaged.py +++ b/scripts/benchmarks/qwen3_grpo_tpaged.py @@ -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")