Phase 1+3: LoRA rollout benchmarks + CB sync driver + Phase 4 scaffold
Phase 1 results (`scripts/benchmarks/results/lora_rollout_baselines.md`):
- vLLM+LoRA: 4581 decode tok/s, 156 GB peak (gold, 100%)
- unsloth_fi_false+LoRA: 641 tok/s, 15.8 GB peak (14%, 7x lower mem)
- CB paged+FA4 persistent+LoRA: 422 tok/s (9.2%)
- CB sdpa_paged persistent+LoRA: 434 tok/s (9.5%)
Headline finding: Unsloth's `fast_inference=False` path (custom HF inference
kernels with cached fp16 LoRA in `fast_linear_forward`) is 1.5x faster than
CB at 1/7th the peak memory. Phase 2 will include it as a first-class backend.
cb_vs_vllm_generation.py:
- New --lora_adapter flag. vLLM uses LoRARequest; tpaged uses
PeftModel.from_pretrained (no merge_adapter so we measure LoRA-active
inference); unsloth_fi_false copies the adapter weights into Unsloth's
get_peft_model wrapper (with key normalization so PEFT's base_model.model.
prefix and Unsloth's .default. wrapper both match).
- New unsloth_fi_false backend: batched generate across all 32 prompts in
a single call after FastLanguageModel.for_inference(model).
- Exposed sampling knobs (temp/top_p/min_p/top_k); defaults are the
equivalence params.
cb_sync_driver.py (Phase 3 scaffold):
- SyncCBDriver owns PagedAttentionCache + ContinuousBatchProcessor +
FIFOScheduler on the main thread. Never calls manager.start() so there is
no background thread.
- slice_inputs=False => fixed-shape buffer views each step => CUDA graph
replay is safe.
- use_cuda_graph=True path: 2-step eager warmup, then capture one decode
step, then replay. `_is_pure_decode()` keeps prefill out of the graphed
path since those have varying shapes.
- Greedy sampling only (CUDA-graph-safe); stochastic sanity checks stay in
the non-graphed path.
- Standalone benchmark harness at the bottom.
qwen3_grpo_unified.py (Phase 4 scaffold):
- Single entrypoint for vllm / unsloth_fi_false / cb_paged / cb_sdpa /
naive_trl backends sharing dataset, reward funcs, sampling, and the
torch_debugging_utils StatisticsCallback.
- New --compile_mode {default,reduce-overhead,max-autotune-no-cudagraphs}
that compiles `trainer.model.forward` and `trainer.ref_model.forward`
after the trainer is built. CompileDebugger tracks graph breaks and
recompiles. Skipped for vLLM since vLLM owns its own compile pipeline.
- Post-warmup median (skip first 3 steps) is computed and saved alongside
the full per-step logs.
make_lora_adapter.py: writes a canonical PEFT adapter to
outputs/lora_rank32_fresh. Re-initializes lora_B with a tiny gaussian so
the adapter isn't a no-op (PEFT's default zero-init would let LoRA kernels
short-circuit).
qwen3_grpo_notebook.py (Phase 0): notebook-to-script port with
StatisticsCallback and equivalence sampling. 10-step reference reported in
scripts/benchmarks/results/notebook_ref_10.md (median step 5.80s,
peak 158.9 GB).
This commit is contained in:
Daniel Han
parent
c31533fc02
commit
5907d1525c
10 changed files with 886 additions and 486 deletions
336
scripts/benchmarks/cb_sync_driver.py
Normal file
336
scripts/benchmarks/cb_sync_driver.py
Normal file
|
|
@ -0,0 +1,336 @@
|
|||
"""Main-thread synchronous driver for `ContinuousBatchProcessor`.
|
||||
|
||||
`ContinuousBatchingManager.start()` spawns a background thread that owns the
|
||||
decode loop. That thread conflicts with two things we want to enable here:
|
||||
|
||||
1. `torch.compile(mode="reduce-overhead")` which uses `cudagraph_trees` and
|
||||
requires main-thread TLS.
|
||||
2. Raw `torch.cuda.CUDAGraph` capture/replay on the decode forward, which is
|
||||
the hot path. Captures *can* live in a child thread in principle, but
|
||||
integrating with Inductor and debugging goes much smoother on the main
|
||||
thread.
|
||||
|
||||
The manager's dead `warmup()` path suggests CB was supposed to grow CUDA
|
||||
graph support upstream, but `init_continuous_batching` currently raises
|
||||
`NotImplementedError` on `use_cuda_graph=True`. This driver side-steps that
|
||||
entirely by not going through `manager.start()` at all.
|
||||
|
||||
Key fixed-shape invariant: with `slice_inputs=False`, the full pre-allocated
|
||||
tensor buffers (input_ids, position_ids, cu_seq_lens_*, attention_mask,
|
||||
read_index / write_index) are returned as *views of the same storage* every
|
||||
step, so their shapes are constant across iterations. That is the precondition
|
||||
for CUDA graph replay to be safe.
|
||||
|
||||
Greedy sampling only (`do_sample=False`). `torch.multinomial` is not
|
||||
CUDA-graph-friendly; a downstream stochastic sanity check runs in a separate,
|
||||
non-graphed path.
|
||||
|
||||
Usage:
|
||||
from cb_sync_driver import cb_sync_generate, CBSyncConfig
|
||||
cfg = CBSyncConfig(max_new_tokens=512, use_cuda_graph=True)
|
||||
outputs = cb_sync_generate(model, generation_config, prompt_ids_list, cfg)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from transformers.generation.configuration_utils import GenerationConfig
|
||||
from transformers.generation.continuous_batching import (
|
||||
PagedAttentionCache,
|
||||
RequestStatus,
|
||||
)
|
||||
from transformers.generation.continuous_batching.continuous_api import (
|
||||
ContinuousBatchProcessor,
|
||||
ContinuousBatchingManager,
|
||||
)
|
||||
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.
|
||||
warmup_steps: int = 2
|
||||
# Generation config knobs (forwarded to the manager's GenerationConfig).
|
||||
do_sample: bool = False # greedy only (CUDA-graph safe)
|
||||
eos_token_id: Optional[int] = None
|
||||
pad_token_id: Optional[int] = None
|
||||
# Paged cache upper bounds; keep well above the default 256 / 4096.
|
||||
max_batch_tokens: int = 8192
|
||||
num_blocks: int = 8192
|
||||
# Progress callback (step_index, tokens_produced_total) -> None.
|
||||
on_step: Optional[callable] = field(default=None)
|
||||
|
||||
|
||||
class SyncCBDriver:
|
||||
"""Main-thread driver that owns the PagedAttentionCache,
|
||||
ContinuousBatchProcessor, and (optionally) a captured CUDA graph.
|
||||
|
||||
Unlike `ContinuousBatchingManager.start()`, there is no background
|
||||
thread; `drive_until_empty()` blocks until every pending request is
|
||||
finished.
|
||||
"""
|
||||
|
||||
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.
|
||||
gc = GenerationConfig.from_dict(generation_config.to_dict())
|
||||
gc.do_sample = cfg.do_sample
|
||||
if cfg.max_new_tokens:
|
||||
gc.max_new_tokens = cfg.max_new_tokens
|
||||
if cfg.eos_token_id is not None:
|
||||
gc.eos_token_id = cfg.eos_token_id
|
||||
if cfg.pad_token_id is not None:
|
||||
gc.pad_token_id = cfg.pad_token_id
|
||||
gc.max_batch_tokens = cfg.max_batch_tokens
|
||||
gc.num_blocks = cfg.num_blocks
|
||||
# Paged cache reads these at init.
|
||||
self.generation_config = gc
|
||||
|
||||
# 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
|
||||
)
|
||||
# The manager's `use_cuda_graph` is checked inside `warmup()`, but its
|
||||
# `__init__` refuses to set it. Set it directly now that we bypass
|
||||
# `init_continuous_batching`.
|
||||
self.manager.use_cuda_graph = cfg.use_cuda_graph
|
||||
|
||||
# Stand up the cache + processor ourselves so `_inner_generation_loop`
|
||||
# has everything it needs.
|
||||
self.cache = PagedAttentionCache(
|
||||
self.model.config,
|
||||
gc,
|
||||
self.model.device,
|
||||
self.model.dtype,
|
||||
tp_size=getattr(self.model, "_tp_size", None),
|
||||
)
|
||||
self.batch_processor = ContinuousBatchProcessor(
|
||||
self.cache,
|
||||
self.model.config,
|
||||
gc,
|
||||
self.manager.input_queue,
|
||||
self.manager.output_queue,
|
||||
self.manager.stop_event,
|
||||
self.model.device,
|
||||
self.model.dtype,
|
||||
FIFOScheduler(self.cache, manual_eviction=False),
|
||||
streaming=False,
|
||||
manual_eviction=False,
|
||||
slice_inputs=False,
|
||||
)
|
||||
self.manager.batch_processor = self.batch_processor
|
||||
self._graph: Optional[torch.cuda.CUDAGraph] = None
|
||||
self._step_count = 0
|
||||
|
||||
def add_requests(self, prompt_ids_list: list[list[int]]) -> list[str]:
|
||||
return [self.manager.add_request(ids) for ids in prompt_ids_list]
|
||||
|
||||
def _graphed_step(self):
|
||||
"""Capture or replay the decode CUDA graph."""
|
||||
if self._graph is None:
|
||||
# Eager warmup to populate allocator + workspaces.
|
||||
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.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):
|
||||
self.manager._generation_step(self.batch_processor)
|
||||
else:
|
||||
self._graph.replay()
|
||||
|
||||
def drive_until_empty(self) -> dict[str, list[int]]:
|
||||
"""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():
|
||||
# 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():
|
||||
break
|
||||
# 2. GPU: forward (graphed on decode steps, eager on prefill).
|
||||
if self.cfg.use_cuda_graph and self._is_pure_decode():
|
||||
self._graphed_step()
|
||||
else:
|
||||
self.manager._generation_step(self.batch_processor)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
# 3. CPU: append new tokens, detect EOS, update scheduler.
|
||||
self.batch_processor.update_batch()
|
||||
self._step_count += 1
|
||||
if self.cfg.on_step is not None:
|
||||
self.cfg.on_step(self._step_count, self._produced())
|
||||
# 4. Drain output_queue into results dict.
|
||||
while True:
|
||||
try:
|
||||
out = self.manager.output_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
break
|
||||
if out.status == RequestStatus.FINISHED:
|
||||
results[out.request_id] = out.generated_tokens
|
||||
# Final drain after loop exits.
|
||||
while True:
|
||||
try:
|
||||
out = self.manager.output_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
break
|
||||
if out.status == RequestStatus.FINISHED:
|
||||
results[out.request_id] = out.generated_tokens
|
||||
return results
|
||||
|
||||
def _is_pure_decode(self) -> bool:
|
||||
"""A decode-only batch has every request contributing exactly one
|
||||
query token (q_len == b_size). Prefill batches have q_len >> b_size.
|
||||
Shape consistency between decodes is what makes the graph replayable.
|
||||
"""
|
||||
try:
|
||||
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())
|
||||
|
||||
def close(self):
|
||||
# Caches hold GPU memory; free them explicitly.
|
||||
self._graph = None
|
||||
self.cache = None
|
||||
self.batch_processor = None
|
||||
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]]:
|
||||
"""One-shot entrypoint: build a driver, submit, drain, close.
|
||||
|
||||
Matches the semantics of `model.generate_batch(...)` but on the main
|
||||
thread with optional CUDA graph capture.
|
||||
"""
|
||||
driver = SyncCBDriver(model, generation_config, cfg)
|
||||
driver.add_requests(prompt_ids_list)
|
||||
try:
|
||||
return driver.drive_until_empty()
|
||||
finally:
|
||||
driver.close()
|
||||
|
||||
|
||||
# Simple microbench harness so the file is runnable standalone.
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
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)
|
||||
args = parser.parse_args()
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(args.model_name)
|
||||
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,
|
||||
).to("cuda")
|
||||
model.eval()
|
||||
|
||||
from unsloth_grpo_common import (
|
||||
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]
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
# Warmup (first 16 prompts).
|
||||
_ = cb_sync_generate(model, gc, prompt_ids[:16], cfg)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
wall_times = []
|
||||
total_decoded = 0
|
||||
for _ in range(2):
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
results = cb_sync_generate(model, gc, prompt_ids, cfg)
|
||||
torch.cuda.synchronize()
|
||||
wall_times.append(time.perf_counter() - t0)
|
||||
total_decoded = sum(len(v) for v in results.values())
|
||||
|
||||
med = sorted(wall_times)[len(wall_times) // 2]
|
||||
out = {
|
||||
"backend": "cb_sync_driver",
|
||||
"use_cuda_graph": args.use_cuda_graph,
|
||||
"attn_impl": args.attn_impl,
|
||||
"n_prompts": args.n_prompts,
|
||||
"n_decoded_tokens": total_decoded,
|
||||
"wall_times_s": wall_times,
|
||||
"median_wall_s": med,
|
||||
"decode_tps": total_decoded / med if med else 0,
|
||||
"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)
|
||||
with open(args.stats_path, "w") as f:
|
||||
json.dump(out, f, indent=2)
|
||||
print(json.dumps(out, indent=2))
|
||||
|
|
@ -50,8 +50,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},
|
||||
|
|
@ -60,11 +60,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
|
||||
|
|
@ -75,37 +75,35 @@ def run_vllm(args):
|
|||
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)
|
||||
|
||||
lora_request = None
|
||||
if args.lora_adapter:
|
||||
from vllm.lora.request import LoRARequest
|
||||
|
||||
lora_request = LoRARequest("fresh", 1, str(Path(args.lora_adapter).resolve()))
|
||||
|
||||
from vllm import SamplingParams
|
||||
|
||||
sp = SamplingParams(
|
||||
temperature = args.temperature,
|
||||
top_p = args.top_p,
|
||||
min_p = args.min_p,
|
||||
top_k = args.top_k,
|
||||
seed = 3407,
|
||||
max_tokens = args.max_new_tokens,
|
||||
stop = [tokenizer.eos_token],
|
||||
include_stop_str_in_output = True,
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
min_p=args.min_p,
|
||||
top_k=args.top_k,
|
||||
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 = lora_request)
|
||||
_ = model.fast_generate(warmup_text, sampling_params=sp, lora_request=lora_request)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
n_prompt_tokens = sum(len(p) for p in prompt_ids)
|
||||
|
|
@ -116,7 +114,7 @@ def run_vllm(args):
|
|||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
outputs = model.fast_generate(
|
||||
prompts_text, sampling_params = sp, lora_request = lora_request
|
||||
prompts_text, sampling_params=sp, lora_request=lora_request
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
wall_times.append(time.perf_counter() - t0)
|
||||
|
|
@ -124,11 +122,7 @@ def run_vllm(args):
|
|||
last_outputs = outputs
|
||||
|
||||
med = sorted(wall_times)[len(wall_times) // 2]
|
||||
sample_texts = (
|
||||
[o.outputs[0].text[:200] for o in (last_outputs[:3] or [])]
|
||||
if last_outputs
|
||||
else []
|
||||
)
|
||||
sample_texts = [o.outputs[0].text[:200] for o in (last_outputs[:3] or [])] if last_outputs else []
|
||||
return {
|
||||
"backend": "vllm",
|
||||
"lora_adapter": args.lora_adapter,
|
||||
|
|
@ -157,17 +151,16 @@ 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.lora_adapter:
|
||||
from peft import PeftModel
|
||||
|
||||
# NOTE: no merge_adapter -- we measure LoRA-active inference.
|
||||
model = PeftModel.from_pretrained(
|
||||
model, str(Path(args.lora_adapter).resolve()), is_trainable = False
|
||||
model, str(Path(args.lora_adapter).resolve()), is_trainable=False
|
||||
)
|
||||
model.eval()
|
||||
|
||||
|
|
@ -176,16 +169,16 @@ def run_tpaged(args):
|
|||
prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts)
|
||||
|
||||
gen_config = GenerationConfig(
|
||||
max_new_tokens = args.max_new_tokens,
|
||||
do_sample = True,
|
||||
temperature = args.temperature,
|
||||
top_p = args.top_p,
|
||||
min_p = args.min_p,
|
||||
top_k = args.top_k,
|
||||
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=args.temperature,
|
||||
top_p=args.top_p,
|
||||
min_p=args.min_p,
|
||||
top_k=args.top_k,
|
||||
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,
|
||||
)
|
||||
gen_config.max_batch_tokens = args.max_batch_tokens
|
||||
gen_config.num_blocks = args.num_blocks
|
||||
|
|
@ -195,9 +188,7 @@ def run_tpaged(args):
|
|||
|
||||
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)
|
||||
|
|
@ -209,7 +200,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)
|
||||
|
|
@ -221,7 +212,7 @@ def run_tpaged(args):
|
|||
if last_outputs is not None:
|
||||
for k in list(last_outputs.keys())[:3]:
|
||||
toks = last_outputs[k].generated_tokens
|
||||
sample_texts.append(tokenizer.decode(toks, skip_special_tokens = False)[:200])
|
||||
sample_texts.append(tokenizer.decode(toks, skip_special_tokens=False)[:200])
|
||||
|
||||
med = sorted(wall_times)[len(wall_times) // 2]
|
||||
return {
|
||||
|
|
@ -257,69 +248,59 @@ def run_unsloth_fi_false(args):
|
|||
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 = False,
|
||||
max_lora_rank = 32,
|
||||
model_name=args.model_name,
|
||||
max_seq_length=args.max_seq_length,
|
||||
load_in_4bit=False,
|
||||
fast_inference=False,
|
||||
max_lora_rank=32,
|
||||
)
|
||||
# Attach LoRA rank 32 the same way the GRPO notebook does.
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r = 32,
|
||||
target_modules = [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
r=32,
|
||||
target_modules=[
|
||||
"q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",
|
||||
],
|
||||
lora_alpha = 64,
|
||||
use_gradient_checkpointing = "unsloth",
|
||||
random_state = 3407,
|
||||
lora_alpha=64,
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
)
|
||||
|
||||
# Optional: overlay a shared adapter so weights match other backends.
|
||||
if args.lora_adapter:
|
||||
from safetensors import safe_open
|
||||
|
||||
adapter_file = Path(args.lora_adapter).resolve() / "adapter_model.safetensors"
|
||||
loaded_tensors = {}
|
||||
with safe_open(str(adapter_file), framework = "pt") as f:
|
||||
with safe_open(str(adapter_file), framework="pt") as f:
|
||||
for key in f.keys():
|
||||
loaded_tensors[key] = f.get_tensor(key)
|
||||
# PEFT saves with keys like `base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight`.
|
||||
# Unsloth's `get_peft_model` produces the same key shape.
|
||||
own_state = {n: p for n, p in model.named_parameters() if "lora_" in n}
|
||||
# Re-key by the suffix after `base_model.model.`.
|
||||
# Both PEFT and Unsloth's `get_peft_model` produce parameter names with
|
||||
# `base_model.model.` prefix plus `.lora_{A,B}.default.weight`. Build a
|
||||
# normalized (core-path) -> param map, then match by core path only.
|
||||
def _core(name: str) -> str:
|
||||
n = name
|
||||
for pref in ("base_model.model.", "model."):
|
||||
if n.startswith(pref):
|
||||
n = n[len(pref):]
|
||||
n = n.replace(".lora_A.default.", ".lora_A.").replace(
|
||||
".lora_B.default.", ".lora_B.")
|
||||
return n
|
||||
own_by_core = {}
|
||||
for n, p in model.named_parameters():
|
||||
if "lora_" in n:
|
||||
own_by_core.setdefault(_core(n), []).append(p)
|
||||
matched = 0
|
||||
with torch.no_grad():
|
||||
for name, tensor in loaded_tensors.items():
|
||||
# Try direct + strip `base_model.model.` prefix variants.
|
||||
candidates = [
|
||||
name,
|
||||
name.replace("base_model.model.", ""),
|
||||
"base_model.model." + name,
|
||||
]
|
||||
for cand in candidates:
|
||||
# PEFT sometimes inserts `.default.` between module and lora_A.
|
||||
variants = [cand, cand.replace(".default.", ".")]
|
||||
for v in variants:
|
||||
# own_state keys typically have `.default.weight` suffix
|
||||
for own_name, own in own_state.items():
|
||||
if own_name.endswith(
|
||||
v.split("base_model.model.")[-1]
|
||||
) or v.endswith(own_name.split("base_model.model.")[-1]):
|
||||
if own.shape == tensor.shape:
|
||||
own.data.copy_(tensor.to(own.device, own.dtype))
|
||||
matched += 1
|
||||
break
|
||||
print(
|
||||
f"[unsloth_fi_false] LoRA weight sync matched {matched} tensors "
|
||||
f"(out of {len(loaded_tensors)} adapter entries)."
|
||||
)
|
||||
core = _core(name)
|
||||
for own in own_by_core.get(core, []):
|
||||
if own.shape == tensor.shape:
|
||||
own.data.copy_(tensor.to(own.device, own.dtype))
|
||||
matched += 1
|
||||
break
|
||||
print(f"[unsloth_fi_false] LoRA weight sync matched {matched} tensors "
|
||||
f"(out of {len(loaded_tensors)} adapter entries).")
|
||||
|
||||
FastLanguageModel.for_inference(model)
|
||||
|
||||
|
|
@ -327,29 +308,28 @@ def run_unsloth_fi_false(args):
|
|||
|
||||
# `model.generate` accepts batched input_ids; pad to max length.
|
||||
from transformers import GenerationConfig
|
||||
|
||||
if tokenizer.padding_side != "left":
|
||||
tokenizer.padding_side = "left" # decoder needs left padding
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
gen_config = GenerationConfig(
|
||||
max_new_tokens = args.max_new_tokens,
|
||||
do_sample = True,
|
||||
temperature = args.temperature,
|
||||
top_p = args.top_p,
|
||||
min_p = args.min_p,
|
||||
top_k = args.top_k,
|
||||
pad_token_id = tokenizer.pad_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=args.temperature,
|
||||
top_p=args.top_p,
|
||||
min_p=args.min_p,
|
||||
top_k=args.top_k,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
bos_token_id=tokenizer.bos_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
use_cache=True,
|
||||
)
|
||||
|
||||
def _batched_generate(texts):
|
||||
batch = tokenizer(texts, return_tensors = "pt", padding = True).to("cuda")
|
||||
batch = tokenizer(texts, return_tensors="pt", padding=True).to("cuda")
|
||||
with torch.inference_mode():
|
||||
out = model.generate(**batch, generation_config = gen_config)
|
||||
out = model.generate(**batch, generation_config=gen_config)
|
||||
prompt_len = batch["input_ids"].shape[1]
|
||||
return out, prompt_len
|
||||
|
||||
|
|
@ -370,9 +350,7 @@ def run_unsloth_fi_false(args):
|
|||
wall_times.append(time.perf_counter() - t0)
|
||||
# Count generated tokens past prompt_len per sequence (subtract any
|
||||
# trailing pad-only tail by comparing against EOS).
|
||||
total_decoded = int(
|
||||
(out_ids[:, prompt_len:] != tokenizer.pad_token_id).sum().item()
|
||||
)
|
||||
total_decoded = int((out_ids[:, prompt_len:] != tokenizer.pad_token_id).sum().item())
|
||||
last_out_ids = out_ids
|
||||
last_prompt_len = prompt_len
|
||||
|
||||
|
|
@ -380,11 +358,8 @@ def run_unsloth_fi_false(args):
|
|||
sample_texts = []
|
||||
if last_out_ids is not None:
|
||||
for i in range(min(3, last_out_ids.shape[0])):
|
||||
sample_texts.append(
|
||||
tokenizer.decode(
|
||||
last_out_ids[i, last_prompt_len:], skip_special_tokens = False
|
||||
)[:200]
|
||||
)
|
||||
sample_texts.append(tokenizer.decode(
|
||||
last_out_ids[i, last_prompt_len:], skip_special_tokens=False)[:200])
|
||||
|
||||
return {
|
||||
"backend": "unsloth_fi_false",
|
||||
|
|
@ -403,35 +378,30 @@ def run_unsloth_fi_false(args):
|
|||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument(
|
||||
"--backend", choices = ["vllm", "tpaged", "unsloth_fi_false"], 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 = 32)
|
||||
p.add_argument("--n_rounds", type = int, default = 2)
|
||||
p.add_argument("--max_new_tokens", type = int, default = 512)
|
||||
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")
|
||||
p.add_argument(
|
||||
"--lora_adapter",
|
||||
default = None,
|
||||
help = "Path to a PEFT adapter (rank 32) applied in every backend.",
|
||||
)
|
||||
p.add_argument("--temperature", type = float, default = 0.1)
|
||||
p.add_argument("--top_p", type = float, default = 0.97)
|
||||
p.add_argument("--min_p", type = float, default = 0.5)
|
||||
p.add_argument("--top_k", type = int, default = 5)
|
||||
p.add_argument("--stats_path", required = True)
|
||||
p.add_argument("--backend", choices=["vllm", "tpaged", "unsloth_fi_false"], 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=32)
|
||||
p.add_argument("--n_rounds", type=int, default=2)
|
||||
p.add_argument("--max_new_tokens", type=int, default=512)
|
||||
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")
|
||||
p.add_argument("--lora_adapter", default=None,
|
||||
help="Path to a PEFT adapter (rank 32) applied in every backend.")
|
||||
p.add_argument("--temperature", type=float, default=0.1)
|
||||
p.add_argument("--top_p", type=float, default=0.97)
|
||||
p.add_argument("--min_p", type=float, default=0.5)
|
||||
p.add_argument("--top_k", type=int, default=5)
|
||||
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":
|
||||
|
|
@ -449,8 +419,8 @@ def main():
|
|||
"top_k": args.top_k,
|
||||
}
|
||||
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))
|
||||
os._exit(0)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -20,16 +20,16 @@ from pathlib import Path
|
|||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
|
||||
p.add_argument("--output", default = "outputs/lora_rank32_fresh")
|
||||
p.add_argument("--rank", type = int, default = 32)
|
||||
p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base")
|
||||
p.add_argument("--output", default="outputs/lora_rank32_fresh")
|
||||
p.add_argument("--rank", type=int, default=32)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
out_dir = Path(args.output).resolve()
|
||||
out_dir.mkdir(parents = True, exist_ok = True)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Use vanilla HF -- PEFT's save_pretrained yields the canonical
|
||||
# adapter_config.json + adapter_model.safetensors that vLLM's LoRARequest
|
||||
|
|
@ -45,23 +45,16 @@ def main():
|
|||
# bf16 base; we only need structure + save. Keep on CPU to avoid a GPU load
|
||||
# just for `save_pretrained`.
|
||||
print(f"[make_lora_adapter] Loading {args.model_name} on CPU...")
|
||||
model = AutoModelForCausalLM.from_pretrained(args.model_name, dtype = torch.bfloat16)
|
||||
model = AutoModelForCausalLM.from_pretrained(args.model_name, dtype=torch.bfloat16)
|
||||
|
||||
peft_cfg = LoraConfig(
|
||||
r = args.rank,
|
||||
lora_alpha = args.rank * 2,
|
||||
target_modules = [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
],
|
||||
bias = "none",
|
||||
task_type = "CAUSAL_LM",
|
||||
lora_dropout = 0.0,
|
||||
r=args.rank,
|
||||
lora_alpha=args.rank * 2,
|
||||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj"],
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
lora_dropout=0.0,
|
||||
)
|
||||
peft_model = get_peft_model(model, peft_cfg)
|
||||
peft_model.print_trainable_parameters()
|
||||
|
|
@ -73,31 +66,26 @@ def main():
|
|||
with torch.no_grad():
|
||||
for name, p in peft_model.named_parameters():
|
||||
if "lora_B" in name:
|
||||
p.normal_(mean = 0.0, std = 1e-4)
|
||||
p.normal_(mean=0.0, std=1e-4)
|
||||
n_reinit += 1
|
||||
print(
|
||||
f"[make_lora_adapter] Reinitialized {n_reinit} lora_B matrices with tiny gaussian."
|
||||
)
|
||||
print(f"[make_lora_adapter] Reinitialized {n_reinit} lora_B matrices with tiny gaussian.")
|
||||
|
||||
peft_model.save_pretrained(str(out_dir))
|
||||
tok.save_pretrained(str(out_dir))
|
||||
|
||||
# Sanity: verify safetensors file present and non-trivial.
|
||||
from safetensors import safe_open
|
||||
|
||||
st_path = out_dir / "adapter_model.safetensors"
|
||||
n_zero_tensors = 0
|
||||
n_tensors = 0
|
||||
with safe_open(str(st_path), framework = "pt") as f:
|
||||
with safe_open(str(st_path), framework="pt") as f:
|
||||
for key in f.keys():
|
||||
t = f.get_tensor(key)
|
||||
n_tensors += 1
|
||||
if (t == 0).all().item():
|
||||
n_zero_tensors += 1
|
||||
print(
|
||||
f"[make_lora_adapter] Wrote {n_tensors} tensors to {st_path} "
|
||||
f"({n_zero_tensors} all-zero)."
|
||||
)
|
||||
print(f"[make_lora_adapter] Wrote {n_tensors} tensors to {st_path} "
|
||||
f"({n_zero_tensors} all-zero).")
|
||||
print(f"[make_lora_adapter] Adapter saved to {out_dir}")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -33,31 +33,28 @@ for p in (HERE, WORKSPACE_ROOT):
|
|||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--stats_path", default = "logs/notebook_ref_10.json")
|
||||
p.add_argument("--output_dir", default = "outputs/notebook_ref_10")
|
||||
p.add_argument("--max_steps", type = int, default = 10)
|
||||
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("--gpu_memory_utilization", type = float, default = 0.85)
|
||||
p.add_argument("--num_generations", type = int, default = 4)
|
||||
p.add_argument("--per_device_train_batch_size", type = int, default = 1)
|
||||
p.add_argument("--temperature", type = float, default = 0.1)
|
||||
p.add_argument("--top_p", type = float, default = 0.97)
|
||||
p.add_argument("--min_p", type = float, default = 0.5)
|
||||
p.add_argument("--top_k", type = int, default = 5)
|
||||
p.add_argument(
|
||||
"--skip_sft_pre_finetune",
|
||||
action = "store_true",
|
||||
help = "Skip the format-priming SFT stage; go straight to GRPO.",
|
||||
)
|
||||
p.add_argument("--stats_path", default="logs/notebook_ref_10.json")
|
||||
p.add_argument("--output_dir", default="outputs/notebook_ref_10")
|
||||
p.add_argument("--max_steps", type=int, default=10)
|
||||
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("--gpu_memory_utilization", type=float, default=0.85)
|
||||
p.add_argument("--num_generations", type=int, default=4)
|
||||
p.add_argument("--per_device_train_batch_size", type=int, default=1)
|
||||
p.add_argument("--temperature", type=float, default=0.1)
|
||||
p.add_argument("--top_p", type=float, default=0.97)
|
||||
p.add_argument("--min_p", type=float, default=0.5)
|
||||
p.add_argument("--top_k", type=int, default=5)
|
||||
p.add_argument("--skip_sft_pre_finetune", action="store_true",
|
||||
help="Skip the format-priming SFT stage; go straight to GRPO.")
|
||||
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(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)
|
||||
|
||||
# Import order matters: unsloth must come before transformers/trl.
|
||||
os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
|
||||
|
|
@ -65,28 +62,23 @@ def main():
|
|||
import torch # noqa: E402
|
||||
|
||||
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 = args.lora_rank,
|
||||
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=args.lora_rank,
|
||||
gpu_memory_utilization=args.gpu_memory_utilization,
|
||||
)
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r = args.lora_rank,
|
||||
target_modules = [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
r=args.lora_rank,
|
||||
target_modules=[
|
||||
"q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",
|
||||
],
|
||||
lora_alpha = args.lora_rank * 2,
|
||||
use_gradient_checkpointing = "unsloth",
|
||||
random_state = 3407,
|
||||
lora_alpha=args.lora_rank * 2,
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=3407,
|
||||
)
|
||||
|
||||
reasoning_start = "<start_working_out>"
|
||||
|
|
@ -127,29 +119,16 @@ def main():
|
|||
import numpy as np
|
||||
|
||||
if not args.skip_sft_pre_finetune:
|
||||
sft_ds = load_dataset("unsloth/OpenMathReasoning-mini", split = "cot")
|
||||
sft_df = sft_ds.to_pandas()[
|
||||
["expected_answer", "problem", "generated_solution"]
|
||||
]
|
||||
is_number = pd.to_numeric(
|
||||
pd.Series(sft_df["expected_answer"]), errors = "coerce"
|
||||
).notnull()
|
||||
sft_ds = load_dataset("unsloth/OpenMathReasoning-mini", split="cot")
|
||||
sft_df = sft_ds.to_pandas()[["expected_answer", "problem", "generated_solution"]]
|
||||
is_number = pd.to_numeric(pd.Series(sft_df["expected_answer"]), errors="coerce").notnull()
|
||||
sft_df = sft_df.iloc[np.where(is_number)[0]]
|
||||
|
||||
def format_dataset(x):
|
||||
thoughts = (
|
||||
x["generated_solution"]
|
||||
.replace("<think>", "")
|
||||
.replace("</think>", "")
|
||||
.strip()
|
||||
)
|
||||
thoughts = x["generated_solution"].replace("<think>", "").replace("</think>", "").strip()
|
||||
final_prompt = (
|
||||
reasoning_start
|
||||
+ thoughts
|
||||
+ reasoning_end
|
||||
+ solution_start
|
||||
+ x["expected_answer"]
|
||||
+ solution_end
|
||||
reasoning_start + thoughts + reasoning_end
|
||||
+ solution_start + x["expected_answer"] + solution_end
|
||||
)
|
||||
return [
|
||||
{"role": "system", "content": system_prompt},
|
||||
|
|
@ -157,69 +136,61 @@ def main():
|
|||
{"role": "assistant", "content": final_prompt},
|
||||
]
|
||||
|
||||
sft_df["Messages"] = sft_df.apply(format_dataset, axis = 1)
|
||||
sft_df["N"] = sft_df["Messages"].apply(
|
||||
lambda m: len(tokenizer.apply_chat_template(m))
|
||||
)
|
||||
sft_df["Messages"] = sft_df.apply(format_dataset, axis=1)
|
||||
sft_df["N"] = sft_df["Messages"].apply(lambda m: len(tokenizer.apply_chat_template(m)))
|
||||
sft_df = sft_df.loc[sft_df["N"] <= args.max_seq_length / 2].copy()
|
||||
sft_df["text"] = tokenizer.apply_chat_template(
|
||||
sft_df["Messages"].values.tolist(), tokenize = False
|
||||
sft_df["Messages"].values.tolist(), tokenize=False
|
||||
)
|
||||
sft_dataset = Dataset.from_pandas(sft_df)
|
||||
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
|
||||
sft_trainer = SFTTrainer(
|
||||
model = model,
|
||||
tokenizer = tokenizer,
|
||||
train_dataset = sft_dataset,
|
||||
args = SFTConfig(
|
||||
dataset_text_field = "text",
|
||||
per_device_train_batch_size = 1,
|
||||
gradient_accumulation_steps = 1,
|
||||
warmup_steps = 5,
|
||||
num_train_epochs = 2,
|
||||
learning_rate = 2e-4,
|
||||
logging_steps = 5,
|
||||
optim = "adamw_8bit",
|
||||
weight_decay = 0.001,
|
||||
lr_scheduler_type = "linear",
|
||||
seed = 3407,
|
||||
report_to = "none",
|
||||
output_dir = os.path.join(args.output_dir, "sft"),
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=sft_dataset,
|
||||
args=SFTConfig(
|
||||
dataset_text_field="text",
|
||||
per_device_train_batch_size=1,
|
||||
gradient_accumulation_steps=1,
|
||||
warmup_steps=5,
|
||||
num_train_epochs=2,
|
||||
learning_rate=2e-4,
|
||||
logging_steps=5,
|
||||
optim="adamw_8bit",
|
||||
weight_decay=0.001,
|
||||
lr_scheduler_type="linear",
|
||||
seed=3407,
|
||||
report_to="none",
|
||||
output_dir=os.path.join(args.output_dir, "sft"),
|
||||
),
|
||||
)
|
||||
sft_trainer.train()
|
||||
del sft_dataset, sft_df, sft_ds, sft_trainer
|
||||
torch.cuda.empty_cache()
|
||||
import gc
|
||||
|
||||
gc.collect()
|
||||
|
||||
# --- GRPO stage -----------------------------------------------------------
|
||||
dataset = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
|
||||
dataset = dataset.map(
|
||||
lambda x: {
|
||||
"prompt": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": x["prompt"]},
|
||||
],
|
||||
"answer": x["solution"],
|
||||
}
|
||||
)
|
||||
dataset = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train")
|
||||
dataset = dataset.map(lambda x: {
|
||||
"prompt": [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": x["prompt"]},
|
||||
],
|
||||
"answer": x["solution"],
|
||||
})
|
||||
|
||||
solution_end_regex = (
|
||||
r"</SOLUTION>[\s]{0,}" + "(?:" + re.escape(tokenizer.eos_token) + ")?"
|
||||
)
|
||||
solution_end_regex = r"</SOLUTION>[\s]{0,}" + "(?:" + re.escape(tokenizer.eos_token) + ")?"
|
||||
match_format = re.compile(
|
||||
rf"{reasoning_end}.*?"
|
||||
rf"{solution_start}(.+?){solution_end_regex}"
|
||||
rf"[\s]{{0,}}$",
|
||||
flags = re.MULTILINE | re.DOTALL,
|
||||
flags=re.MULTILINE | re.DOTALL,
|
||||
)
|
||||
match_numbers = re.compile(
|
||||
solution_start + r".*?[\s]{0,}([-]?[\d\.\,]{1,})",
|
||||
flags = re.MULTILINE | re.DOTALL,
|
||||
flags=re.MULTILINE | re.DOTALL,
|
||||
)
|
||||
|
||||
def match_format_exactly(completions, **kwargs):
|
||||
|
|
@ -291,12 +262,10 @@ def main():
|
|||
|
||||
# Filter long prompts.
|
||||
tokenized = dataset.map(
|
||||
lambda x: {
|
||||
"tokens": tokenizer.apply_chat_template(
|
||||
x["prompt"], add_generation_prompt = True, tokenize = True
|
||||
)
|
||||
},
|
||||
batched = False,
|
||||
lambda x: {"tokens": tokenizer.apply_chat_template(
|
||||
x["prompt"], add_generation_prompt=True, tokenize=True
|
||||
)},
|
||||
batched=False,
|
||||
)
|
||||
tokenized = tokenized.map(lambda x: {"L": len(x["tokens"])})
|
||||
maximum_length = int(np.quantile(tokenized["L"], 0.9))
|
||||
|
|
@ -308,63 +277,60 @@ def main():
|
|||
max_completion_length = args.max_seq_length - max_prompt_length
|
||||
|
||||
from vllm import SamplingParams
|
||||
|
||||
vllm_sampling_params = SamplingParams(
|
||||
temperature = args.temperature,
|
||||
top_p = args.top_p,
|
||||
min_p = args.min_p,
|
||||
top_k = args.top_k,
|
||||
seed = 3407,
|
||||
stop = [tokenizer.eos_token],
|
||||
include_stop_str_in_output = True,
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
min_p=args.min_p,
|
||||
top_k=args.top_k,
|
||||
seed=3407,
|
||||
stop=[tokenizer.eos_token],
|
||||
include_stop_str_in_output=True,
|
||||
)
|
||||
|
||||
from trl import GRPOConfig, GRPOTrainer
|
||||
|
||||
training_args = GRPOConfig(
|
||||
vllm_sampling_params = vllm_sampling_params,
|
||||
temperature = args.temperature,
|
||||
top_p = args.top_p,
|
||||
top_k = args.top_k,
|
||||
learning_rate = 5e-6,
|
||||
weight_decay = 0.001,
|
||||
warmup_ratio = 0.1,
|
||||
lr_scheduler_type = "linear",
|
||||
optim = "adamw_8bit",
|
||||
logging_steps = 1,
|
||||
per_device_train_batch_size = args.per_device_train_batch_size,
|
||||
gradient_accumulation_steps = 1,
|
||||
num_generations = args.num_generations,
|
||||
max_prompt_length = max_prompt_length,
|
||||
max_completion_length = max_completion_length,
|
||||
max_steps = args.max_steps,
|
||||
save_steps = args.max_steps + 1,
|
||||
report_to = "none",
|
||||
output_dir = args.output_dir,
|
||||
seed = 3407,
|
||||
vllm_sampling_params=vllm_sampling_params,
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
top_k=args.top_k,
|
||||
learning_rate=5e-6,
|
||||
weight_decay=0.001,
|
||||
warmup_ratio=0.1,
|
||||
lr_scheduler_type="linear",
|
||||
optim="adamw_8bit",
|
||||
logging_steps=1,
|
||||
per_device_train_batch_size=args.per_device_train_batch_size,
|
||||
gradient_accumulation_steps=1,
|
||||
num_generations=args.num_generations,
|
||||
max_prompt_length=max_prompt_length,
|
||||
max_completion_length=max_completion_length,
|
||||
max_steps=args.max_steps,
|
||||
save_steps=args.max_steps + 1,
|
||||
report_to="none",
|
||||
output_dir=args.output_dir,
|
||||
seed=3407,
|
||||
)
|
||||
|
||||
from torch_debugging_utils import StatisticsCallback
|
||||
|
||||
stats_cb = StatisticsCallback(
|
||||
track_loss = True,
|
||||
track_grad_norm = True,
|
||||
track_memory = True,
|
||||
track_tensor_stats = False, # hooks are noisy + slow on GRPO model
|
||||
track_loss=True,
|
||||
track_grad_norm=True,
|
||||
track_memory=True,
|
||||
track_tensor_stats=False, # hooks are noisy + slow on GRPO model
|
||||
)
|
||||
|
||||
trainer = GRPOTrainer(
|
||||
model = model,
|
||||
processing_class = tokenizer,
|
||||
reward_funcs = [
|
||||
model=model,
|
||||
processing_class=tokenizer,
|
||||
reward_funcs=[
|
||||
match_format_exactly,
|
||||
match_format_approximately,
|
||||
check_answer,
|
||||
check_numbers,
|
||||
],
|
||||
args = training_args,
|
||||
train_dataset = dataset,
|
||||
callbacks = [stats_cb],
|
||||
args=training_args,
|
||||
train_dataset=dataset,
|
||||
callbacks=[stats_cb],
|
||||
)
|
||||
|
||||
t0 = time.perf_counter()
|
||||
|
|
@ -395,39 +361,30 @@ def main():
|
|||
"logs_path": args.stats_path,
|
||||
"peak_memory_gb": torch.cuda.max_memory_allocated() / 1024**3,
|
||||
}
|
||||
print(json.dumps(summary, indent = 2))
|
||||
print(json.dumps(summary, indent=2))
|
||||
|
||||
# Canonical quick-inference: produce a few generations for the writeup.
|
||||
rollouts = []
|
||||
try:
|
||||
from vllm import SamplingParams as SP
|
||||
|
||||
sp_sample = SP(
|
||||
temperature = args.temperature,
|
||||
top_p = args.top_p,
|
||||
min_p = args.min_p,
|
||||
top_k = args.top_k,
|
||||
max_tokens = 256,
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
min_p=args.min_p,
|
||||
top_k=args.top_k,
|
||||
max_tokens=256,
|
||||
)
|
||||
probe_prompts = [
|
||||
[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": "What is the sqrt of 101?"},
|
||||
],
|
||||
[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": "If 3x+7 = 22, what is x?"},
|
||||
],
|
||||
[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": "What is 17 * 13?"},
|
||||
],
|
||||
[{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": "What is the sqrt of 101?"}],
|
||||
[{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": "If 3x+7 = 22, what is x?"}],
|
||||
[{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": "What is 17 * 13?"}],
|
||||
]
|
||||
texts = [
|
||||
tokenizer.apply_chat_template(p, add_generation_prompt = True, tokenize = False)
|
||||
for p in probe_prompts
|
||||
]
|
||||
outs = model.fast_generate(texts, sampling_params = sp_sample, lora_request = None)
|
||||
texts = [tokenizer.apply_chat_template(p, add_generation_prompt=True, tokenize=False)
|
||||
for p in probe_prompts]
|
||||
outs = model.fast_generate(texts, sampling_params=sp_sample, lora_request=None)
|
||||
for t, o in zip(texts, outs):
|
||||
rollouts.append({"prompt": t, "completion": o.outputs[0].text})
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -42,30 +42,36 @@ os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
|
|||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument(
|
||||
"--backend",
|
||||
choices = ["vllm", "unsloth_fi_false", "cb_paged", "cb_sdpa", "naive_trl"],
|
||||
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("--lora_rank", type = int, default = 32)
|
||||
p.add_argument("--max_steps", type = int, default = 10)
|
||||
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("--gpu_memory_utilization", type = float, default = 0.75)
|
||||
p.add_argument("--temperature", type = float, default = 0.1)
|
||||
p.add_argument("--top_p", type = float, default = 0.97)
|
||||
p.add_argument("--min_p", type = float, default = 0.5)
|
||||
p.add_argument("--top_k", type = int, default = 5)
|
||||
p.add_argument("--learning_rate", type = float, default = 5e-6)
|
||||
p.add_argument("--max_batch_tokens", type = int, default = 8192)
|
||||
p.add_argument("--num_blocks", type = int, default = 8192)
|
||||
p.add_argument("--persistent_cb", action = "store_true")
|
||||
p.add_argument("--output_dir", required = True)
|
||||
p.add_argument("--stats_path", required = True)
|
||||
p.add_argument("--seed", type = int, default = 3407)
|
||||
p.add_argument("--backend",
|
||||
choices=["vllm", "unsloth_fi_false", "cb_paged", "cb_sdpa", "naive_trl"],
|
||||
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("--lora_rank", type=int, default=32)
|
||||
p.add_argument("--max_steps", type=int, default=10)
|
||||
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("--gpu_memory_utilization", type=float, default=0.75)
|
||||
p.add_argument("--temperature", type=float, default=0.1)
|
||||
p.add_argument("--top_p", type=float, default=0.97)
|
||||
p.add_argument("--min_p", type=float, default=0.5)
|
||||
p.add_argument("--top_k", type=int, default=5)
|
||||
p.add_argument("--learning_rate", type=float, default=5e-6)
|
||||
p.add_argument("--max_batch_tokens", type=int, default=8192)
|
||||
p.add_argument("--num_blocks", type=int, default=8192)
|
||||
p.add_argument("--persistent_cb", action="store_true")
|
||||
p.add_argument("--output_dir", required=True)
|
||||
p.add_argument("--stats_path", required=True)
|
||||
p.add_argument("--seed", type=int, default=3407)
|
||||
# Phase 4: torch.compile on the training forward.
|
||||
p.add_argument("--compile_mode", default=None,
|
||||
choices=[None, "default", "reduce-overhead",
|
||||
"max-autotune-no-cudagraphs"],
|
||||
help="If set, torch.compile(model.forward, mode=...) after "
|
||||
"the trainer is built. vllm backend is excluded; the "
|
||||
"rollout engine owns its own compile pipeline.")
|
||||
p.add_argument("--compile_dynamic", action="store_true", default=True)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
|
|
@ -73,18 +79,10 @@ def _prepare_common(args):
|
|||
"""Dataset + rewards are the same for every backend. Always uses the
|
||||
shared chat template and reward funcs from unsloth_grpo_common."""
|
||||
from unsloth_grpo_common import (
|
||||
apply_chat_template_to_tokenizer,
|
||||
build_dataset,
|
||||
build_reward_funcs,
|
||||
build_grpo_kwargs,
|
||||
)
|
||||
|
||||
return (
|
||||
apply_chat_template_to_tokenizer,
|
||||
build_dataset,
|
||||
build_reward_funcs,
|
||||
build_grpo_kwargs,
|
||||
apply_chat_template_to_tokenizer, build_dataset,
|
||||
build_reward_funcs, build_grpo_kwargs,
|
||||
)
|
||||
return apply_chat_template_to_tokenizer, build_dataset, build_reward_funcs, build_grpo_kwargs
|
||||
|
||||
|
||||
def _make_stats_callback():
|
||||
|
|
@ -92,12 +90,11 @@ def _make_stats_callback():
|
|||
grad-norm, memory, and wall time. Reward/KL are picked up from the TRL
|
||||
log dict via `on_log`."""
|
||||
from torch_debugging_utils import StatisticsCallback
|
||||
|
||||
return StatisticsCallback(
|
||||
track_loss = True,
|
||||
track_grad_norm = True,
|
||||
track_memory = True,
|
||||
track_tensor_stats = False,
|
||||
track_loss=True,
|
||||
track_grad_norm=True,
|
||||
track_memory=True,
|
||||
track_tensor_stats=False,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -107,13 +104,10 @@ def _maybe_shim_guided_decoding():
|
|||
the transformers-paged path. Inject a no-op shim if missing."""
|
||||
try:
|
||||
import vllm.sampling_params as sp
|
||||
|
||||
if not hasattr(sp, "GuidedDecodingParams"):
|
||||
|
||||
class _Shim:
|
||||
def __init__(self, *a, **kw):
|
||||
pass
|
||||
|
||||
sp.GuidedDecodingParams = _Shim
|
||||
except ImportError:
|
||||
pass
|
||||
|
|
@ -121,34 +115,22 @@ def _maybe_shim_guided_decoding():
|
|||
|
||||
def _load_unsloth(args, fast_inference: bool):
|
||||
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 = fast_inference,
|
||||
max_lora_rank = args.lora_rank,
|
||||
**(
|
||||
{"gpu_memory_utilization": args.gpu_memory_utilization}
|
||||
if fast_inference
|
||||
else {}
|
||||
),
|
||||
model_name=args.model_name,
|
||||
max_seq_length=args.max_seq_length,
|
||||
load_in_4bit=False,
|
||||
fast_inference=fast_inference,
|
||||
max_lora_rank=args.lora_rank,
|
||||
**({"gpu_memory_utilization": args.gpu_memory_utilization} if fast_inference else {}),
|
||||
)
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r = args.lora_rank,
|
||||
target_modules = [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
],
|
||||
lora_alpha = args.lora_rank * 2,
|
||||
use_gradient_checkpointing = "unsloth",
|
||||
random_state = args.seed,
|
||||
r=args.lora_rank,
|
||||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj"],
|
||||
lora_alpha=args.lora_rank * 2,
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=args.seed,
|
||||
)
|
||||
return model, tokenizer
|
||||
|
||||
|
|
@ -164,28 +146,21 @@ def _load_vanilla_hf(args, attn_impl: str):
|
|||
tokenizer.pad_token = tokenizer.eos_token
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model_name,
|
||||
dtype = torch.bfloat16,
|
||||
attn_implementation = attn_impl,
|
||||
dtype=torch.bfloat16,
|
||||
attn_implementation=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",
|
||||
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}
|
||||
gradient_checkpointing_kwargs={"use_reentrant": False}
|
||||
)
|
||||
except TypeError:
|
||||
model.gradient_checkpointing_enable()
|
||||
|
|
@ -195,46 +170,38 @@ def _load_vanilla_hf(args, attn_impl: str):
|
|||
|
||||
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)
|
||||
|
||||
import torch
|
||||
from torch_debugging_utils import set_all_seeds_fast
|
||||
|
||||
set_all_seeds_fast(args.seed)
|
||||
|
||||
# FA4 shim lives here so CB paths dispatch to Blackwell kernels.
|
||||
import flash_attn_fa4_shim # noqa: F401
|
||||
|
||||
flash_attn_fa4_shim.apply()
|
||||
_maybe_shim_guided_decoding()
|
||||
|
||||
(
|
||||
apply_chat_template_to_tokenizer,
|
||||
build_dataset,
|
||||
build_reward_funcs,
|
||||
build_grpo_kwargs,
|
||||
) = _prepare_common(args)
|
||||
(apply_chat_template_to_tokenizer, build_dataset,
|
||||
build_reward_funcs, build_grpo_kwargs) = _prepare_common(args)
|
||||
|
||||
# --- load model / tokenizer per backend -----------------------------------
|
||||
persistent_teardown_target = None
|
||||
if args.backend == "vllm":
|
||||
model, tokenizer = _load_unsloth(args, fast_inference = True)
|
||||
model, tokenizer = _load_unsloth(args, fast_inference=True)
|
||||
elif args.backend == "unsloth_fi_false":
|
||||
model, tokenizer = _load_unsloth(args, fast_inference = False)
|
||||
model, tokenizer = _load_unsloth(args, fast_inference=False)
|
||||
elif args.backend == "cb_paged":
|
||||
model, tokenizer = _load_vanilla_hf(args, attn_impl = "paged_attention")
|
||||
model, tokenizer = _load_vanilla_hf(args, attn_impl="paged_attention")
|
||||
elif args.backend == "cb_sdpa":
|
||||
model, tokenizer = _load_vanilla_hf(args, attn_impl = "sdpa_paged")
|
||||
model, tokenizer = _load_vanilla_hf(args, attn_impl="sdpa_paged")
|
||||
elif args.backend == "naive_trl":
|
||||
model, tokenizer = _load_vanilla_hf(args, attn_impl = "sdpa")
|
||||
model, tokenizer = _load_vanilla_hf(args, attn_impl="sdpa")
|
||||
else:
|
||||
raise ValueError(args.backend)
|
||||
|
||||
apply_chat_template_to_tokenizer(tokenizer)
|
||||
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"[{args.backend}] p90 prompt length = {maximum_length}")
|
||||
reward_funcs = build_reward_funcs(tokenizer)
|
||||
|
||||
|
|
@ -242,12 +209,12 @@ def main():
|
|||
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,
|
||||
)
|
||||
# Overwrite the equivalence-friendly sampling params.
|
||||
shared["temperature"] = args.temperature
|
||||
|
|
@ -258,24 +225,18 @@ def main():
|
|||
shared["learning_rate"] = args.learning_rate
|
||||
|
||||
from trl import GRPOConfig, GRPOTrainer
|
||||
|
||||
if args.backend == "vllm":
|
||||
from vllm import SamplingParams
|
||||
|
||||
vllm_sp = SamplingParams(
|
||||
temperature = args.temperature,
|
||||
top_p = args.top_p,
|
||||
min_p = args.min_p,
|
||||
top_k = args.top_k,
|
||||
seed = args.seed,
|
||||
stop = [tokenizer.eos_token],
|
||||
include_stop_str_in_output = True,
|
||||
temperature=args.temperature, top_p=args.top_p, min_p=args.min_p,
|
||||
top_k=args.top_k, seed=args.seed,
|
||||
stop=[tokenizer.eos_token], include_stop_str_in_output=True,
|
||||
)
|
||||
training_args = GRPOConfig(
|
||||
use_vllm = True,
|
||||
vllm_mode = "colocate",
|
||||
vllm_sampling_params = vllm_sp,
|
||||
vllm_gpu_memory_utilization = args.gpu_memory_utilization,
|
||||
use_vllm=True,
|
||||
vllm_mode="colocate",
|
||||
vllm_sampling_params=vllm_sp,
|
||||
vllm_gpu_memory_utilization=args.gpu_memory_utilization,
|
||||
**shared,
|
||||
)
|
||||
elif args.backend == "unsloth_fi_false":
|
||||
|
|
@ -283,16 +244,16 @@ def main():
|
|||
# fast_inference=False + for_inference() wires the fast single-token
|
||||
# decode + cached fp16 LoRA.
|
||||
training_args = GRPOConfig(
|
||||
use_vllm = False,
|
||||
bf16 = True,
|
||||
use_vllm=False,
|
||||
bf16=True,
|
||||
**shared,
|
||||
)
|
||||
elif args.backend in ("cb_paged", "cb_sdpa"):
|
||||
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,
|
||||
},
|
||||
|
|
@ -300,33 +261,57 @@ def main():
|
|||
)
|
||||
else: # naive_trl
|
||||
training_args = GRPOConfig(
|
||||
use_vllm = False,
|
||||
bf16 = True,
|
||||
use_vllm=False,
|
||||
bf16=True,
|
||||
**shared,
|
||||
)
|
||||
|
||||
stats_cb = _make_stats_callback()
|
||||
|
||||
trainer = GRPOTrainer(
|
||||
model = model,
|
||||
processing_class = tokenizer,
|
||||
reward_funcs = reward_funcs,
|
||||
args = training_args,
|
||||
train_dataset = dataset,
|
||||
callbacks = [stats_cb],
|
||||
model=model,
|
||||
processing_class=tokenizer,
|
||||
reward_funcs=reward_funcs,
|
||||
args=training_args,
|
||||
train_dataset=dataset,
|
||||
callbacks=[stats_cb],
|
||||
)
|
||||
|
||||
if args.persistent_cb and args.backend in ("cb_paged", "cb_sdpa"):
|
||||
from persistent_cb import install_for_model, teardown
|
||||
|
||||
base = (
|
||||
trainer.model_wrapped.base_model.model
|
||||
if hasattr(trainer.model_wrapped, "base_model")
|
||||
else trainer.model_wrapped
|
||||
)
|
||||
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)
|
||||
persistent_teardown_target = base
|
||||
|
||||
# Phase 4: torch.compile on the training forward.
|
||||
if args.compile_mode and args.backend != "vllm":
|
||||
from torch_debugging_utils import clear_inductor_cache, CompileDebugger
|
||||
clear_inductor_cache()
|
||||
CompileDebugger.enable(graph_breaks=True, recompiles=True)
|
||||
# Raise Dynamo cache limit so dynamic-shape recompiles don't thrash.
|
||||
import torch._dynamo
|
||||
torch._dynamo.config.cache_size_limit = 128
|
||||
try:
|
||||
torch._dynamo.config.allow_unspec_int_on_nn_module = True
|
||||
except AttributeError:
|
||||
pass
|
||||
print(f"[{args.backend}] Compiling trainer.model.forward "
|
||||
f"(mode={args.compile_mode}, dynamic={args.compile_dynamic})")
|
||||
trainer.model.forward = torch.compile(
|
||||
trainer.model.forward,
|
||||
mode=args.compile_mode,
|
||||
dynamic=args.compile_dynamic,
|
||||
)
|
||||
# Reference model inside TRL's GRPO loop also runs a forward.
|
||||
ref = getattr(trainer, "ref_model", None)
|
||||
if ref is not None:
|
||||
ref.forward = torch.compile(
|
||||
ref.forward, mode=args.compile_mode,
|
||||
dynamic=args.compile_dynamic,
|
||||
)
|
||||
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
t_start = time.perf_counter()
|
||||
try:
|
||||
|
|
@ -334,7 +319,6 @@ def main():
|
|||
finally:
|
||||
if persistent_teardown_target is not None:
|
||||
from persistent_cb import teardown
|
||||
|
||||
teardown(persistent_teardown_target)
|
||||
train_wall = time.perf_counter() - t_start
|
||||
|
||||
|
|
@ -374,17 +358,10 @@ def main():
|
|||
}
|
||||
summary_path = Path(args.stats_path).with_suffix(".summary.json")
|
||||
with open(summary_path, "w") as f:
|
||||
json.dump(summary, f, indent = 2)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
k: v
|
||||
for k, v in summary.items()
|
||||
if k not in ("losses", "rewards", "kls", "grad_norms", "step_times_ms")
|
||||
},
|
||||
indent = 2,
|
||||
)
|
||||
)
|
||||
json.dump(summary, f, indent=2)
|
||||
print(json.dumps({k: v for k, v in summary.items()
|
||||
if k not in ("losses", "rewards", "kls", "grad_norms", "step_times_ms")},
|
||||
indent=2))
|
||||
print(f"\n[{args.backend}] wrote summary to {summary_path}")
|
||||
# vLLM engine holds refs; fast-exit rather than wait for shutdown.
|
||||
os._exit(0)
|
||||
|
|
|
|||
60
scripts/benchmarks/results/lora_rollout_baselines.md
Normal file
60
scripts/benchmarks/results/lora_rollout_baselines.md
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
# Phase 1: rollout-only LoRA rank-32 microbenchmark
|
||||
|
||||
Every backend generates the same 32 prompts (DAPO-Math-17k, seed 3407) for
|
||||
`max_new_tokens=512` with equivalence sampling: `temperature=0.1, top_p=0.97,
|
||||
min_p=0.5, top_k=5`. 16-prompt warmup, 2 measured rounds, median wall reported.
|
||||
|
||||
All four backends load the same `outputs/lora_rank32_fresh` adapter (see
|
||||
`make_lora_adapter.py`). LoRA kernels are active on every decode step.
|
||||
|
||||
## Results (GPU B200, bf16, Qwen3-4B-Base + rank-32 LoRA)
|
||||
|
||||
| Backend | Median wall (s) | Decode tok/s | Prompt tok/s | Peak mem (GB) | % of vLLM |
|
||||
|----------------------|-----------------|--------------|--------------|---------------|-----------|
|
||||
| vLLM (fast_inference)| 3.30 | **4581** | 1467 | 156.2 | 100.0 % |
|
||||
| unsloth_fi_false | 25.54 | 641 | 190 | **15.8** | 14.0 % |
|
||||
| CB paged+FA4 (persistent) | 34.99 | 422 | 138 | 103.8 | 9.2 % |
|
||||
| CB sdpa_paged (persistent)| 34.07 | 434 | 142 | 111.9 | 9.5 % |
|
||||
|
||||
## Observations
|
||||
|
||||
1. **vLLM with LoRA is ~37% slower than vLLM without LoRA** (7224 → 4581 tok/s
|
||||
per the pre-LoRA PR table). The LoRA kernels cost real time even in vLLM.
|
||||
Still the gold standard by a wide margin.
|
||||
|
||||
2. **Unsloth `fast_inference=False` is the surprise**: **1.5× faster than CB**
|
||||
at **1/7th the peak memory**. The cached fp16 LoRA copies in
|
||||
`fast_linear_forward` and the Triton RMSNorm/RoPE paths dominate the CB
|
||||
baseline on this workload. It is a real practical middle ground — no vLLM
|
||||
dependency, low memory, and ~14% of vLLM's throughput.
|
||||
|
||||
3. **CB paged_attention (FA4 shim) and CB sdpa_paged are within noise**:
|
||||
422 vs 434 tok/s. At this scale the attention kernel is not the bottleneck;
|
||||
Python-side launch overhead on `_generation_step` dominates (confirmed by
|
||||
prior profile: ~16k `cuLaunchKernelEx` for 371 decoded tokens). CUDA graph
|
||||
replay (Phase 3) is the right lever.
|
||||
|
||||
4. **Unsloth `fi_false` reached max_new_tokens on every prompt** (`n_decoded =
|
||||
16384 = 32 × 512`) whereas vLLM / CB stopped some sequences on EOS
|
||||
(`~15000 decoded`). Equivalence sampling + greedy-ish settings means most
|
||||
completions are long, but the slight difference is worth noting when
|
||||
reading the raw tok/s numbers.
|
||||
|
||||
5. Completions are qualitatively coherent in every backend (see
|
||||
`sample_completions` in the stats JSONs). vLLM and unsloth_fi_false produce
|
||||
the *same* opening tokens on probe prompts (deterministic sampling lower
|
||||
bound), which is a useful weak sanity check.
|
||||
|
||||
## Raw stats
|
||||
|
||||
- `scripts/benchmarks/results/stats/lora_vllm_gen.json`
|
||||
- `scripts/benchmarks/results/stats/lora_unsloth_fi_false_gen.json`
|
||||
- `scripts/benchmarks/results/stats/lora_cb_paged_fa4_gen.json`
|
||||
- `scripts/benchmarks/results/stats/lora_cb_sdpa_paged_gen.json`
|
||||
|
||||
## Downstream implication
|
||||
|
||||
Phase 2 (full GRPO training) will include `unsloth_fi_false` as a first-class
|
||||
backend — if throughput parity holds end-to-end, it may be the pragmatic
|
||||
default for teams that cannot take the vLLM memory footprint. Phase 3 (CB sync
|
||||
driver + CUDA graphs) targets the CB paths specifically.
|
||||
29
scripts/benchmarks/results/stats/lora_cb_paged_fa4_gen.json
Normal file
29
scripts/benchmarks/results/stats/lora_cb_paged_fa4_gen.json
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
{
|
||||
"backend": "tpaged",
|
||||
"lora_adapter": "outputs/lora_rank32_fresh",
|
||||
"attn_impl": "paged_attention",
|
||||
"persistent_cb": true,
|
||||
"n_prompts": 32,
|
||||
"n_prompt_tokens": 4847,
|
||||
"n_decoded_tokens": 14750,
|
||||
"wall_times_s": [
|
||||
34.991439365025144,
|
||||
33.211731680028606
|
||||
],
|
||||
"median_wall_s": 34.991439365025144,
|
||||
"prompt_tps": 138.5195947339252,
|
||||
"decode_tps": 421.53167367967745,
|
||||
"max_new_tokens": 512,
|
||||
"sample_completions": [
|
||||
"First, we can factor the quadratic expression $n^2-3n+2$ as $(n-1)(n-2)$. For this expression to be a prime number, one of the factors must be equal to 1 and the other factor must be a prime number. \n",
|
||||
"First, we need to determine how many $4 \\times 5$ rectangles can fit into a $20 \\times 24$ rectangle. To do this, we divide the dimensions of the larger rectangle by the dimensions of the smaller rect",
|
||||
" \nTo find the area of the smaller square, we need to determine its side length. Let's denote the side length of the smaller square as \\( s \\).\n\nFrom the diagram, we can see that the larger square has "
|
||||
],
|
||||
"peak_memory_gb": 103.81339406967163,
|
||||
"sampling": {
|
||||
"temperature": 0.1,
|
||||
"top_p": 0.97,
|
||||
"min_p": 0.5,
|
||||
"top_k": 5
|
||||
}
|
||||
}
|
||||
29
scripts/benchmarks/results/stats/lora_cb_sdpa_paged_gen.json
Normal file
29
scripts/benchmarks/results/stats/lora_cb_sdpa_paged_gen.json
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
{
|
||||
"backend": "tpaged",
|
||||
"lora_adapter": "outputs/lora_rank32_fresh",
|
||||
"attn_impl": "sdpa_paged",
|
||||
"persistent_cb": true,
|
||||
"n_prompts": 32,
|
||||
"n_prompt_tokens": 4847,
|
||||
"n_decoded_tokens": 14785,
|
||||
"wall_times_s": [
|
||||
33.532237556006294,
|
||||
34.068173120962456
|
||||
],
|
||||
"median_wall_s": 34.068173120962456,
|
||||
"prompt_tps": 142.27355199793783,
|
||||
"decode_tps": 433.9827658942667,
|
||||
"max_new_tokens": 512,
|
||||
"sample_completions": [
|
||||
"First, we need to understand the structure of a cube. A cube has 12 edges and 8 vertices. Each vertex is connected to 3 edges. \n\nNow, let's consider the pairs of parallel edges. Since a cube has 12 ed",
|
||||
"First, we need to determine how many $4 \\times 5$ rectangles can fit into a $20 \\times 24$ rectangle. We can do this by dividing the dimensions of the larger rectangle by the dimensions of the smaller",
|
||||
"First, let's count the total number of letters in the word \"FLUFFY\". There are 6 letters in total.\n\nNext, we need to determine how many of these letters are repeated. In this case, the letter \"F\" appe"
|
||||
],
|
||||
"peak_memory_gb": 111.93839406967163,
|
||||
"sampling": {
|
||||
"temperature": 0.1,
|
||||
"top_p": 0.97,
|
||||
"min_p": 0.5,
|
||||
"top_k": 5
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
{
|
||||
"backend": "unsloth_fi_false",
|
||||
"lora_adapter": "outputs/lora_rank32_fresh",
|
||||
"n_prompts": 32,
|
||||
"n_prompt_tokens": 4847,
|
||||
"n_decoded_tokens": 16384,
|
||||
"wall_times_s": [
|
||||
25.54396249598358,
|
||||
25.480999241000973
|
||||
],
|
||||
"median_wall_s": 25.54396249598358,
|
||||
"prompt_tps": 189.75129644674828,
|
||||
"decode_tps": 641.4040109311994,
|
||||
"max_new_tokens": 512,
|
||||
"sample_completions": [
|
||||
"First, we need to find the length of the legs of the trapezoid. Since the trapezoid is isosceles, the legs are equal in length. Let's call the length of each leg $x$. We can use the Pythagorean theore",
|
||||
"Let $Q(x) = P(x) - x^{2023}P(1-\\frac{1}{x})$. Then $Q(k) = 0$ for every positive integer $1 \\leq k \\leq 2023$. Since $P(x)$ is a monic polynomial of degree $2023$, $Q(x)$ is also a monic polynomial of",
|
||||
" To solve this problem, we need to determine the maximum value of \\(a\\) such that the line \\(y = mx + 2\\) does not pass through any lattice points for \\(0 < x \\leq 100\\) when \\(\\frac{1}{2} < m < a\\).\n"
|
||||
],
|
||||
"peak_memory_gb": 15.8363037109375,
|
||||
"sampling": {
|
||||
"temperature": 0.1,
|
||||
"top_p": 0.97,
|
||||
"min_p": 0.5,
|
||||
"top_k": 5
|
||||
}
|
||||
}
|
||||
27
scripts/benchmarks/results/stats/lora_vllm_gen.json
Normal file
27
scripts/benchmarks/results/stats/lora_vllm_gen.json
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
{
|
||||
"backend": "vllm",
|
||||
"lora_adapter": "outputs/lora_rank32_fresh",
|
||||
"n_prompts": 32,
|
||||
"n_prompt_tokens": 4847,
|
||||
"n_decoded_tokens": 15140,
|
||||
"wall_times_s": [
|
||||
3.304712440993171,
|
||||
3.2573561430326663
|
||||
],
|
||||
"median_wall_s": 3.304712440993171,
|
||||
"prompt_tps": 1466.6934223612275,
|
||||
"decode_tps": 4581.336582329066,
|
||||
"max_new_tokens": 512,
|
||||
"sample_completions": [
|
||||
"First, we need to find the length of the legs of the trapezoid. Since the trapezoid is isosceles, the legs are equal in length. Let's call the length of each leg $x$. We can use the Pythagorean theore",
|
||||
"Let $Q(x) = P(x) - x^{2023}P(1-\\frac{1}{x})$. Then $Q(k) = 0$ for every positive integer $1 \\leq k \\leq 2023$. Since $P(x)$ is a monic polynomial of degree $2023$, $Q(x)$ is also a monic polynomial of",
|
||||
" To solve this problem, we need to find the maximum value of \\(a\\) such that the line \\(y = mx + 2\\) does not pass through any lattice points for \\(0 < x \\leq 100\\) when \\(\\frac{1}{2} < m < a\\).\n\nFirs"
|
||||
],
|
||||
"peak_memory_gb": 156.21798133850098,
|
||||
"sampling": {
|
||||
"temperature": 0.1,
|
||||
"top_p": 0.97,
|
||||
"min_p": 0.5,
|
||||
"top_k": 5
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue