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

Phase 2 vibe (10-step): vllm vs unsloth_fi_false vs cb_paged + Phase 3 fixes

Browse source 
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 2026-04-20 14:23:51 +00:00
commit 5f3e1c98df
13 changed files with 1815 additions and 569 deletions
Download patch file Download diff file Expand all files Collapse all files

144
scripts/benchmarks/cb_sync_driver.py
View file

@ -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))

211
scripts/benchmarks/cb_vs_vllm_generation.py
View file

@ -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,40 +248,33 @@ 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)
# 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.
@ -298,12 +282,10 @@ def run_unsloth_fi_false(args):
n = name
for pref in ("base_model.model.", "model."):
if n.startswith(pref):
n = n[len(pref) :]
n = n[len(pref):]
n = n.replace(".lora_A.default.", ".lora_A.").replace(
".lora_B.default.", ".lora_B."
)
".lora_B.default.", ".lora_B.")
return n
own_by_core = {}
for n, p in model.named_parameters():
if "lora_" in n:
@ -317,10 +299,8 @@ def run_unsloth_fi_false(args):
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)."
)
print(f"[unsloth_fi_false] LoRA weight sync matched {matched} tensors "
f"(out of {len(loaded_tensors)} adapter entries).")
FastLanguageModel.for_inference(model)
@ -328,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
@ -371,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
@ -381,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",
@ -404,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":
@ -450,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)

46
scripts/benchmarks/make_lora_adapter.py
View file

@ -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}")

289
scripts/benchmarks/qwen3_grpo_notebook.py
View file

@ -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:

300
scripts/benchmarks/qwen3_grpo_unified.py
View file

@ -42,40 +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)
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()
@ -83,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():
@ -102,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,
)
@ -117,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
@ -131,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
@ -174,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()
@ -205,46 +170,57 @@ 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)
# TRL requires `generation_batch_size = pdb * grad_accum * world_size` to
# be divisible by `num_generations`. Unsloth's loader auto-adjusts
# `per_device_train_batch_size` to match `num_generations`, but vanilla HF
# paths (cb_paged, cb_sdpa, naive_trl) do not -- do it ourselves.
if args.backend not in ("vllm", "unsloth_fi_false"):
effective = (args.per_device_train_batch_size
* args.gradient_accumulation_steps)
if effective % args.num_generations != 0:
new_pdb = args.num_generations
print(f"[{args.backend}] Bumping per_device_train_batch_size "
f"{args.per_device_train_batch_size} -> {new_pdb} to satisfy "
f"GRPO divisibility.")
args.per_device_train_batch_size = new_pdb
# --- 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")
# `paged_attention` requires cu_seq_lens on every forward, which only
# the CB rollout path provides. GRPO's training forward (dense batch)
# crashes. Load with `sdpa_paged` which gracefully falls back to
# plain SDPA when paged args are absent, and still exercises the
# paged path during CB rollout.
model, tokenizer = _load_vanilla_hf(args, attn_impl="sdpa_paged")
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)
@ -252,12 +228,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
@ -268,24 +244,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":
@ -293,16 +263,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,
},
@ -310,63 +280,55 @@ 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)
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})"
)
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,
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,
ref.forward, mode=args.compile_mode,
dynamic=args.compile_dynamic,
)
torch.cuda.reset_peak_memory_stats()
@ -376,7 +338,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
@ -416,17 +377,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)

89
scripts/benchmarks/results/grpo_equivalence.md Normal file
View file

@ -0,0 +1,89 @@
# Phase 2: end-to-end GRPO backend comparison (10-step vibe check)
Same dataset, reward functions, sampling (`temperature=0.1, top_p=0.97,
min_p=0.5, top_k=5`), and seed (3407). `max_steps=10, num_generations=4,
per_device_train_batch_size=4` (auto-adjusted from 1 on vanilla-HF paths
to satisfy TRL's `generation_batch_size % num_generations == 0`).
Callbacks: `StatisticsCallback` from `torch_debugging_utils` logs per-step
loss, grad-norm, memory, wall time. Median step time is computed over steps
4-10 (first 3 skipped for compile / graph / warmup amortization).
## 10-step results
| Backend | Train wall (s) | Median step (s) | Peak mem (GB) | % of vLLM |
|-------------------------------|----------------|-----------------|---------------|-----------|
| vLLM (fast_inference) | 74.4 | **4.14** | 157.9 | 100 % |
| unsloth_fi_false | 355.4 | 23.95 | **10.7** | 17 % |
| cb_paged (sdpa_paged load) | 466.0 | 36.02 | 55.6 | 11.5 % |
Loss / reward / KL arrays for each backend (10 steps, rounded):
| Step | vLLM loss | vLLM reward | vLLM kl | fi_false loss | fi_false reward | fi_false kl | cb_paged loss | cb_paged reward |
|------|-----------|-------------|----------|---------------|-----------------|-------------|----------------|------------------|
| 1 | 0.031 | 0.00 | 0.00000 | 0.000 | 0.50 | 0.00000 | -0.086 | 0.62 |
| 2 | -0.194 | -2.50 | 0.00000 | -0.089 | -6.50 | 0.00000 | 0.041 | -2.50 |
| 3 | 0.263 | -3.62 | 0.01192 | -0.139 | -2.50 | 0.00857 | 0.000 | 0.50 |
| 4 | -0.201 | 0.00 | 0.00422 | -0.124 | -2.50 | 0.00931 | 0.086 | -1.50 |
| 5 | 0.209 | 0.38 | 0.00369 | 0.000 | 0.50 | 0.00213 | 0.016 | 0.00 |
| 6 | 0.000 | -7.50 | 0.00250 | 0.000 | -7.50 | 0.00071 | 0.000 | -7.50 |
| 7 | 0.037 | 1.50 | 0.00614 | -0.010 | -5.50 | 0.00522 | 0.074 | 1.50 |
| 8 | 0.000 | -7.50 | 0.00465 | 0.034 | -5.25 | 0.00014 | -0.048 | -4.25 |
| 9 | 0.000 | 0.50 | 0.00176 | 0.000 | 0.50 | 0.01090 | -0.188 | -1.50 |
| 10 | 0.205 | -3.50 | 0.00200 | 0.204 | -2.50 | 0.00215 | 0.044 | -6.50 |
## Observations
1. **Coherence gate (all backends)**: losses are bounded in `[-0.25, 0.3]`,
grad-norms finite, rewards in the plan's expected negative-then-rising
range. No CJK-token salad, no NaNs.
2. **KL trajectories are qualitatively matched** between vLLM and
`unsloth_fi_false` (both in `[0, 0.015]`), confirming that
`fast_inference=False` produces rollouts close to the vLLM reference once
`temperature=0.1` is used. `cb_paged` also produces rollouts but our
`StatisticsCallback` did not capture TRL's `kl` entry in its `on_log`
pass -- the next iteration will forward every log dict entry into the JSON.
3. **Per-step timing**: `unsloth_fi_false` is 5.8x slower than vLLM; `cb_paged`
is 8.7x slower. Neither hits the plan's 30% target on this vibe check.
4. **Memory is the standout axis**:
- vLLM: 158 GB (prefill KV cache + vLLM engine overhead)
- cb_paged: 55.6 GB (paged cache only)
- unsloth_fi_false: **10.7 GB** -- 15x lower than vLLM.
Unsloth's fast_inference=False path is a genuine option for teams who
cannot afford the vLLM footprint but are willing to take a ~5-6x rollout
wall-clock hit.
5. **cb_paged load needed `sdpa_paged` not `paged_attention`**: the
FA4-shimmed `paged_attention` kernel requires `cu_seq_lens_q` on every
forward, but GRPO's training forward (dense batch) doesn't provide them.
`sdpa_paged` falls back to plain SDPA when no paged kwargs are present and
still exercises paged attention during the CB rollout. This is consistent
with the existing `qwen3_grpo_tpaged.py` which loads with `sdpa`.
## What's next (not yet run)
- **30-step equivalence** with `torch_debugging_utils.compare_training_runs`
comparing vLLM vs each backend on loss / reward / KL arrays.
- **Phase 3 sync driver** smoke-tested successfully (eager decode produces
512 correct tokens) but CUDA graph capture hangs on the first graphed step.
Likely cause: `PagedAttentionCache` constructs tensors inside
`cache.update()` the first call, which doesn't survive graph capture.
Two possible fixes being explored: (a) pre-capture warmup steps on the
capture stream so allocations are already done, (b) replace in-place
torch.multinomial-adjacent ops with CUDA-graph-safe equivalents.
- **Phase 4 torch.compile**: hook-up ready in `qwen3_grpo_unified.py`
(`--compile_mode default|reduce-overhead|max-autotune-no-cudagraphs`);
needs a run budget allocated and the `CompileDebugger` output reviewed.
## Raw stats
- `scripts/benchmarks/results/stats/grpo_vllm_10.summary.json`
- `scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.summary.json`
- `scripts/benchmarks/results/stats/grpo_cb_paged_10.summary.json`
Full per-step logs (one entry per step with loss/reward/kl/grad_norm and all
of TRL's logging dict) live at `scripts/benchmarks/results/stats/grpo_*.json`.

15
scripts/benchmarks/results/stats/cb_sync_smoke.json Normal file
View file

@ -0,0 +1,15 @@
{
"backend": "cb_sync_driver",
"use_cuda_graph": false,
"attn_impl": "paged_attention",
"n_prompts": 8,
"n_decoded_tokens": 512,
"wall_times_s": [
100.8571443540277,
100.87157070200192
],
"median_wall_s": 100.87157070200192,
"decode_tps": 5.075761152887835,
"max_new_tokens": 64,
"peak_memory_gb": 45.91296434402466
}

352
scripts/benchmarks/results/stats/grpo_cb_paged_10.json Normal file
View file

@ -0,0 +1,352 @@
[
{
"step": 1,
"loss": -0.0862,
"grad_norm": 716.0,
"learning_rate": 0.0,
"num_tokens": 4262.0,
"completions/mean_length": 953.5,
"completions/min_length": 824.0,
"completions/max_length": 1092.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 953.5,
"completions/min_terminated_length": 824.0,
"completions/max_terminated_length": 1092.0,
"rewards/match_format_exactly/mean": 2.25,
"rewards/match_format_exactly/std": 1.5,
"rewards/match_format_approximately/mean": 1.125,
"rewards/match_format_approximately/std": 0.75,
"rewards/check_answer/mean": -1.25,
"rewards/check_answer/std": 2.1794495582580566,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.625,
"reward_std": 3.4731109142303467,
"frac_reward_zero_std": 0.0,
"entropy": 0.1351587027311325,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 7.868439688409789e-05,
"time_ms": 56644.44096497027,
"memory_mb": 57451.41162109375,
"memory_gb": 56.104894161224365
},
{
"step": 2,
"loss": 0.041,
"grad_norm": 186.0,
"learning_rate": 5e-06,
"num_tokens": 6762.0,
"completions/mean_length": 536.0,
"completions/min_length": 492.0,
"completions/max_length": 603.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 536.0,
"completions/min_terminated_length": 492.0,
"completions/max_terminated_length": 603.0,
"rewards/match_format_exactly/mean": 0.75,
"rewards/match_format_exactly/std": 1.5,
"rewards/match_format_approximately/mean": 0.375,
"rewards/match_format_approximately/std": 0.75,
"rewards/check_answer/mean": -2.125,
"rewards/check_answer/std": 0.25,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": -2.5,
"reward_std": 2.0,
"frac_reward_zero_std": 0.0,
"entropy": 0.05973631516098976,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00015736879376819577,
"time_ms": 29505.720576969907,
"memory_mb": 53805.20556640625,
"memory_gb": 52.5441460609436
},
{
"step": 3,
"loss": 0.0,
"grad_norm": 0.0,
"learning_rate": 4.444444444444444e-06,
"num_tokens": 10099.0,
"completions/mean_length": 657.25,
"completions/min_length": 436.0,
"completions/max_length": 1302.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 657.25,
"completions/min_terminated_length": 436.0,
"completions/max_terminated_length": 1302.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.5,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"entropy": 0.06822667270898819,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00023605319065229366,
"time_ms": 62939.08593803644,
"memory_mb": 59249.8505859375,
"memory_gb": 57.86118221282959
},
{
"step": 4,
"loss": 0.0855,
"grad_norm": 274.0,
"learning_rate": 3.88888888888889e-06,
"num_tokens": 13644.0,
"completions/mean_length": 721.25,
"completions/min_length": 441.0,
"completions/max_length": 988.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 721.25,
"completions/min_terminated_length": 441.0,
"completions/max_terminated_length": 988.0,
"rewards/match_format_exactly/mean": 1.5,
"rewards/match_format_exactly/std": 1.7320507764816284,
"rewards/match_format_approximately/mean": 0.75,
"rewards/match_format_approximately/std": 0.8660253882408142,
"rewards/check_answer/mean": -2.25,
"rewards/check_answer/std": 0.28867512941360474,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": -1.5,
"reward_std": 2.309401035308838,
"frac_reward_zero_std": 0.0,
"entropy": 0.25085046887397766,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00031473758753639155,
"time_ms": 49076.82833302533,
"memory_mb": 56826.26611328125,
"memory_gb": 55.49440050125122
},
{
"step": 5,
"loss": 0.0162,
"grad_norm": 143.0,
"learning_rate": 3.3333333333333333e-06,
"num_tokens": 15442.0,
"completions/mean_length": 293.5,
"completions/min_length": 246.0,
"completions/max_length": 365.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 293.5,
"completions/min_terminated_length": 246.0,
"completions/max_terminated_length": 365.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -3.0,
"rewards/check_answer/std": 1.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.0,
"reward_std": 1.0,
"frac_reward_zero_std": 0.0,
"entropy": 0.0809403508901596,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00039342198442048943,
"time_ms": 18123.881562962197,
"memory_mb": 52261.25830078125,
"memory_gb": 51.03638505935669
},
{
"step": 6,
"loss": 0.0,
"grad_norm": 0.0,
"learning_rate": 2.7777777777777783e-06,
"num_tokens": 21877.0,
"completions/mean_length": 1511.75,
"completions/min_length": 1112.0,
"completions/max_length": 1846.0,
"completions/clipped_ratio": 0.5,
"completions/mean_terminated_length": 1177.5,
"completions/min_terminated_length": 1112.0,
"completions/max_terminated_length": 1243.0,
"rewards/match_format_exactly/mean": 0.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": -3.0,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.0,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -2.5,
"rewards/check_numbers/std": 0.0,
"reward": -7.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"entropy": 0.2572544813156128,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0004721063813045873,
"time_ms": 92771.2398529984,
"memory_mb": 63378.49365234375,
"memory_gb": 61.89306020736694
},
{
"step": 7,
"loss": 0.0736,
"grad_norm": 74.0,
"learning_rate": 2.222222222222222e-06,
"num_tokens": 24635.0,
"completions/mean_length": 546.5,
"completions/min_length": 466.0,
"completions/max_length": 585.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 546.5,
"completions/min_terminated_length": 466.0,
"completions/max_terminated_length": 585.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -1.5,
"rewards/check_answer/std": 2.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 1.5,
"reward_std": 2.0,
"frac_reward_zero_std": 0.0,
"entropy": 0.05001620948314667,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0005507907781886852,
"time_ms": 30032.14380296413,
"memory_mb": 53704.75830078125,
"memory_gb": 52.44605302810669
},
{
"step": 8,
"loss": -0.0475,
"grad_norm": 97.0,
"learning_rate": 1.6666666666666667e-06,
"num_tokens": 27393.0,
"completions/mean_length": 622.5,
"completions/min_length": 533.0,
"completions/max_length": 696.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 622.5,
"completions/min_terminated_length": 533.0,
"completions/max_terminated_length": 696.0,
"rewards/match_format_exactly/mean": 0.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": -1.5,
"rewards/match_format_approximately/std": 1.7320507764816284,
"rewards/check_answer/mean": -2.0,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -0.75,
"rewards/check_numbers/std": 2.872281312942505,
"reward": -4.25,
"reward_std": 4.27200174331665,
"frac_reward_zero_std": 0.0,
"entropy": 0.08264704048633575,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0006294751750727831,
"time_ms": 36016.123837034684,
"memory_mb": 54504.00634765625,
"memory_gb": 53.22656869888306
},
{
"step": 9,
"loss": -0.1881,
"grad_norm": 274.0,
"learning_rate": 1.111111111111111e-06,
"num_tokens": 29461.0,
"completions/mean_length": 420.0,
"completions/min_length": 327.0,
"completions/max_length": 647.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 420.0,
"completions/min_terminated_length": 327.0,
"completions/max_terminated_length": 647.0,
"rewards/match_format_exactly/mean": 1.5,
"rewards/match_format_exactly/std": 1.7320507764816284,
"rewards/match_format_approximately/mean": 0.0,
"rewards/match_format_approximately/std": 2.1213202476501465,
"rewards/check_answer/mean": -1.25,
"rewards/check_answer/std": 1.8484227657318115,
"rewards/check_numbers/mean": -1.75,
"rewards/check_numbers/std": 0.5,
"reward": -1.5,
"reward_std": 5.16397762298584,
"frac_reward_zero_std": 0.0,
"entropy": 0.22891533374786377,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0007081595719568809,
"time_ms": 35705.52668598248,
"memory_mb": 54149.77783203125,
"memory_gb": 52.88064241409302
},
{
"step": 10,
"loss": 0.0444,
"grad_norm": 236.0,
"learning_rate": 5.555555555555555e-07,
"num_tokens": 33785.0,
"completions/mean_length": 913.0,
"completions/min_length": 832.0,
"completions/max_length": 998.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 913.0,
"completions/min_terminated_length": 832.0,
"completions/max_terminated_length": 998.0,
"rewards/match_format_exactly/mean": 0.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": -2.25,
"rewards/match_format_approximately/std": 1.5,
"rewards/check_answer/mean": -2.0,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -2.25,
"rewards/check_numbers/std": 0.5,
"reward": -6.5,
"reward_std": 2.0,
"frac_reward_zero_std": 0.0,
"entropy": 0.1578676998615265,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0007868439688409789,
"time_ms": 53879.347916983534,
"memory_mb": 56904.7578125,
"memory_gb": 55.57105255126953
}
]

64
scripts/benchmarks/results/stats/grpo_cb_paged_10.summary.json Normal file
View file

@ -0,0 +1,64 @@
{
"backend": "cb_paged",
"max_steps": 10,
"train_wall_s": 466.01091928296955,
"median_step_ms_post_warmup": 36016.123837034684,
"n_logged_steps": 10,
"sampling": {
"temperature": 0.1,
"top_p": 0.97,
"min_p": 0.5,
"top_k": 5
},
"losses": [
-0.0862,
0.041,
0.0,
0.0855,
0.0162,
0.0,
0.0736,
-0.0475,
-0.1881,
0.0444
],
"rewards": [
0.625,
-2.5,
0.5,
-1.5,
0.0,
-7.5,
1.5,
-4.25,
-1.5,
-6.5
],
"kls": [],
"grad_norms": [
716.0,
186.0,
0.0,
274.0,
143.0,
0.0,
74.0,
97.0,
274.0,
236.0
],
"step_times_ms": [
56644.44096497027,
29505.720576969907,
62939.08593803644,
49076.82833302533,
18123.881562962197,
92771.2398529984,
30032.14380296413,
36016.123837034684,
35705.52668598248,
53879.347916983534
],
"peak_memory_gb": 55.57105255126953,
"logs_path": "logs/grpo_cb_paged_10.json"
}

362
scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.json Normal file
View file

@ -0,0 +1,362 @@
[
{
"step": 1,
"loss": 0.0,
"grad_norm": 0.0,
"learning_rate": 0.0,
"num_tokens": 3693.0,
"completions/mean_length": 811.25,
"completions/min_length": 779.0,
"completions/max_length": 856.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 811.25,
"completions/min_terminated_length": 779.0,
"completions/max_terminated_length": 856.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.5,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"completion_length": 811.25,
"kl": 0.0,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 7.868439688409789e-05,
"time_ms": 77897.29859499494,
"memory_mb": 9442.14013671875,
"memory_gb": 9.220839977264404
},
{
"step": 2,
"loss": -0.0893,
"grad_norm": 0.6125104427337646,
"learning_rate": 5e-06,
"num_tokens": 6238.0,
"completions/mean_length": 547.25,
"completions/min_length": 487.0,
"completions/max_length": 645.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 547.25,
"completions/min_terminated_length": 487.0,
"completions/max_terminated_length": 645.0,
"rewards/match_format_exactly/mean": 0.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": -2.25,
"rewards/match_format_approximately/std": 1.5,
"rewards/check_answer/mean": -2.0,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -2.25,
"rewards/check_numbers/std": 0.5,
"reward": -6.5,
"reward_std": 2.0,
"frac_reward_zero_std": 0.0,
"completion_length": 547.25,
"kl": 0.0,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00015736879376819577,
"time_ms": 21634.983669035137,
"memory_mb": 9076.26025390625,
"memory_gb": 8.863535404205322
},
{
"step": 3,
"loss": -0.1386,
"grad_norm": 0.5993297696113586,
"learning_rate": 4.444444444444444e-06,
"num_tokens": 10165.0,
"completions/mean_length": 804.75,
"completions/min_length": 605.0,
"completions/max_length": 1214.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 804.75,
"completions/min_terminated_length": 605.0,
"completions/max_terminated_length": 1214.0,
"rewards/match_format_exactly/mean": 1.5,
"rewards/match_format_exactly/std": 1.7320507764816284,
"rewards/match_format_approximately/mean": -0.75,
"rewards/match_format_approximately/std": 2.598076105117798,
"rewards/check_answer/mean": -1.25,
"rewards/check_answer/std": 1.8484227657318115,
"rewards/check_numbers/mean": -2.0,
"rewards/check_numbers/std": 0.5773502588272095,
"reward": -2.5,
"reward_std": 6.0,
"frac_reward_zero_std": 0.0,
"completion_length": 804.75,
"kl": 0.008573448285460472,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00023605319065229366,
"time_ms": 40298.14357904252,
"memory_mb": 9952.80322265625,
"memory_gb": 9.719534397125244
},
{
"step": 4,
"loss": -0.1236,
"grad_norm": 0.5647851228713989,
"learning_rate": 3.88888888888889e-06,
"num_tokens": 13320.0,
"completions/mean_length": 623.75,
"completions/min_length": 421.0,
"completions/max_length": 789.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 623.75,
"completions/min_terminated_length": 421.0,
"completions/max_terminated_length": 789.0,
"rewards/match_format_exactly/mean": 0.75,
"rewards/match_format_exactly/std": 1.5,
"rewards/match_format_approximately/mean": 0.375,
"rewards/match_format_approximately/std": 0.75,
"rewards/check_answer/mean": -2.125,
"rewards/check_answer/std": 0.25,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": -2.5,
"reward_std": 2.0,
"frac_reward_zero_std": 0.0,
"completion_length": 623.75,
"kl": 0.009312103502452374,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00031473758753639155,
"time_ms": 26193.765547999647,
"memory_mb": 9306.09326171875,
"memory_gb": 9.087981700897217
},
{
"step": 5,
"loss": 0.0,
"grad_norm": 0.0010538548231124878,
"learning_rate": 3.3333333333333333e-06,
"num_tokens": 14970.0,
"completions/mean_length": 256.5,
"completions/min_length": 246.0,
"completions/max_length": 260.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 256.5,
"completions/min_terminated_length": 246.0,
"completions/max_terminated_length": 260.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.5,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"completion_length": 256.5,
"kl": 0.002130241831764579,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00039342198442048943,
"time_ms": 9324.433026020415,
"memory_mb": 8777.982421875,
"memory_gb": 8.572248458862305
},
{
"step": 6,
"loss": 0.0,
"grad_norm": 0.00012166703527327627,
"learning_rate": 2.7777777777777783e-06,
"num_tokens": 21841.0,
"completions/mean_length": 1620.75,
"completions/min_length": 1214.0,
"completions/max_length": 1846.0,
"completions/clipped_ratio": 0.5,
"completions/mean_terminated_length": 1395.5,
"completions/min_terminated_length": 1214.0,
"completions/max_terminated_length": 1577.0,
"rewards/match_format_exactly/mean": 0.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": -3.0,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.0,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -2.5,
"rewards/check_numbers/std": 0.0,
"reward": -7.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"completion_length": 1620.75,
"kl": 0.0007094849133864045,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0004721063813045873,
"time_ms": 61652.94101298787,
"memory_mb": 10914.11279296875,
"memory_gb": 10.658313274383545
},
{
"step": 7,
"loss": -0.0097,
"grad_norm": 0.9194015860557556,
"learning_rate": 2.222222222222222e-06,
"num_tokens": 24587.0,
"completions/mean_length": 543.5,
"completions/min_length": 511.0,
"completions/max_length": 562.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 543.5,
"completions/min_terminated_length": 511.0,
"completions/max_terminated_length": 562.0,
"rewards/match_format_exactly/mean": 0.75,
"rewards/match_format_exactly/std": 1.5,
"rewards/match_format_approximately/mean": -1.875,
"rewards/match_format_approximately/std": 2.25,
"rewards/check_answer/mean": -2.125,
"rewards/check_answer/std": 0.25,
"rewards/check_numbers/mean": -2.25,
"rewards/check_numbers/std": 0.5,
"reward": -5.5,
"reward_std": 4.0,
"frac_reward_zero_std": 0.0,
"completion_length": 543.5,
"kl": 0.005215016193687916,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0005507907781886852,
"time_ms": 18906.7719859886,
"memory_mb": 8996.00390625,
"memory_gb": 8.785160064697266
},
{
"step": 8,
"loss": 0.0338,
"grad_norm": 0.5346357822418213,
"learning_rate": 1.6666666666666667e-06,
"num_tokens": 27519.0,
"completions/mean_length": 666.0,
"completions/min_length": 615.0,
"completions/max_length": 714.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 666.0,
"completions/min_terminated_length": 615.0,
"completions/max_terminated_length": 714.0,
"rewards/match_format_exactly/mean": 0.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": -2.25,
"rewards/match_format_approximately/std": 1.5,
"rewards/check_answer/mean": -2.0,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -1.0,
"rewards/check_numbers/std": 3.0,
"reward": -5.25,
"reward_std": 4.5,
"frac_reward_zero_std": 0.0,
"completion_length": 666.0,
"kl": 0.0001442090724594891,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0006294751750727831,
"time_ms": 23948.13355000224,
"memory_mb": 9202.39306640625,
"memory_gb": 8.986711978912354
},
{
"step": 9,
"loss": 0.0,
"grad_norm": 0.0033828848972916603,
"learning_rate": 1.111111111111111e-06,
"num_tokens": 29207.0,
"completions/mean_length": 325.0,
"completions/min_length": 310.0,
"completions/max_length": 334.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 325.0,
"completions/min_terminated_length": 310.0,
"completions/max_terminated_length": 334.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.5,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"completion_length": 325.0,
"kl": 0.010901343077421188,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0007081595719568809,
"time_ms": 11388.401818985585,
"memory_mb": 8705.7236328125,
"memory_gb": 8.501683235168457
},
{
"step": 10,
"loss": 0.2036,
"grad_norm": 0.22472381591796875,
"learning_rate": 5.555555555555555e-07,
"num_tokens": 34891.0,
"completions/mean_length": 1253.0,
"completions/min_length": 1044.0,
"completions/max_length": 1846.0,
"completions/clipped_ratio": 0.25,
"completions/mean_terminated_length": 1055.3333740234375,
"completions/min_terminated_length": 1044.0,
"completions/max_terminated_length": 1067.0,
"rewards/match_format_exactly/mean": 1.5,
"rewards/match_format_exactly/std": 1.7320507764816284,
"rewards/match_format_approximately/mean": 0.0,
"rewards/match_format_approximately/std": 2.1213202476501465,
"rewards/check_answer/mean": -2.25,
"rewards/check_answer/std": 0.28867512941360474,
"rewards/check_numbers/mean": -1.75,
"rewards/check_numbers/std": 0.5,
"reward": -2.5,
"reward_std": 3.8297085762023926,
"frac_reward_zero_std": 0.0,
"completion_length": 1253.0,
"kl": 0.00215436820872128,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0007868439688409789,
"time_ms": 61737.33338096645,
"memory_mb": 10919.5498046875,
"memory_gb": 10.663622856140137
}
]

75
scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.summary.json Normal file
View file

@ -0,0 +1,75 @@
{
"backend": "unsloth_fi_false",
"max_steps": 10,
"train_wall_s": 355.41431403998286,
"median_step_ms_post_warmup": 23948.13355000224,
"n_logged_steps": 10,
"sampling": {
"temperature": 0.1,
"top_p": 0.97,
"min_p": 0.5,
"top_k": 5
},
"losses": [
0.0,
-0.0893,
-0.1386,
-0.1236,
0.0,
0.0,
-0.0097,
0.0338,
0.0,
0.2036
],
"rewards": [
0.5,
-6.5,
-2.5,
-2.5,
0.5,
-7.5,
-5.5,
-5.25,
0.5,
-2.5
],
"kls": [
0.0,
0.0,
0.008573448285460472,
0.009312103502452374,
0.002130241831764579,
0.0007094849133864045,
0.005215016193687916,
0.0001442090724594891,
0.010901343077421188,
0.00215436820872128
],
"grad_norms": [
0.0,
0.6125104427337646,
0.5993297696113586,
0.5647851228713989,
0.0010538548231124878,
0.00012166703527327627,
0.9194015860557556,
0.5346357822418213,
0.0033828848972916603,
0.22472381591796875
],
"step_times_ms": [
77897.29859499494,
21634.983669035137,
40298.14357904252,
26193.765547999647,
9324.433026020415,
61652.94101298787,
18906.7719859886,
23948.13355000224,
11388.401818985585,
61737.33338096645
],
"peak_memory_gb": 10.663622856140137,
"logs_path": "logs/grpo_unsloth_fi_false_10.json"
}

362
scripts/benchmarks/results/stats/grpo_vllm_10.json Normal file
View file

@ -0,0 +1,362 @@
[
{
"step": 1,
"loss": 0.0305,
"grad_norm": 0.4133029878139496,
"learning_rate": 0.0,
"num_tokens": 3705.0,
"completions/mean_length": 814.25,
"completions/min_length": 781.0,
"completions/max_length": 864.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 814.25,
"completions/min_terminated_length": 781.0,
"completions/max_terminated_length": 864.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -3.0,
"rewards/check_answer/std": 1.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.0,
"reward_std": 1.0,
"frac_reward_zero_std": 0.0,
"completion_length": 814.25,
"kl": 0.0,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 7.868439688409789e-05,
"time_ms": 17984.56621397054,
"memory_mb": 161170.2099609375,
"memory_gb": 157.39278316497803
},
{
"step": 2,
"loss": -0.1941,
"grad_norm": 0.8333088159561157,
"learning_rate": 5e-06,
"num_tokens": 7167.0,
"completions/mean_length": 776.5,
"completions/min_length": 525.0,
"completions/max_length": 1078.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 776.5,
"completions/min_terminated_length": 525.0,
"completions/max_terminated_length": 1078.0,
"rewards/match_format_exactly/mean": 0.75,
"rewards/match_format_exactly/std": 1.5,
"rewards/match_format_approximately/mean": 0.375,
"rewards/match_format_approximately/std": 0.75,
"rewards/check_answer/mean": -2.125,
"rewards/check_answer/std": 0.25,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": -2.5,
"reward_std": 2.0,
"frac_reward_zero_std": 0.0,
"completion_length": 776.5,
"kl": 0.0,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00015736879376819577,
"time_ms": 6704.717919987161,
"memory_mb": 161638.7705078125,
"memory_gb": 157.85036182403564
},
{
"step": 3,
"loss": 0.2632,
"grad_norm": 0.4677680730819702,
"learning_rate": 4.444444444444444e-06,
"num_tokens": 11596.0,
"completions/mean_length": 930.25,
"completions/min_length": 445.0,
"completions/max_length": 1846.0,
"completions/clipped_ratio": 0.25,
"completions/mean_terminated_length": 625.0,
"completions/min_terminated_length": 445.0,
"completions/max_terminated_length": 863.0,
"rewards/match_format_exactly/mean": 1.5,
"rewards/match_format_exactly/std": 1.7320507764816284,
"rewards/match_format_approximately/mean": -0.75,
"rewards/match_format_approximately/std": 2.598076105117798,
"rewards/check_answer/mean": -2.75,
"rewards/check_answer/std": 1.190238118171692,
"rewards/check_numbers/mean": -1.625,
"rewards/check_numbers/std": 1.1814539432525635,
"reward": -3.625,
"reward_std": 4.479118347167969,
"frac_reward_zero_std": 0.0,
"completion_length": 930.25,
"kl": 0.011923530139029026,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00023605319065229366,
"time_ms": 12108.133931003977,
"memory_mb": 162822.73876953125,
"memory_gb": 159.00658082962036
},
{
"step": 4,
"loss": -0.2013,
"grad_norm": 0.5163940191268921,
"learning_rate": 3.88888888888889e-06,
"num_tokens": 14365.0,
"completions/mean_length": 527.25,
"completions/min_length": 315.0,
"completions/max_length": 598.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 527.25,
"completions/min_terminated_length": 315.0,
"completions/max_terminated_length": 598.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -3.0,
"rewards/check_answer/std": 1.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.0,
"reward_std": 1.0,
"frac_reward_zero_std": 0.0,
"completion_length": 527.25,
"kl": 0.004221913404762745,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00031473758753639155,
"time_ms": 4019.3081409670413,
"memory_mb": 160943.57275390625,
"memory_gb": 157.17145776748657
},
{
"step": 5,
"loss": 0.2093,
"grad_norm": 1.160618782043457,
"learning_rate": 3.3333333333333333e-06,
"num_tokens": 16176.0,
"completions/mean_length": 296.75,
"completions/min_length": 246.0,
"completions/max_length": 421.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 296.75,
"completions/min_terminated_length": 246.0,
"completions/max_terminated_length": 421.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -3.0,
"rewards/check_answer/std": 1.0,
"rewards/check_numbers/mean": -1.125,
"rewards/check_numbers/std": 0.75,
"reward": 0.375,
"reward_std": 0.25,
"frac_reward_zero_std": 0.0,
"completion_length": 296.75,
"kl": 0.003692739875987172,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.00039342198442048943,
"time_ms": 3241.851194994524,
"memory_mb": 160665.36376953125,
"memory_gb": 156.89976930618286
},
{
"step": 6,
"loss": 0.0,
"grad_norm": 0.0007444396032951772,
"learning_rate": 2.7777777777777783e-06,
"num_tokens": 23948.0,
"completions/mean_length": 1846.0,
"completions/min_length": 1846.0,
"completions/max_length": 1846.0,
"completions/clipped_ratio": 1.0,
"completions/mean_terminated_length": 0.0,
"completions/min_terminated_length": 0.0,
"completions/max_terminated_length": 0.0,
"rewards/match_format_exactly/mean": 0.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": -3.0,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.0,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -2.5,
"rewards/check_numbers/std": 0.0,
"reward": -7.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"completion_length": 1846.0,
"kl": 0.0025038770399987698,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0004721063813045873,
"time_ms": 10860.668059962336,
"memory_mb": 162817.607421875,
"memory_gb": 159.0015697479248
},
{
"step": 7,
"loss": 0.0371,
"grad_norm": 2.262518882751465,
"learning_rate": 2.222222222222222e-06,
"num_tokens": 26728.0,
"completions/mean_length": 552.0,
"completions/min_length": 506.0,
"completions/max_length": 627.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 552.0,
"completions/min_terminated_length": 506.0,
"completions/max_terminated_length": 627.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -1.5,
"rewards/check_answer/std": 2.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 1.5,
"reward_std": 2.0,
"frac_reward_zero_std": 0.0,
"completion_length": 552.0,
"kl": 0.006138760130852461,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0005507907781886852,
"time_ms": 4138.088690000586,
"memory_mb": 160989.1611328125,
"memory_gb": 157.2159776687622
},
{
"step": 8,
"loss": 0.0,
"grad_norm": 0.001562082557938993,
"learning_rate": 1.6666666666666667e-06,
"num_tokens": 29420.0,
"completions/mean_length": 606.0,
"completions/min_length": 560.0,
"completions/max_length": 636.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 606.0,
"completions/min_terminated_length": 560.0,
"completions/max_terminated_length": 636.0,
"rewards/match_format_exactly/mean": 0.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": -3.0,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.0,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -2.5,
"rewards/check_numbers/std": 0.0,
"reward": -7.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"completion_length": 606.0,
"kl": 0.004652692936360836,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0006294751750727831,
"time_ms": 4177.52773797838,
"memory_mb": 160996.86376953125,
"memory_gb": 157.22349977493286
},
{
"step": 9,
"loss": 0.0,
"grad_norm": 0.00027447607135400176,
"learning_rate": 1.111111111111111e-06,
"num_tokens": 31353.0,
"completions/mean_length": 386.25,
"completions/min_length": 334.0,
"completions/max_length": 464.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 386.25,
"completions/min_terminated_length": 334.0,
"completions/max_terminated_length": 464.0,
"rewards/match_format_exactly/mean": 3.0,
"rewards/match_format_exactly/std": 0.0,
"rewards/match_format_approximately/mean": 1.5,
"rewards/match_format_approximately/std": 0.0,
"rewards/check_answer/mean": -2.5,
"rewards/check_answer/std": 0.0,
"rewards/check_numbers/mean": -1.5,
"rewards/check_numbers/std": 0.0,
"reward": 0.5,
"reward_std": 0.0,
"frac_reward_zero_std": 1.0,
"completion_length": 386.25,
"kl": 0.0017617446137592196,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0007081595719568809,
"time_ms": 3263.2006779895164,
"memory_mb": 160745.5380859375,
"memory_gb": 156.97806453704834
},
{
"step": 10,
"loss": 0.2052,
"grad_norm": 0.4320540428161621,
"learning_rate": 5.555555555555555e-07,
"num_tokens": 35257.0,
"completions/mean_length": 808.0,
"completions/min_length": 650.0,
"completions/max_length": 1119.0,
"completions/clipped_ratio": 0.0,
"completions/mean_terminated_length": 808.0,
"completions/min_terminated_length": 650.0,
"completions/max_terminated_length": 1119.0,
"rewards/match_format_exactly/mean": 1.5,
"rewards/match_format_exactly/std": 1.7320507764816284,
"rewards/match_format_approximately/mean": 0.0,
"rewards/match_format_approximately/std": 2.1213202476501465,
"rewards/check_answer/mean": -3.25,
"rewards/check_answer/std": 1.4433757066726685,
"rewards/check_numbers/mean": -1.75,
"rewards/check_numbers/std": 0.5,
"reward": -3.5,
"reward_std": 2.8284270763397217,
"frac_reward_zero_std": 0.0,
"completion_length": 808.0,
"kl": 0.001998987514525652,
"clip_ratio/low_mean": 0.0,
"clip_ratio/low_min": 0.0,
"clip_ratio/high_mean": 0.0,
"clip_ratio/high_max": 0.0,
"clip_ratio/region_mean": 0.0,
"epoch": 0.0007868439688409789,
"time_ms": 6874.866564990953,
"memory_mb": 161736.77734375,
"memory_gb": 157.94607162475586
}
]

75
scripts/benchmarks/results/stats/grpo_vllm_10.summary.json Normal file
View file

@ -0,0 +1,75 @@
{
"backend": "vllm",
"max_steps": 10,
"train_wall_s": 74.41919421299826,
"median_step_ms_post_warmup": 4138.088690000586,
"n_logged_steps": 10,
"sampling": {
"temperature": 0.1,
"top_p": 0.97,
"min_p": 0.5,
"top_k": 5
},
"losses": [
0.0305,
-0.1941,
0.2632,
-0.2013,
0.2093,
0.0,
0.0371,
0.0,
0.0,
0.2052
],
"rewards": [
0.0,
-2.5,
-3.625,
0.0,
0.375,
-7.5,
1.5,
-7.5,
0.5,
-3.5
],
"kls": [
0.0,
0.0,
0.011923530139029026,
0.004221913404762745,
0.003692739875987172,
0.0025038770399987698,
0.006138760130852461,
0.004652692936360836,
0.0017617446137592196,
0.001998987514525652
],
"grad_norms": [
0.4133029878139496,
0.8333088159561157,
0.4677680730819702,
0.5163940191268921,
1.160618782043457,
0.0007444396032951772,
2.262518882751465,
0.001562082557938993,
0.00027447607135400176,
0.4320540428161621
],
"step_times_ms": [
17984.56621397054,
6704.717919987161,
12108.133931003977,
4019.3081409670413,
3241.851194994524,
10860.668059962336,
4138.088690000586,
4177.52773797838,
3263.2006779895164,
6874.866564990953
],
"peak_memory_gb": 157.94607162475586,
"logs_path": "logs/grpo_vllm_10.json"
}