From a68d346e77bcc41af1f91517c1f6f3d526b58d3d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 21 Apr 2026 06:18:23 +0000 Subject: [PATCH] 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. --- scripts/benchmarks/README.md | 7 + scripts/benchmarks/qwen3_grpo_naive.py | 97 ++--- scripts/benchmarks/qwen3_grpo_notebook.py | 473 ---------------------- scripts/benchmarks/qwen3_grpo_tpaged.py | 103 ++--- scripts/benchmarks/qwen3_grpo_unified.py | 456 --------------------- scripts/benchmarks/qwen3_grpo_vllm.py | 57 +-- scripts/benchmarks/unsloth_grpo_common.py | 135 ++++++ 7 files changed, 235 insertions(+), 1093 deletions(-) delete mode 100644 scripts/benchmarks/qwen3_grpo_notebook.py delete mode 100644 scripts/benchmarks/qwen3_grpo_unified.py diff --git a/scripts/benchmarks/README.md b/scripts/benchmarks/README.md index 0d6ac62e4f..c9cd60b671 100644 --- a/scripts/benchmarks/README.md +++ b/scripts/benchmarks/README.md @@ -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 diff --git a/scripts/benchmarks/qwen3_grpo_naive.py b/scripts/benchmarks/qwen3_grpo_naive.py index b6b079b1f1..3d8a5dcb9e 100644 --- a/scripts/benchmarks/qwen3_grpo_naive.py +++ b/scripts/benchmarks/qwen3_grpo_naive.py @@ -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") diff --git a/scripts/benchmarks/qwen3_grpo_notebook.py b/scripts/benchmarks/qwen3_grpo_notebook.py deleted file mode 100644 index de01cd6344..0000000000 --- a/scripts/benchmarks/qwen3_grpo_notebook.py +++ /dev/null @@ -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 = "" - reasoning_end = "" - solution_start = "" - solution_end = "" - - 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("", "") - .replace("", "") - .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"[\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() diff --git a/scripts/benchmarks/qwen3_grpo_tpaged.py b/scripts/benchmarks/qwen3_grpo_tpaged.py index 13d6743427..c2b8839704 100644 --- a/scripts/benchmarks/qwen3_grpo_tpaged.py +++ b/scripts/benchmarks/qwen3_grpo_tpaged.py @@ -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") diff --git a/scripts/benchmarks/qwen3_grpo_unified.py b/scripts/benchmarks/qwen3_grpo_unified.py deleted file mode 100644 index 9cf412d27d..0000000000 --- a/scripts/benchmarks/qwen3_grpo_unified.py +++ /dev/null @@ -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() diff --git a/scripts/benchmarks/qwen3_grpo_vllm.py b/scripts/benchmarks/qwen3_grpo_vllm.py index 4dd8073d74..d78b131c46 100644 --- a/scripts/benchmarks/qwen3_grpo_vllm.py +++ b/scripts/benchmarks/qwen3_grpo_vllm.py @@ -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") diff --git a/scripts/benchmarks/unsloth_grpo_common.py b/scripts/benchmarks/unsloth_grpo_common.py index ab9921f540..a015331c84 100644 --- a/scripts/benchmarks/unsloth_grpo_common.py +++ b/scripts/benchmarks/unsloth_grpo_common.py @@ -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