Phase 2 vibe (10-step): vllm vs unsloth_fi_false vs cb_paged + Phase 3 fixes
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.
This commit is contained in:
Daniel Han
parent
0557f9151c
commit
5f3e1c98df
13 changed files with 1815 additions and 569 deletions
|
|
@ -55,7 +55,6 @@ from transformers.generation.continuous_batching.scheduler import FIFOScheduler
|
|||
@dataclass
|
||||
class CBSyncConfig:
|
||||
"""Tunables for the sync driver."""
|
||||
|
||||
max_new_tokens: int = 512
|
||||
use_cuda_graph: bool = True
|
||||
# Number of eager warmup steps before capturing a CUDA graph.
|
||||
|
|
@ -68,7 +67,7 @@ class CBSyncConfig:
|
|||
max_batch_tokens: int = 8192
|
||||
num_blocks: int = 8192
|
||||
# Progress callback (step_index, tokens_produced_total) -> None.
|
||||
on_step: Optional[callable] = field(default = None)
|
||||
on_step: Optional[callable] = field(default=None)
|
||||
|
||||
|
||||
class SyncCBDriver:
|
||||
|
|
@ -80,12 +79,8 @@ class SyncCBDriver:
|
|||
finished.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: torch.nn.Module,
|
||||
generation_config: GenerationConfig,
|
||||
cfg: CBSyncConfig,
|
||||
):
|
||||
def __init__(self, model: torch.nn.Module, generation_config: GenerationConfig,
|
||||
cfg: CBSyncConfig):
|
||||
self.model = model.eval()
|
||||
self.cfg = cfg
|
||||
# Force-greedy + upper-bound overrides on a copy.
|
||||
|
|
@ -105,11 +100,11 @@ class SyncCBDriver:
|
|||
# We reuse the Manager's methods but never call `.start()`. Its
|
||||
# constructor builds: logit processor, do_sample flag, etc.
|
||||
self.manager = ContinuousBatchingManager(
|
||||
model = self.model,
|
||||
generation_config = gc,
|
||||
manual_eviction = False,
|
||||
streaming = False,
|
||||
slice_inputs = False, # fixed-shape views -> CUDA-graph safe
|
||||
model=self.model,
|
||||
generation_config=gc,
|
||||
manual_eviction=False,
|
||||
streaming=False,
|
||||
slice_inputs=False, # fixed-shape views -> CUDA-graph safe
|
||||
)
|
||||
# The manager's `use_cuda_graph` is checked inside `warmup()`, but its
|
||||
# `__init__` refuses to set it. Set it directly now that we bypass
|
||||
|
|
@ -123,7 +118,7 @@ class SyncCBDriver:
|
|||
gc,
|
||||
self.model.device,
|
||||
self.model.dtype,
|
||||
tp_size = getattr(self.model, "_tp_size", None),
|
||||
tp_size=getattr(self.model, "_tp_size", None),
|
||||
)
|
||||
self.batch_processor = ContinuousBatchProcessor(
|
||||
self.cache,
|
||||
|
|
@ -134,10 +129,10 @@ class SyncCBDriver:
|
|||
self.manager.stop_event,
|
||||
self.model.device,
|
||||
self.model.dtype,
|
||||
FIFOScheduler(self.cache, manual_eviction = False),
|
||||
streaming = False,
|
||||
manual_eviction = False,
|
||||
slice_inputs = False,
|
||||
FIFOScheduler(self.cache),
|
||||
streaming=False,
|
||||
manual_eviction=False,
|
||||
slice_inputs=False,
|
||||
)
|
||||
self.manager.batch_processor = self.batch_processor
|
||||
self._graph: Optional[torch.cuda.CUDAGraph] = None
|
||||
|
|
@ -153,13 +148,13 @@ class SyncCBDriver:
|
|||
for _ in range(self.cfg.warmup_steps):
|
||||
self.manager._generation_step(self.batch_processor)
|
||||
torch.cuda.synchronize()
|
||||
stream = torch.cuda.Stream(device = self.model.device)
|
||||
stream = torch.cuda.Stream(device=self.model.device)
|
||||
stream.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(stream):
|
||||
self.manager._generation_step(self.batch_processor)
|
||||
torch.cuda.current_stream().wait_stream(stream)
|
||||
self._graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(self._graph, stream = stream):
|
||||
with torch.cuda.graph(self._graph, stream=stream):
|
||||
self.manager._generation_step(self.batch_processor)
|
||||
else:
|
||||
self._graph.replay()
|
||||
|
|
@ -168,12 +163,24 @@ class SyncCBDriver:
|
|||
"""Run the decode loop until every request finishes. Returns a dict
|
||||
{request_id: generated_token_ids}."""
|
||||
results: dict[str, list[int]] = {}
|
||||
while self.batch_processor.has_pending_requests():
|
||||
# prepare_next_batch drains self.input_queue into the scheduler; we
|
||||
# have to call it at least once before has_pending_requests() can
|
||||
# return True. Loop until both the input_queue is empty AND the
|
||||
# scheduler has nothing queued/active.
|
||||
while True:
|
||||
input_empty = self.manager.input_queue.empty()
|
||||
nothing_scheduled = not self.batch_processor.has_pending_requests()
|
||||
if input_empty and nothing_scheduled:
|
||||
break
|
||||
# 1. CPU: schedule the next batch (prepare_next_batch reads the
|
||||
# input_queue, packs shapes).
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
if not self.batch_processor.prepare_next_batch():
|
||||
# prepare_next_batch returns False if both the input queue
|
||||
# drained empty AND the scheduler has no active requests. If
|
||||
# we reach here with items still in input_queue, something is
|
||||
# wrong -- bail to avoid an infinite loop.
|
||||
break
|
||||
# 2. GPU: forward (graphed on decode steps, eager on prefill).
|
||||
if self.cfg.use_cuda_graph and self._is_pure_decode():
|
||||
|
|
@ -211,20 +218,14 @@ class SyncCBDriver:
|
|||
Shape consistency between decodes is what makes the graph replayable.
|
||||
"""
|
||||
try:
|
||||
return (
|
||||
self.batch_processor.total_query_length
|
||||
== self.batch_processor.total_batch_size
|
||||
)
|
||||
return (self.batch_processor.total_query_length
|
||||
== self.batch_processor.total_batch_size)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _produced(self) -> int:
|
||||
return sum(
|
||||
len(r.generated_tokens)
|
||||
for r in getattr(
|
||||
self.batch_processor.scheduler, "active_requests", {}
|
||||
).values()
|
||||
)
|
||||
return sum(len(r.generated_tokens) for r
|
||||
in getattr(self.batch_processor.scheduler, "active_requests", {}).values())
|
||||
|
||||
def close(self):
|
||||
# Caches hold GPU memory; free them explicitly.
|
||||
|
|
@ -234,12 +235,9 @@ class SyncCBDriver:
|
|||
self.manager.batch_processor = None
|
||||
|
||||
|
||||
def cb_sync_generate(
|
||||
model: torch.nn.Module,
|
||||
generation_config: GenerationConfig,
|
||||
prompt_ids_list: list[list[int]],
|
||||
cfg: CBSyncConfig,
|
||||
) -> dict[str, list[int]]:
|
||||
def cb_sync_generate(model: torch.nn.Module, generation_config: GenerationConfig,
|
||||
prompt_ids_list: list[list[int]],
|
||||
cfg: CBSyncConfig) -> dict[str, list[int]]:
|
||||
"""One-shot entrypoint: build a driver, submit, drain, close.
|
||||
|
||||
Matches the semantics of `model.generate_batch(...)` but on the main
|
||||
|
|
@ -265,18 +263,17 @@ if __name__ == "__main__":
|
|||
sys.path.insert(0, str(HERE))
|
||||
|
||||
import flash_attn_fa4_shim # noqa: E402
|
||||
|
||||
flash_attn_fa4_shim.apply()
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
|
||||
parser.add_argument("--n_prompts", type = int, default = 32)
|
||||
parser.add_argument("--max_new_tokens", type = int, default = 512)
|
||||
parser.add_argument("--attn_impl", default = "paged_attention")
|
||||
parser.add_argument("--use_cuda_graph", action = "store_true")
|
||||
parser.add_argument("--max_batch_tokens", type = int, default = 8192)
|
||||
parser.add_argument("--num_blocks", type = int, default = 8192)
|
||||
parser.add_argument("--stats_path", required = True)
|
||||
parser.add_argument("--model_name", default="unsloth/Qwen3-4B-Base")
|
||||
parser.add_argument("--n_prompts", type=int, default=32)
|
||||
parser.add_argument("--max_new_tokens", type=int, default=512)
|
||||
parser.add_argument("--attn_impl", default="paged_attention")
|
||||
parser.add_argument("--use_cuda_graph", action="store_true")
|
||||
parser.add_argument("--max_batch_tokens", type=int, default=8192)
|
||||
parser.add_argument("--num_blocks", type=int, default=8192)
|
||||
parser.add_argument("--stats_path", required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
|
||||
|
|
@ -285,49 +282,36 @@ if __name__ == "__main__":
|
|||
if tok.pad_token is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model_name,
|
||||
dtype = torch.bfloat16,
|
||||
attn_implementation = args.attn_impl,
|
||||
args.model_name, dtype=torch.bfloat16,
|
||||
attn_implementation=args.attn_impl,
|
||||
).to("cuda")
|
||||
model.eval()
|
||||
|
||||
from unsloth_grpo_common import (
|
||||
SYSTEM_PROMPT,
|
||||
apply_chat_template_to_tokenizer,
|
||||
SYSTEM_PROMPT, apply_chat_template_to_tokenizer,
|
||||
)
|
||||
from datasets import load_dataset
|
||||
|
||||
apply_chat_template_to_tokenizer(tok)
|
||||
ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
|
||||
ds = ds.shuffle(seed = 3407).select(range(args.n_prompts))
|
||||
messages = [
|
||||
[
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": x["prompt"]},
|
||||
]
|
||||
for x in ds
|
||||
]
|
||||
prompt_ids = [
|
||||
tok.apply_chat_template(m, add_generation_prompt = True, tokenize = True)
|
||||
for m in messages
|
||||
]
|
||||
ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train")
|
||||
ds = ds.shuffle(seed=3407).select(range(args.n_prompts))
|
||||
messages = [[{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": x["prompt"]}] for x in ds]
|
||||
prompt_ids = [tok.apply_chat_template(m, add_generation_prompt=True, tokenize=True)
|
||||
for m in messages]
|
||||
|
||||
gc = GenerationConfig(
|
||||
max_new_tokens = args.max_new_tokens,
|
||||
do_sample = False,
|
||||
pad_token_id = tok.pad_token_id,
|
||||
bos_token_id = tok.bos_token_id,
|
||||
eos_token_id = tok.eos_token_id,
|
||||
use_cache = True,
|
||||
max_new_tokens=args.max_new_tokens, do_sample=False,
|
||||
pad_token_id=tok.pad_token_id, bos_token_id=tok.bos_token_id,
|
||||
eos_token_id=tok.eos_token_id, use_cache=True,
|
||||
)
|
||||
|
||||
cfg = CBSyncConfig(
|
||||
max_new_tokens = args.max_new_tokens,
|
||||
use_cuda_graph = args.use_cuda_graph,
|
||||
max_batch_tokens = args.max_batch_tokens,
|
||||
num_blocks = args.num_blocks,
|
||||
eos_token_id = tok.eos_token_id,
|
||||
pad_token_id = tok.pad_token_id or tok.eos_token_id,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
use_cuda_graph=args.use_cuda_graph,
|
||||
max_batch_tokens=args.max_batch_tokens,
|
||||
num_blocks=args.num_blocks,
|
||||
eos_token_id=tok.eos_token_id,
|
||||
pad_token_id=tok.pad_token_id or tok.eos_token_id,
|
||||
)
|
||||
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
|
|
@ -358,7 +342,7 @@ if __name__ == "__main__":
|
|||
"max_new_tokens": args.max_new_tokens,
|
||||
"peak_memory_gb": torch.cuda.max_memory_allocated() / 1024**3,
|
||||
}
|
||||
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)
|
||||
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))
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue