benchmarks: consolidate GRPO entrypoints and extract shared helpers
Delete two unreferenced drivers (qwen3_grpo_notebook.py, qwen3_grpo_unified.py)
that duplicated the canonical trio. Port the --compile_mode / --compile_dynamic
flags from unified into qwen3_grpo_naive.py and qwen3_grpo_tpaged.py before
deletion so the torch.compile path is preserved on the training-side backends
(vLLM is excluded because it owns its own inference graph).
Extract the 20-line StepTimer TrainerCallback, the per-step stats JSON writer,
the vLLM GuidedDecodingParams shim, and the optional torch.compile wrapper
into unsloth_grpo_common.py so the three canonical drivers
(qwen3_grpo_{vllm,naive,tpaged}.py) share one implementation. Stats schema is
unchanged: backend, train_wall_s, peak_memory_gb, step_wall_s, losses, rewards,
max_prompt_length, max_completion_length, num_generations, max_steps, plus
backend-specific extras (attn_impl, persistent_cb) passed through write_stats's
extra kwarg.
Add a short paragraph to scripts/benchmarks/README.md describing the new
--compile_mode flag.
Verified:
- python -m py_compile on all four modified files.
- --help on all three drivers shows --compile_mode on naive + tpaged only.
- 2-step tpaged smoke (flash_attention_2, num_generations=2, pdb=2) runs to
completion on B200. Stats JSON schema matches the pre-refactor output exactly.
Net: 7 files changed, +235 / -1093, 21 -> 19 benchmark files.
This commit is contained in:
Daniel Han
parent
8792e5da7b
commit
a68d346e77
7 changed files with 235 additions and 1093 deletions
|
|
@ -141,6 +141,13 @@ CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/qwen3_grpo_tpaged.py \
|
|||
--max_batch_tokens 16384 --num_blocks 16384
|
||||
```
|
||||
|
||||
`qwen3_grpo_naive.py` and `qwen3_grpo_tpaged.py` accept an optional
|
||||
`--compile_mode {default,reduce-overhead,max-autotune-no-cudagraphs}` flag.
|
||||
When set, `trainer.model.forward` (and `trainer.ref_model.forward`, if present)
|
||||
are wrapped with `torch.compile` after trainer construction. The vLLM driver
|
||||
has no such flag because vLLM owns its own inference graph. `--compile_dynamic`
|
||||
(default on) toggles dynamic-shape compilation.
|
||||
|
||||
## Known integration notes for transformers continuous batching + TRL + Unsloth
|
||||
|
||||
These are the sharp edges you hit going down the continuous-batching path and
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@ Run:
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
|
@ -28,32 +27,23 @@ from pathlib import Path
|
|||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE))
|
||||
|
||||
# Same vLLM sampling-params shim as the tpaged script so TRL imports cleanly
|
||||
# even when vLLM is installed but the GuidedDecodingParams symbol has moved.
|
||||
try:
|
||||
import vllm.sampling_params as _vllm_sp
|
||||
from unsloth_grpo_common import ( # noqa: E402
|
||||
StepTimer,
|
||||
apply_chat_template_to_tokenizer,
|
||||
build_dataset,
|
||||
build_grpo_kwargs,
|
||||
build_reward_funcs,
|
||||
install_vllm_sampling_shim,
|
||||
maybe_compile_trainer_forwards,
|
||||
write_stats,
|
||||
)
|
||||
|
||||
if not hasattr(_vllm_sp, "GuidedDecodingParams"):
|
||||
|
||||
class _GuidedDecodingParamsShim: # pragma: no cover
|
||||
def __init__(self, *a, **kw):
|
||||
pass
|
||||
|
||||
_vllm_sp.GuidedDecodingParams = _GuidedDecodingParamsShim
|
||||
except ImportError:
|
||||
pass
|
||||
install_vllm_sampling_shim()
|
||||
|
||||
import torch # noqa: E402
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402
|
||||
from peft import LoraConfig, get_peft_model # noqa: E402
|
||||
|
||||
from unsloth_grpo_common import ( # noqa: E402
|
||||
apply_chat_template_to_tokenizer,
|
||||
build_dataset,
|
||||
build_reward_funcs,
|
||||
build_grpo_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
|
|
@ -71,6 +61,13 @@ def parse_args():
|
|||
)
|
||||
p.add_argument("--output_dir", default = "outputs/grpo_naive")
|
||||
p.add_argument("--stats_path", default = "logs/naive_stats.json")
|
||||
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.",
|
||||
)
|
||||
p.add_argument("--compile_dynamic", action = "store_true", default = True)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
|
|
@ -149,30 +146,7 @@ def main():
|
|||
**shared,
|
||||
)
|
||||
|
||||
from transformers import TrainerCallback
|
||||
|
||||
timings = {"step_wall": [], "loss": [], "reward": []}
|
||||
|
||||
class StepTimer(TrainerCallback):
|
||||
def __init__(self):
|
||||
self.t0 = None
|
||||
|
||||
def on_step_begin(self, _args, state, control, **kwargs):
|
||||
torch.cuda.synchronize()
|
||||
self.t0 = time.perf_counter()
|
||||
|
||||
def on_log(self, _args, state, control, logs = None, **kwargs):
|
||||
if logs is None:
|
||||
return
|
||||
if "loss" in logs:
|
||||
timings["loss"].append(float(logs["loss"]))
|
||||
if "reward" in logs:
|
||||
timings["reward"].append(float(logs["reward"]))
|
||||
|
||||
def on_step_end(self, _args, state, control, **kwargs):
|
||||
if self.t0 is not None:
|
||||
torch.cuda.synchronize()
|
||||
timings["step_wall"].append(time.perf_counter() - self.t0)
|
||||
timer = StepTimer()
|
||||
|
||||
trainer = GRPOTrainer(
|
||||
model = model,
|
||||
|
|
@ -180,7 +154,11 @@ def main():
|
|||
reward_funcs = reward_funcs,
|
||||
args = training_args,
|
||||
train_dataset = dataset,
|
||||
callbacks = [StepTimer()],
|
||||
callbacks = [timer],
|
||||
)
|
||||
|
||||
maybe_compile_trainer_forwards(
|
||||
trainer, args.compile_mode, dynamic = args.compile_dynamic, tag = "naive"
|
||||
)
|
||||
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
|
|
@ -190,21 +168,18 @@ def main():
|
|||
|
||||
peak = torch.cuda.max_memory_allocated() / 1024**3
|
||||
|
||||
stats = {
|
||||
"backend": "naive_trl",
|
||||
"attn_impl": args.attn_impl,
|
||||
"train_wall_s": t_train,
|
||||
"peak_memory_gb": peak,
|
||||
"step_wall_s": timings["step_wall"],
|
||||
"losses": timings["loss"],
|
||||
"rewards": timings["reward"],
|
||||
"max_prompt_length": shared["max_prompt_length"],
|
||||
"max_completion_length": shared["max_completion_length"],
|
||||
"num_generations": args.num_generations,
|
||||
"max_steps": args.max_steps,
|
||||
}
|
||||
with open(args.stats_path, "w") as f:
|
||||
json.dump(stats, f, indent = 2)
|
||||
write_stats(
|
||||
args.stats_path,
|
||||
backend = "naive_trl",
|
||||
timer = timer,
|
||||
train_wall_s = t_train,
|
||||
peak_memory_gb = peak,
|
||||
max_prompt_length = shared["max_prompt_length"],
|
||||
max_completion_length = shared["max_completion_length"],
|
||||
num_generations = args.num_generations,
|
||||
max_steps = args.max_steps,
|
||||
extra = {"attn_impl": args.attn_impl},
|
||||
)
|
||||
print(f"[naive] Wrote stats to {args.stats_path}")
|
||||
print(f"[naive] Total train wall: {t_train:.1f}s Peak mem: {peak:.2f} GB")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,473 +0,0 @@
|
|||
"""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()
|
||||
|
|
@ -15,7 +15,6 @@ Run:
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
|
@ -24,23 +23,23 @@ from pathlib import Path
|
|||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE))
|
||||
|
||||
# Minimal shim so TRL's GRPOTrainer imports cleanly against newer vLLM
|
||||
# releases where `GuidedDecodingParams` has moved or been removed. We do NOT
|
||||
# `install_vllm_sampling_shim()` shims `vllm.sampling_params.GuidedDecodingParams`
|
||||
# for newer vLLM releases so TRL's GRPOTrainer imports cleanly. We do NOT
|
||||
# import `unsloth` here because that replaces TRL's GRPOTrainer with an
|
||||
# Unsloth-compiled variant that assumes the model has `for_training()` /
|
||||
# `for_inference()` hooks, which a vanilla HF model does not.
|
||||
try:
|
||||
import vllm.sampling_params as _vllm_sp
|
||||
from unsloth_grpo_common import ( # noqa: E402
|
||||
StepTimer,
|
||||
apply_chat_template_to_tokenizer,
|
||||
build_dataset,
|
||||
build_grpo_kwargs,
|
||||
build_reward_funcs,
|
||||
install_vllm_sampling_shim,
|
||||
maybe_compile_trainer_forwards,
|
||||
write_stats,
|
||||
)
|
||||
|
||||
if not hasattr(_vllm_sp, "GuidedDecodingParams"):
|
||||
|
||||
class _GuidedDecodingParamsShim: # pragma: no cover - used only if TRL asks
|
||||
def __init__(self, *a, **kw):
|
||||
pass
|
||||
|
||||
_vllm_sp.GuidedDecodingParams = _GuidedDecodingParamsShim
|
||||
except ImportError:
|
||||
pass
|
||||
install_vllm_sampling_shim()
|
||||
|
||||
import torch # noqa: E402
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402
|
||||
|
|
@ -50,13 +49,6 @@ import flash_attn_fa4_shim # noqa: E402
|
|||
|
||||
flash_attn_fa4_shim.apply()
|
||||
|
||||
from unsloth_grpo_common import ( # noqa: E402
|
||||
apply_chat_template_to_tokenizer,
|
||||
build_dataset,
|
||||
build_reward_funcs,
|
||||
build_grpo_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
|
|
@ -92,6 +84,13 @@ def parse_args():
|
|||
help = "Reuse one ContinuousBatchingManager across every training step instead "
|
||||
"of letting TRL's generate_batch rebuild it (and the paged cache) each step.",
|
||||
)
|
||||
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.",
|
||||
)
|
||||
p.add_argument("--compile_dynamic", action = "store_true", default = True)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
|
|
@ -176,30 +175,7 @@ def main():
|
|||
)
|
||||
|
||||
# 4. Timing callback.
|
||||
from transformers import TrainerCallback
|
||||
|
||||
timings = {"step_wall": [], "loss": [], "reward": []}
|
||||
|
||||
class StepTimer(TrainerCallback):
|
||||
def __init__(self):
|
||||
self.t0 = None
|
||||
|
||||
def on_step_begin(self, _args, state, control, **kwargs):
|
||||
torch.cuda.synchronize()
|
||||
self.t0 = time.perf_counter()
|
||||
|
||||
def on_log(self, _args, state, control, logs = None, **kwargs):
|
||||
if logs is None:
|
||||
return
|
||||
if "loss" in logs:
|
||||
timings["loss"].append(float(logs["loss"]))
|
||||
if "reward" in logs:
|
||||
timings["reward"].append(float(logs["reward"]))
|
||||
|
||||
def on_step_end(self, _args, state, control, **kwargs):
|
||||
if self.t0 is not None:
|
||||
torch.cuda.synchronize()
|
||||
timings["step_wall"].append(time.perf_counter() - self.t0)
|
||||
timer = StepTimer()
|
||||
|
||||
trainer = GRPOTrainer(
|
||||
model = model,
|
||||
|
|
@ -207,7 +183,11 @@ def main():
|
|||
reward_funcs = reward_funcs,
|
||||
args = training_args,
|
||||
train_dataset = dataset,
|
||||
callbacks = [StepTimer()],
|
||||
callbacks = [timer],
|
||||
)
|
||||
|
||||
maybe_compile_trainer_forwards(
|
||||
trainer, args.compile_mode, dynamic = args.compile_dynamic, tag = "tpaged"
|
||||
)
|
||||
|
||||
if args.persistent_cb:
|
||||
|
|
@ -239,22 +219,21 @@ def main():
|
|||
|
||||
peak = torch.cuda.max_memory_allocated() / 1024**3
|
||||
|
||||
stats = {
|
||||
"backend": "transformers_paged",
|
||||
"attn_impl": args.attn_impl,
|
||||
"train_wall_s": t_train,
|
||||
"peak_memory_gb": peak,
|
||||
"step_wall_s": timings["step_wall"],
|
||||
"losses": timings["loss"],
|
||||
"rewards": timings["reward"],
|
||||
"max_prompt_length": shared["max_prompt_length"],
|
||||
"max_completion_length": shared["max_completion_length"],
|
||||
"num_generations": args.num_generations,
|
||||
"max_steps": args.max_steps,
|
||||
"persistent_cb": args.persistent_cb,
|
||||
}
|
||||
with open(args.stats_path, "w") as f:
|
||||
json.dump(stats, f, indent = 2)
|
||||
write_stats(
|
||||
args.stats_path,
|
||||
backend = "transformers_paged",
|
||||
timer = timer,
|
||||
train_wall_s = t_train,
|
||||
peak_memory_gb = peak,
|
||||
max_prompt_length = shared["max_prompt_length"],
|
||||
max_completion_length = shared["max_completion_length"],
|
||||
num_generations = args.num_generations,
|
||||
max_steps = args.max_steps,
|
||||
extra = {
|
||||
"attn_impl": args.attn_impl,
|
||||
"persistent_cb": args.persistent_cb,
|
||||
},
|
||||
)
|
||||
print(f"[tpaged] Wrote stats to {args.stats_path}")
|
||||
print(f"[tpaged] Total train wall: {t_train:.1f}s Peak mem: {peak:.2f} GB")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,456 +0,0 @@
|
|||
"""Unified entrypoint for Qwen3-4B GRPO backend comparison.
|
||||
|
||||
Single script, N backends. Identical dataset / reward functions / sampling /
|
||||
callbacks so per-step loss, reward, KL, and grad-norm arrays are directly
|
||||
comparable across runs.
|
||||
|
||||
Backends (pick one via `--backend`):
|
||||
vllm : Unsloth fast_inference=True (vLLM colocated).
|
||||
unsloth_fi_false : Unsloth fast_inference=False (custom HF inference
|
||||
kernels + cached fp16 LoRA in fast_linear_forward).
|
||||
Uses trainer's default (non-vLLM, non-CB) rollout path.
|
||||
cb_paged : Vanilla HF + PEFT LoRA + transformers continuous
|
||||
batching with `attn_implementation="paged_attention"`
|
||||
(FA4 shim active).
|
||||
cb_sdpa : Same but with `attn_implementation="sdpa_paged"`.
|
||||
naive_trl : Vanilla HF + PEFT LoRA, no CB, no vLLM (TRL's naive
|
||||
generate path).
|
||||
|
||||
Run:
|
||||
CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/qwen3_grpo_unified.py \
|
||||
--backend vllm --max_steps 10 \
|
||||
--output_dir outputs/grpo_vllm_10 \
|
||||
--stats_path logs/grpo_vllm_10.json
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
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))
|
||||
|
||||
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)
|
||||
# Phase 4: torch.compile on the training forward.
|
||||
p.add_argument(
|
||||
"--compile_mode",
|
||||
default = None,
|
||||
choices = [None, "default", "reduce-overhead", "max-autotune-no-cudagraphs"],
|
||||
help = "If set, torch.compile(model.forward, mode=...) after "
|
||||
"the trainer is built. vllm backend is excluded; the "
|
||||
"rollout engine owns its own compile pipeline.",
|
||||
)
|
||||
p.add_argument("--compile_dynamic", action = "store_true", default = True)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
def _make_stats_callback():
|
||||
"""StatisticsCallback from torch_debugging_utils. Logs per-step loss,
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
def _maybe_shim_guided_decoding():
|
||||
"""Newer vLLM releases have moved GuidedDecodingParams out of
|
||||
`vllm.sampling_params`; TRL's GRPOTrainer still tries to import it on
|
||||
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
|
||||
|
||||
|
||||
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 = 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,
|
||||
)
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
def _load_vanilla_hf(args, attn_impl: str):
|
||||
"""Vanilla HF + PEFT LoRA. Used by cb_paged / cb_sdpa / naive_trl."""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from peft import LoraConfig, get_peft_model
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model_name,
|
||||
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",
|
||||
)
|
||||
model = get_peft_model(model, lora)
|
||||
try:
|
||||
model.gradient_checkpointing_enable(
|
||||
gradient_checkpointing_kwargs = {"use_reentrant": False}
|
||||
)
|
||||
except TypeError:
|
||||
model.gradient_checkpointing_enable()
|
||||
model.enable_input_require_grads()
|
||||
return model, tokenizer
|
||||
|
||||
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
# 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)
|
||||
elif args.backend == "unsloth_fi_false":
|
||||
model, tokenizer = _load_unsloth(args, fast_inference = False)
|
||||
elif args.backend == "cb_paged":
|
||||
# `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")
|
||||
elif args.backend == "naive_trl":
|
||||
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
|
||||
)
|
||||
print(f"[{args.backend}] p90 prompt length = {maximum_length}")
|
||||
reward_funcs = build_reward_funcs(tokenizer)
|
||||
|
||||
# --- GRPOConfig: shared core, backend-specific flags ----------------------
|
||||
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,
|
||||
)
|
||||
# Overwrite the equivalence-friendly sampling params.
|
||||
shared["temperature"] = args.temperature
|
||||
shared["top_p"] = args.top_p
|
||||
shared["min_p"] = args.min_p
|
||||
# TRL's TopKLogitsWarper rejects -1; accept an int >=0 only.
|
||||
shared["top_k"] = args.top_k if args.top_k and args.top_k > 0 else None
|
||||
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,
|
||||
)
|
||||
training_args = GRPOConfig(
|
||||
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":
|
||||
# Trainer's default rollout path: model.generate. Unsloth's
|
||||
# fast_inference=False + for_inference() wires the fast single-token
|
||||
# decode + cached fp16 LoRA.
|
||||
training_args = GRPOConfig(
|
||||
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 = {
|
||||
"max_batch_tokens": args.max_batch_tokens,
|
||||
"num_blocks": args.num_blocks,
|
||||
},
|
||||
**shared,
|
||||
)
|
||||
else: # naive_trl
|
||||
training_args = GRPOConfig(
|
||||
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],
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
install_for_model(base, trainer.generation_config)
|
||||
persistent_teardown_target = base
|
||||
|
||||
# Phase 4: torch.compile on the training forward.
|
||||
if args.compile_mode and args.backend != "vllm":
|
||||
from torch_debugging_utils import clear_inductor_cache, CompileDebugger
|
||||
|
||||
clear_inductor_cache()
|
||||
CompileDebugger.enable(graph_breaks = True, recompiles = True)
|
||||
# Raise Dynamo cache limit so dynamic-shape recompiles don't thrash.
|
||||
import torch._dynamo
|
||||
|
||||
torch._dynamo.config.cache_size_limit = 128
|
||||
try:
|
||||
torch._dynamo.config.allow_unspec_int_on_nn_module = True
|
||||
except AttributeError:
|
||||
pass
|
||||
print(
|
||||
f"[{args.backend}] Compiling trainer.model.forward "
|
||||
f"(mode={args.compile_mode}, dynamic={args.compile_dynamic})"
|
||||
)
|
||||
trainer.model.forward = torch.compile(
|
||||
trainer.model.forward,
|
||||
mode = args.compile_mode,
|
||||
dynamic = args.compile_dynamic,
|
||||
)
|
||||
# Reference model inside TRL's GRPO loop also runs a forward.
|
||||
ref = getattr(trainer, "ref_model", None)
|
||||
if ref is not None:
|
||||
ref.forward = torch.compile(
|
||||
ref.forward,
|
||||
mode = args.compile_mode,
|
||||
dynamic = args.compile_dynamic,
|
||||
)
|
||||
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
t_start = time.perf_counter()
|
||||
try:
|
||||
trainer.train()
|
||||
finally:
|
||||
if persistent_teardown_target is not None:
|
||||
from persistent_cb import teardown
|
||||
|
||||
teardown(persistent_teardown_target)
|
||||
train_wall = time.perf_counter() - t_start
|
||||
|
||||
stats_cb.save_logs(args.stats_path)
|
||||
|
||||
times = [l["time_ms"] for l in stats_cb.logs if "time_ms" in l]
|
||||
losses = [l["loss"] for l in stats_cb.logs if "loss" in l]
|
||||
rewards = [l.get("reward") for l in stats_cb.logs if "reward" in l]
|
||||
kls = [l.get("kl") for l in stats_cb.logs if "kl" in l]
|
||||
grad_norms = [l.get("grad_norm") for l in stats_cb.logs if "grad_norm" in l]
|
||||
|
||||
# Post-warmup (skip first 3 steps) median.
|
||||
median_step_ms = None
|
||||
if len(times) > 3:
|
||||
post = sorted(times[3:])
|
||||
median_step_ms = post[len(post) // 2]
|
||||
|
||||
summary = {
|
||||
"backend": args.backend,
|
||||
"max_steps": args.max_steps,
|
||||
"train_wall_s": train_wall,
|
||||
"median_step_ms_post_warmup": median_step_ms,
|
||||
"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,
|
||||
},
|
||||
"losses": losses,
|
||||
"rewards": rewards,
|
||||
"kls": kls,
|
||||
"grad_norms": grad_norms,
|
||||
"step_times_ms": times,
|
||||
"peak_memory_gb": torch.cuda.max_memory_allocated() / 1024**3,
|
||||
"logs_path": args.stats_path,
|
||||
}
|
||||
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,
|
||||
)
|
||||
)
|
||||
print(f"\n[{args.backend}] wrote summary to {summary_path}")
|
||||
# vLLM engine holds refs; fast-exit rather than wait for shutdown.
|
||||
os._exit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -9,7 +9,6 @@ Run:
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
|
@ -26,10 +25,12 @@ from unsloth import FastLanguageModel # noqa: E402
|
|||
import torch # noqa: E402
|
||||
|
||||
from unsloth_grpo_common import ( # noqa: E402
|
||||
StepTimer,
|
||||
apply_chat_template_to_tokenizer,
|
||||
build_dataset,
|
||||
build_reward_funcs,
|
||||
build_grpo_kwargs,
|
||||
build_reward_funcs,
|
||||
write_stats,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -121,30 +122,7 @@ def main():
|
|||
)
|
||||
|
||||
# 5. Timing callback.
|
||||
from transformers import TrainerCallback
|
||||
|
||||
timings = {"step_wall": [], "loss": [], "reward": []}
|
||||
|
||||
class StepTimer(TrainerCallback):
|
||||
def __init__(self):
|
||||
self.t0 = None
|
||||
|
||||
def on_step_begin(self, _args, state, control, **kwargs):
|
||||
torch.cuda.synchronize()
|
||||
self.t0 = time.perf_counter()
|
||||
|
||||
def on_log(self, _args, state, control, logs = None, **kwargs):
|
||||
if logs is None:
|
||||
return
|
||||
if "loss" in logs:
|
||||
timings["loss"].append(float(logs["loss"]))
|
||||
if "reward" in logs:
|
||||
timings["reward"].append(float(logs["reward"]))
|
||||
|
||||
def on_step_end(self, _args, state, control, **kwargs):
|
||||
if self.t0 is not None:
|
||||
torch.cuda.synchronize()
|
||||
timings["step_wall"].append(time.perf_counter() - self.t0)
|
||||
timer = StepTimer()
|
||||
|
||||
trainer = GRPOTrainer(
|
||||
model = model,
|
||||
|
|
@ -152,7 +130,7 @@ def main():
|
|||
reward_funcs = reward_funcs,
|
||||
args = training_args,
|
||||
train_dataset = dataset,
|
||||
callbacks = [StepTimer()],
|
||||
callbacks = [timer],
|
||||
)
|
||||
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
|
|
@ -162,20 +140,17 @@ def main():
|
|||
|
||||
peak = torch.cuda.max_memory_allocated() / 1024**3
|
||||
|
||||
stats = {
|
||||
"backend": "vllm_colocated",
|
||||
"train_wall_s": t_train,
|
||||
"peak_memory_gb": peak,
|
||||
"step_wall_s": timings["step_wall"],
|
||||
"losses": timings["loss"],
|
||||
"rewards": timings["reward"],
|
||||
"max_prompt_length": shared["max_prompt_length"],
|
||||
"max_completion_length": shared["max_completion_length"],
|
||||
"num_generations": args.num_generations,
|
||||
"max_steps": args.max_steps,
|
||||
}
|
||||
with open(args.stats_path, "w") as f:
|
||||
json.dump(stats, f, indent = 2)
|
||||
write_stats(
|
||||
args.stats_path,
|
||||
backend = "vllm_colocated",
|
||||
timer = timer,
|
||||
train_wall_s = t_train,
|
||||
peak_memory_gb = peak,
|
||||
max_prompt_length = shared["max_prompt_length"],
|
||||
max_completion_length = shared["max_completion_length"],
|
||||
num_generations = args.num_generations,
|
||||
max_steps = args.max_steps,
|
||||
)
|
||||
print(f"[vllm] Wrote stats to {args.stats_path}")
|
||||
print(f"[vllm] Total train wall: {t_train:.1f}s Peak mem: {peak:.2f} GB")
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,9 @@ Exports:
|
|||
build_dataset(tokenizer, max_seq_length=2048)
|
||||
build_reward_funcs(tokenizer)
|
||||
build_grpo_kwargs(tokenizer, maximum_length, max_seq_length)
|
||||
StepTimer (TrainerCallback recording per-step wall time, loss, reward)
|
||||
write_stats(path, backend, timer, ...)
|
||||
install_vllm_sampling_shim()
|
||||
|
||||
Keeps dataset loading, chat template, formatting rewards, and GRPO hparams
|
||||
identical between the vLLM baseline script and the transformers-CB candidate.
|
||||
|
|
@ -13,7 +16,10 @@ identical between the vLLM baseline script and the transformers-CB candidate.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from datasets import Dataset, load_dataset
|
||||
|
|
@ -240,3 +246,132 @@ def build_grpo_kwargs(
|
|||
output_dir = output_dir,
|
||||
seed = 3407,
|
||||
)
|
||||
|
||||
|
||||
import torch
|
||||
from transformers import TrainerCallback
|
||||
|
||||
|
||||
class StepTimer(TrainerCallback):
|
||||
"""Per-step wall time / loss / reward recorder.
|
||||
|
||||
Shared verbatim across qwen3_grpo_{vllm,naive,tpaged}. Records step wall
|
||||
time in `self.step_wall`, and picks up `loss` / `reward` from the TRL log
|
||||
dict in `on_log`.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.t0 = None
|
||||
self.step_wall = []
|
||||
self.loss = []
|
||||
self.reward = []
|
||||
|
||||
def on_step_begin(self, _args, state, control, **kwargs):
|
||||
torch.cuda.synchronize()
|
||||
self.t0 = time.perf_counter()
|
||||
|
||||
def on_log(self, _args, state, control, logs = None, **kwargs):
|
||||
if logs is None:
|
||||
return
|
||||
if "loss" in logs:
|
||||
self.loss.append(float(logs["loss"]))
|
||||
if "reward" in logs:
|
||||
self.reward.append(float(logs["reward"]))
|
||||
|
||||
def on_step_end(self, _args, state, control, **kwargs):
|
||||
if self.t0 is not None:
|
||||
torch.cuda.synchronize()
|
||||
self.step_wall.append(time.perf_counter() - self.t0)
|
||||
|
||||
|
||||
def write_stats(
|
||||
path: str,
|
||||
backend: str,
|
||||
timer: StepTimer,
|
||||
*,
|
||||
train_wall_s: float,
|
||||
peak_memory_gb: float,
|
||||
max_prompt_length: int,
|
||||
max_completion_length: int,
|
||||
num_generations: int,
|
||||
max_steps: int,
|
||||
extra: dict | None = None,
|
||||
) -> None:
|
||||
"""Dump the per-step stats dict used by all three GRPO drivers.
|
||||
|
||||
Schema matches the pre-refactor output exactly: `backend`, `train_wall_s`,
|
||||
`peak_memory_gb`, `step_wall_s`, `losses`, `rewards`, `max_prompt_length`,
|
||||
`max_completion_length`, `num_generations`, `max_steps`, plus any
|
||||
backend-specific keys passed in `extra` (e.g. `attn_impl`, `persistent_cb`).
|
||||
"""
|
||||
stats = {
|
||||
"backend": backend,
|
||||
"train_wall_s": train_wall_s,
|
||||
"peak_memory_gb": peak_memory_gb,
|
||||
"step_wall_s": timer.step_wall,
|
||||
"losses": timer.loss,
|
||||
"rewards": timer.reward,
|
||||
"max_prompt_length": max_prompt_length,
|
||||
"max_completion_length": max_completion_length,
|
||||
"num_generations": num_generations,
|
||||
"max_steps": max_steps,
|
||||
}
|
||||
if extra:
|
||||
stats.update(extra)
|
||||
with open(path, "w") as f:
|
||||
json.dump(stats, f, indent = 2)
|
||||
|
||||
|
||||
def maybe_compile_trainer_forwards(trainer, compile_mode, *, dynamic: bool = True, tag: str = ""):
|
||||
"""torch.compile wrap `trainer.model.forward` and (if present)
|
||||
`trainer.ref_model.forward`. No-op if `compile_mode` is falsy.
|
||||
|
||||
Ported from the old `qwen3_grpo_unified.py` compile path, minus the
|
||||
out-of-tree `torch_debugging_utils` imports that were dev-only.
|
||||
"""
|
||||
if not compile_mode:
|
||||
return
|
||||
import torch._dynamo
|
||||
|
||||
torch._dynamo.config.cache_size_limit = 128
|
||||
try:
|
||||
torch._dynamo.config.allow_unspec_int_on_nn_module = True
|
||||
except AttributeError:
|
||||
pass
|
||||
prefix = f"[{tag}] " if tag else ""
|
||||
print(f"{prefix}Compiling trainer.model.forward (mode={compile_mode}, dynamic={dynamic})")
|
||||
trainer.model.forward = torch.compile(
|
||||
trainer.model.forward,
|
||||
mode = compile_mode,
|
||||
dynamic = dynamic,
|
||||
)
|
||||
ref = getattr(trainer, "ref_model", None)
|
||||
if ref is not None:
|
||||
ref.forward = torch.compile(
|
||||
ref.forward,
|
||||
mode = compile_mode,
|
||||
dynamic = dynamic,
|
||||
)
|
||||
|
||||
|
||||
def install_vllm_sampling_shim():
|
||||
"""Shim `vllm.sampling_params.GuidedDecodingParams` for newer vLLM releases.
|
||||
|
||||
TRL's `GRPOTrainer` imports `GuidedDecodingParams` from
|
||||
`vllm.sampling_params`; newer vLLM versions have moved or removed it. Inject
|
||||
a no-op class so the import succeeds even on the non-vLLM training paths
|
||||
(naive, tpaged). No-op if vLLM is not installed or already exposes the
|
||||
symbol.
|
||||
"""
|
||||
try:
|
||||
import vllm.sampling_params as _vllm_sp
|
||||
except ImportError:
|
||||
return
|
||||
if hasattr(_vllm_sp, "GuidedDecodingParams"):
|
||||
return
|
||||
|
||||
class _GuidedDecodingParamsShim: # pragma: no cover - used only if TRL asks
|
||||
def __init__(self, *a, **kw):
|
||||
pass
|
||||
|
||||
_vllm_sp.GuidedDecodingParams = _GuidedDecodingParamsShim
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue