5f3e1c98df
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.
430 lines
16 KiB
Python
430 lines
16 KiB
Python
"""Canonical reference run of Unsloth's Qwen3-4B GRPO notebook.
|
|
|
|
Ports `Qwen3_(4B)-GRPO.ipynb` to a single script with three deviations from the
|
|
notebook:
|
|
|
|
1. `max_steps = 10` (vibe check; escalate to 30/100 later).
|
|
2. Equivalence sampling params (`temperature=0.1, top_p=0.97, min_p=0.5,
|
|
top_k=5`) so KL/reward trajectories across backends can be compared.
|
|
3. `StatisticsCallback` from `torch_debugging_utils` logs per-step loss, reward,
|
|
grad-norm, KL, memory, and step wall time to `--stats_path`.
|
|
|
|
Run:
|
|
CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/qwen3_grpo_notebook.py \
|
|
--stats_path logs/notebook_ref_10.json --max_steps 10
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
# torch_debugging_utils + the shared benchmark helpers live at workspace root.
|
|
HERE = Path(__file__).resolve().parent
|
|
WORKSPACE_ROOT = Path("/mnt/disks/unslothai/ubuntu/workspace_31")
|
|
for p in (HERE, WORKSPACE_ROOT):
|
|
sys.path.insert(0, str(p))
|
|
|
|
|
|
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.")
|
|
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)
|
|
|
|
# Import order matters: unsloth must come before transformers/trl.
|
|
os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
|
|
from unsloth import FastLanguageModel # noqa: E402
|
|
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 = 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=3407,
|
|
)
|
|
|
|
reasoning_start = "<start_working_out>"
|
|
reasoning_end = "<end_working_out>"
|
|
solution_start = "<SOLUTION>"
|
|
solution_end = "</SOLUTION>"
|
|
|
|
system_prompt = (
|
|
"You are given a problem.\n"
|
|
"Think about the problem and provide your working out.\n"
|
|
f"Place it between {reasoning_start} and {reasoning_end}.\n"
|
|
f"Then, provide your solution between {solution_start}{solution_end}"
|
|
)
|
|
|
|
chat_template = (
|
|
"{% if messages[0]['role'] == 'system' %}"
|
|
"{{ messages[0]['content'] + eos_token }}"
|
|
"{% set loop_messages = messages[1:] %}"
|
|
"{% else %}"
|
|
f"{{{{ '{system_prompt}' + eos_token }}}}"
|
|
"{% set loop_messages = messages %}"
|
|
"{% endif %}"
|
|
"{% for message in loop_messages %}"
|
|
"{% if message['role'] == 'user' %}"
|
|
"{{ message['content'] }}"
|
|
"{% elif message['role'] == 'assistant' %}"
|
|
"{{ message['content'] + eos_token }}"
|
|
"{% endif %}"
|
|
"{% endfor %}"
|
|
f"{{% if add_generation_prompt %}}{{{{ '{reasoning_start}' }}}}"
|
|
"{% endif %}"
|
|
)
|
|
tokenizer.chat_template = chat_template
|
|
|
|
# --- pre fine-tune SFT stage (format priming) -----------------------------
|
|
from datasets import Dataset, load_dataset
|
|
import pandas as pd
|
|
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_df = sft_df.iloc[np.where(is_number)[0]]
|
|
|
|
def format_dataset(x):
|
|
thoughts = x["generated_solution"].replace("<think>", "").replace("</think>", "").strip()
|
|
final_prompt = (
|
|
reasoning_start + thoughts + reasoning_end
|
|
+ solution_start + x["expected_answer"] + solution_end
|
|
)
|
|
return [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": x["problem"]},
|
|
{"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 = 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_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"),
|
|
),
|
|
)
|
|
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"],
|
|
})
|
|
|
|
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,
|
|
)
|
|
match_numbers = re.compile(
|
|
solution_start + r".*?[\s]{0,}([-]?[\d\.\,]{1,})",
|
|
flags=re.MULTILINE | re.DOTALL,
|
|
)
|
|
|
|
def match_format_exactly(completions, **kwargs):
|
|
scores = []
|
|
for completion in completions:
|
|
response = completion[0]["content"]
|
|
scores.append(3.0 if match_format.search(response) is not None else 0.0)
|
|
return scores
|
|
|
|
def match_format_approximately(completions, **kwargs):
|
|
scores = []
|
|
for completion in completions:
|
|
response = completion[0]["content"]
|
|
score = 0.0
|
|
score += 0.5 if response.count(reasoning_end) == 1 else -1.0
|
|
score += 0.5 if response.count(solution_start) == 1 else -1.0
|
|
score += 0.5 if response.count(solution_end) == 1 else -1.0
|
|
scores.append(score)
|
|
return scores
|
|
|
|
def check_answer(prompts, completions, answer, **kwargs):
|
|
responses = [c[0]["content"] for c in completions]
|
|
extracted = [
|
|
g.group(1) if (g := match_format.search(r)) is not None else None
|
|
for r in responses
|
|
]
|
|
scores = []
|
|
for guess, true_answer in zip(extracted, answer):
|
|
if guess is None:
|
|
scores.append(-2.0)
|
|
continue
|
|
score = 0.0
|
|
if guess == true_answer:
|
|
score += 5.0
|
|
elif guess.strip() == true_answer.strip():
|
|
score += 3.5
|
|
else:
|
|
try:
|
|
ratio = float(guess) / float(true_answer)
|
|
if 0.9 <= ratio <= 1.1:
|
|
score += 2.0
|
|
elif 0.8 <= ratio <= 1.2:
|
|
score += 1.5
|
|
else:
|
|
score -= 2.5
|
|
except Exception:
|
|
score -= 4.5
|
|
scores.append(score)
|
|
return scores
|
|
|
|
def check_numbers(prompts, completions, answer, **kwargs):
|
|
responses = [c[0]["content"] for c in completions]
|
|
extracted = [
|
|
g.group(1) if (g := match_numbers.search(r)) is not None else None
|
|
for r in responses
|
|
]
|
|
scores = []
|
|
for guess, true_answer in zip(extracted, answer):
|
|
if guess is None:
|
|
scores.append(-2.5)
|
|
continue
|
|
try:
|
|
true_answer = float(true_answer.strip())
|
|
guess = float(guess.strip().replace(",", ""))
|
|
scores.append(3.5 if guess == true_answer else -1.5)
|
|
except Exception:
|
|
scores.append(0.0)
|
|
return scores
|
|
|
|
# Filter long prompts.
|
|
tokenized = dataset.map(
|
|
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))
|
|
print(f"Max prompt length (90th pct): {maximum_length}")
|
|
dataset = dataset.select(np.where(np.array(tokenized["L"]) <= maximum_length)[0])
|
|
del tokenized
|
|
|
|
max_prompt_length = maximum_length + 1
|
|
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,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
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
|
|
)
|
|
|
|
trainer = GRPOTrainer(
|
|
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],
|
|
)
|
|
|
|
t0 = time.perf_counter()
|
|
trainer.train()
|
|
train_wall = time.perf_counter() - t0
|
|
|
|
stats_cb.save_logs(args.stats_path)
|
|
|
|
# Post-warmup median step wall (skip first 3 steps).
|
|
times = [l["time_ms"] for l in stats_cb.logs if "time_ms" in l]
|
|
med_after_warmup = None
|
|
if len(times) > 3:
|
|
post = sorted(times[3:])
|
|
med_after_warmup = post[len(post) // 2]
|
|
|
|
summary = {
|
|
"backend": "unsloth_fast_inference_vllm",
|
|
"max_steps": args.max_steps,
|
|
"train_wall_s": train_wall,
|
|
"median_step_ms_post_warmup": med_after_warmup,
|
|
"n_logged_steps": len(stats_cb.logs),
|
|
"sampling": {
|
|
"temperature": args.temperature,
|
|
"top_p": args.top_p,
|
|
"min_p": args.min_p,
|
|
"top_k": args.top_k,
|
|
},
|
|
"logs_path": args.stats_path,
|
|
"peak_memory_gb": torch.cuda.max_memory_allocated() / 1024**3,
|
|
}
|
|
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,
|
|
)
|
|
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?"}],
|
|
]
|
|
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:
|
|
print(f"[warn] probe generation skipped: {e}")
|
|
|
|
# Emit the Phase 0 markdown report.
|
|
md_path = Path(args.output_dir) / "summary.md"
|
|
lines = [
|
|
f"# Phase 0 reference run: Qwen3-4B GRPO (Unsloth fast_inference=True)\n",
|
|
f"- max_steps: `{args.max_steps}`",
|
|
f"- sampling: `temperature={args.temperature}, top_p={args.top_p}, min_p={args.min_p}, top_k={args.top_k}`",
|
|
f"- train_wall_s: `{train_wall:.2f}`",
|
|
f"- median_step_ms (steps 4+): `{med_after_warmup}`",
|
|
f"- peak_memory_gb: `{summary['peak_memory_gb']:.2f}`\n",
|
|
"## Per-step logs\n",
|
|
"| step | loss | reward | kl | grad_norm | time_ms | mem_gb |",
|
|
"|---|---|---|---|---|---|---|",
|
|
]
|
|
for l in stats_cb.logs:
|
|
lines.append(
|
|
f"| {l.get('step','?')} | "
|
|
f"{l.get('loss','')} | "
|
|
f"{l.get('reward','')} | "
|
|
f"{l.get('kl','')} | "
|
|
f"{l.get('grad_norm','')} | "
|
|
f"{l.get('time_ms','')} | "
|
|
f"{l.get('memory_gb','')} |"
|
|
)
|
|
if rollouts:
|
|
lines.append("\n## Sample rollouts (post-training)\n")
|
|
for i, r in enumerate(rollouts[:3]):
|
|
lines.append(f"### Prompt {i+1}\n")
|
|
lines.append(f"```\n{r['prompt']}\n```\n")
|
|
lines.append(f"**Completion:**\n\n```\n{r['completion']}\n```\n")
|
|
md_path.write_text("\n".join(lines))
|
|
print(f"\nWrote {md_path}")
|
|
|
|
# Release vLLM engine and exit cleanly.
|
|
os._exit(0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|