381 lines
12 KiB
Python
381 lines
12 KiB
Python
"""Shared helpers for the Qwen3-4B GRPO comparison scripts.
|
|
|
|
Exports:
|
|
REASONING_START, REASONING_END, SOLUTION_START, SOLUTION_END, SYSTEM_PROMPT
|
|
CHAT_TEMPLATE
|
|
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.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import time
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from datasets import Dataset, load_dataset
|
|
|
|
|
|
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 %}"
|
|
"{{ '%%%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 %}"
|
|
"{% if add_generation_prompt %}{{ '%%%REASONING_START%%%' }}"
|
|
"{% endif %}"
|
|
)
|
|
|
|
|
|
def apply_chat_template_to_tokenizer(tokenizer):
|
|
tmpl = CHAT_TEMPLATE.replace("%%%SYSTEM_PROMPT%%%", SYSTEM_PROMPT)
|
|
tmpl = tmpl.replace("%%%REASONING_START%%%", REASONING_START)
|
|
tokenizer.chat_template = tmpl
|
|
return tokenizer
|
|
|
|
|
|
def build_dataset(tokenizer, *, max_seq_length: int = 2048):
|
|
"""Build the DAPO-Math-17k GRPO dataset with the prompt formatting from the
|
|
notebook. Returns `(dataset, maximum_prompt_length)`.
|
|
"""
|
|
ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
|
|
|
|
def _map_row(x):
|
|
return {
|
|
"prompt": [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": x["prompt"]},
|
|
],
|
|
"answer": x["solution"],
|
|
}
|
|
|
|
ds = ds.map(_map_row)
|
|
|
|
# Tokenize for length measurement (batched for speed).
|
|
def _tokenize(batch):
|
|
return {
|
|
"tokens": tokenizer.apply_chat_template(
|
|
batch["prompt"],
|
|
add_generation_prompt = True,
|
|
tokenize = True,
|
|
)
|
|
}
|
|
|
|
tokenized = ds.map(_tokenize, batched = True)
|
|
tokenized = tokenized.map(lambda x: {"L": len(x["tokens"])})
|
|
lengths = np.array(tokenized["L"])
|
|
maximum_length = int(np.quantile(lengths, 0.9))
|
|
ds = ds.select(np.where(lengths <= maximum_length)[0])
|
|
return ds, maximum_length
|
|
|
|
|
|
def build_reward_funcs(tokenizer):
|
|
"""Return the 4 reward functions used in the notebook, wired to `tokenizer`."""
|
|
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:
|
|
score = 0.0
|
|
response = completion[0]["content"]
|
|
if match_format.search(response) is not None:
|
|
score += 3.0
|
|
scores.append(score)
|
|
return scores
|
|
|
|
def match_format_approximately(completions, **kwargs):
|
|
scores = []
|
|
for completion in completions:
|
|
score = 0.0
|
|
response = completion[0]["content"]
|
|
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 = [
|
|
guess.group(1) if (guess := match_format.search(r)) is not None else None
|
|
for r in responses
|
|
]
|
|
scores = []
|
|
for guess, true_answer in zip(extracted, answer):
|
|
score = 0.0
|
|
if guess is None:
|
|
scores.append(-2.0)
|
|
continue
|
|
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
|
|
|
|
_printed_state = {"n": 0, "every": 5}
|
|
|
|
def check_numbers(prompts, completions, answer, **kwargs):
|
|
question = prompts[0][-1]["content"]
|
|
responses = [c[0]["content"] for c in completions]
|
|
extracted = [
|
|
guess.group(1) if (guess := match_numbers.search(r)) is not None else None
|
|
for r in responses
|
|
]
|
|
if _printed_state["n"] % _printed_state["every"] == 0:
|
|
print(
|
|
"*" * 20 + f"Question:\n{question}",
|
|
f"\nAnswer:\n{answer[0]}",
|
|
f"\nResponse:\n{responses[0]}",
|
|
f"\nExtracted:\n{extracted[0]}",
|
|
)
|
|
_printed_state["n"] += 1
|
|
|
|
scores = []
|
|
for guess, true_answer in zip(extracted, answer):
|
|
if guess is None:
|
|
scores.append(-2.5)
|
|
continue
|
|
try:
|
|
t = float(true_answer.strip())
|
|
g = float(guess.strip().replace(",", ""))
|
|
scores.append(3.5 if g == t else -1.5)
|
|
except Exception:
|
|
scores.append(0.0)
|
|
return scores
|
|
|
|
return [
|
|
match_format_exactly,
|
|
match_format_approximately,
|
|
check_answer,
|
|
check_numbers,
|
|
]
|
|
|
|
|
|
def build_grpo_kwargs(
|
|
tokenizer,
|
|
maximum_length: int,
|
|
*,
|
|
max_seq_length: int = 2048,
|
|
max_steps: int = 100,
|
|
num_generations: int = 4,
|
|
per_device_train_batch_size: int = 1,
|
|
gradient_accumulation_steps: int = 1,
|
|
output_dir: str = "outputs",
|
|
):
|
|
"""Return the shared dict of GRPOConfig kwargs used by both backends.
|
|
|
|
Caller adds backend-specific keys (use_vllm / use_transformers_paged / etc).
|
|
"""
|
|
max_prompt_length = maximum_length + 1
|
|
max_completion_length = max_seq_length - max_prompt_length
|
|
|
|
return dict(
|
|
temperature = 1.0,
|
|
top_p = 1.0,
|
|
top_k = -1,
|
|
min_p = 0.1,
|
|
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 = per_device_train_batch_size,
|
|
gradient_accumulation_steps = gradient_accumulation_steps,
|
|
num_generations = num_generations,
|
|
max_prompt_length = max_prompt_length,
|
|
max_completion_length = max_completion_length,
|
|
max_steps = max_steps,
|
|
save_steps = max_steps,
|
|
report_to = "none",
|
|
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
|