8dfa076ee4
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.
115 lines
4 KiB
Python
115 lines
4 KiB
Python
"""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
|