From 5f3e1c98df6d3d0367ecca2919dc58e82008c089 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Apr 2026 14:23:51 +0000 Subject: [PATCH] Phase 2 vibe (10-step): vllm vs unsloth_fi_false vs cb_paged + Phase 3 fixes Phase 2 results (`scripts/benchmarks/results/grpo_equivalence.md`): - vLLM: 74.4s train, 4.14s median step, 158 GB peak (100%) - unsloth_fi_false: 355s train, 23.95s median, 10.7 GB peak (17%) - cb_paged (via sdpa_paged load): 466s train, 36s median, 55.6 GB peak (11.5%) Coherence gate passes on all three backends: losses finite, rewards in the expected early-GRPO range, KL trajectories qualitatively matched between vLLM and unsloth_fi_false in [0, 0.015]. Memory story is striking: unsloth_fi_false uses 15x less memory than vLLM. qwen3_grpo_unified.py fixes: - Auto-adjust per_device_train_batch_size -> num_generations for vanilla-HF backends (Unsloth's loader does this automatically; TRL on the HF path doesn't and crashes on the divisibility check). - cb_paged now loads with sdpa_paged (not paged_attention). The FA4 paged_attention kernel requires cu_seq_lens_q on every forward, but the GRPO training forward feeds a dense batch without them. sdpa_paged gracefully falls back to plain SDPA in that case and still exercises the paged path during the CB rollout. cb_sync_driver.py fixes: - FIFOScheduler no longer accepts manual_eviction in its signature; dropped. - drive_until_empty used to check has_pending_requests() before calling prepare_next_batch(), which returned False at startup because nothing had yet been pulled from the input_queue. Now the loop drains the input queue first and exits only when both queues + scheduler are empty. Smoke test on GPU 1 (8 prompts, 64 tokens): eager path produces 512 correct tokens; CUDA-graph path hangs during first-step capture (PagedAttentionCache probably allocates on first use). Tracked for the next commit. --- scripts/benchmarks/cb_sync_driver.py | 144 ++++--- scripts/benchmarks/cb_vs_vllm_generation.py | 211 +++++----- scripts/benchmarks/make_lora_adapter.py | 46 +-- scripts/benchmarks/qwen3_grpo_notebook.py | 289 ++++++-------- scripts/benchmarks/qwen3_grpo_unified.py | 300 ++++++--------- .../benchmarks/results/grpo_equivalence.md | 89 +++++ .../results/stats/cb_sync_smoke.json | 15 + .../results/stats/grpo_cb_paged_10.json | 352 +++++++++++++++++ .../stats/grpo_cb_paged_10.summary.json | 64 ++++ .../stats/grpo_unsloth_fi_false_10.json | 362 ++++++++++++++++++ .../grpo_unsloth_fi_false_10.summary.json | 75 ++++ .../results/stats/grpo_vllm_10.json | 362 ++++++++++++++++++ .../results/stats/grpo_vllm_10.summary.json | 75 ++++ 13 files changed, 1815 insertions(+), 569 deletions(-) create mode 100644 scripts/benchmarks/results/grpo_equivalence.md create mode 100644 scripts/benchmarks/results/stats/cb_sync_smoke.json create mode 100644 scripts/benchmarks/results/stats/grpo_cb_paged_10.json create mode 100644 scripts/benchmarks/results/stats/grpo_cb_paged_10.summary.json create mode 100644 scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.json create mode 100644 scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.summary.json create mode 100644 scripts/benchmarks/results/stats/grpo_vllm_10.json create mode 100644 scripts/benchmarks/results/stats/grpo_vllm_10.summary.json diff --git a/scripts/benchmarks/cb_sync_driver.py b/scripts/benchmarks/cb_sync_driver.py index afafaa4705..0e8da56e7c 100644 --- a/scripts/benchmarks/cb_sync_driver.py +++ b/scripts/benchmarks/cb_sync_driver.py @@ -55,7 +55,6 @@ from transformers.generation.continuous_batching.scheduler import FIFOScheduler @dataclass class CBSyncConfig: """Tunables for the sync driver.""" - max_new_tokens: int = 512 use_cuda_graph: bool = True # Number of eager warmup steps before capturing a CUDA graph. @@ -68,7 +67,7 @@ class CBSyncConfig: max_batch_tokens: int = 8192 num_blocks: int = 8192 # Progress callback (step_index, tokens_produced_total) -> None. - on_step: Optional[callable] = field(default = None) + on_step: Optional[callable] = field(default=None) class SyncCBDriver: @@ -80,12 +79,8 @@ class SyncCBDriver: finished. """ - def __init__( - self, - model: torch.nn.Module, - generation_config: GenerationConfig, - cfg: CBSyncConfig, - ): + def __init__(self, model: torch.nn.Module, generation_config: GenerationConfig, + cfg: CBSyncConfig): self.model = model.eval() self.cfg = cfg # Force-greedy + upper-bound overrides on a copy. @@ -105,11 +100,11 @@ class SyncCBDriver: # We reuse the Manager's methods but never call `.start()`. Its # constructor builds: logit processor, do_sample flag, etc. self.manager = ContinuousBatchingManager( - model = self.model, - generation_config = gc, - manual_eviction = False, - streaming = False, - slice_inputs = False, # fixed-shape views -> CUDA-graph safe + model=self.model, + generation_config=gc, + manual_eviction=False, + streaming=False, + slice_inputs=False, # fixed-shape views -> CUDA-graph safe ) # The manager's `use_cuda_graph` is checked inside `warmup()`, but its # `__init__` refuses to set it. Set it directly now that we bypass @@ -123,7 +118,7 @@ class SyncCBDriver: gc, self.model.device, self.model.dtype, - tp_size = getattr(self.model, "_tp_size", None), + tp_size=getattr(self.model, "_tp_size", None), ) self.batch_processor = ContinuousBatchProcessor( self.cache, @@ -134,10 +129,10 @@ class SyncCBDriver: self.manager.stop_event, self.model.device, self.model.dtype, - FIFOScheduler(self.cache, manual_eviction = False), - streaming = False, - manual_eviction = False, - slice_inputs = False, + FIFOScheduler(self.cache), + streaming=False, + manual_eviction=False, + slice_inputs=False, ) self.manager.batch_processor = self.batch_processor self._graph: Optional[torch.cuda.CUDAGraph] = None @@ -153,13 +148,13 @@ class SyncCBDriver: for _ in range(self.cfg.warmup_steps): self.manager._generation_step(self.batch_processor) torch.cuda.synchronize() - stream = torch.cuda.Stream(device = self.model.device) + stream = torch.cuda.Stream(device=self.model.device) stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): self.manager._generation_step(self.batch_processor) torch.cuda.current_stream().wait_stream(stream) self._graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(self._graph, stream = stream): + with torch.cuda.graph(self._graph, stream=stream): self.manager._generation_step(self.batch_processor) else: self._graph.replay() @@ -168,12 +163,24 @@ class SyncCBDriver: """Run the decode loop until every request finishes. Returns a dict {request_id: generated_token_ids}.""" results: dict[str, list[int]] = {} - while self.batch_processor.has_pending_requests(): + # prepare_next_batch drains self.input_queue into the scheduler; we + # have to call it at least once before has_pending_requests() can + # return True. Loop until both the input_queue is empty AND the + # scheduler has nothing queued/active. + while True: + input_empty = self.manager.input_queue.empty() + nothing_scheduled = not self.batch_processor.has_pending_requests() + if input_empty and nothing_scheduled: + break # 1. CPU: schedule the next batch (prepare_next_batch reads the # input_queue, packs shapes). if torch.cuda.is_available(): torch.cuda.synchronize() if not self.batch_processor.prepare_next_batch(): + # prepare_next_batch returns False if both the input queue + # drained empty AND the scheduler has no active requests. If + # we reach here with items still in input_queue, something is + # wrong -- bail to avoid an infinite loop. break # 2. GPU: forward (graphed on decode steps, eager on prefill). if self.cfg.use_cuda_graph and self._is_pure_decode(): @@ -211,20 +218,14 @@ class SyncCBDriver: Shape consistency between decodes is what makes the graph replayable. """ try: - return ( - self.batch_processor.total_query_length - == self.batch_processor.total_batch_size - ) + return (self.batch_processor.total_query_length + == self.batch_processor.total_batch_size) except Exception: return False def _produced(self) -> int: - return sum( - len(r.generated_tokens) - for r in getattr( - self.batch_processor.scheduler, "active_requests", {} - ).values() - ) + return sum(len(r.generated_tokens) for r + in getattr(self.batch_processor.scheduler, "active_requests", {}).values()) def close(self): # Caches hold GPU memory; free them explicitly. @@ -234,12 +235,9 @@ class SyncCBDriver: self.manager.batch_processor = None -def cb_sync_generate( - model: torch.nn.Module, - generation_config: GenerationConfig, - prompt_ids_list: list[list[int]], - cfg: CBSyncConfig, -) -> dict[str, list[int]]: +def cb_sync_generate(model: torch.nn.Module, generation_config: GenerationConfig, + prompt_ids_list: list[list[int]], + cfg: CBSyncConfig) -> dict[str, list[int]]: """One-shot entrypoint: build a driver, submit, drain, close. Matches the semantics of `model.generate_batch(...)` but on the main @@ -265,18 +263,17 @@ if __name__ == "__main__": sys.path.insert(0, str(HERE)) import flash_attn_fa4_shim # noqa: E402 - flash_attn_fa4_shim.apply() parser = argparse.ArgumentParser() - parser.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base") - parser.add_argument("--n_prompts", type = int, default = 32) - parser.add_argument("--max_new_tokens", type = int, default = 512) - parser.add_argument("--attn_impl", default = "paged_attention") - parser.add_argument("--use_cuda_graph", action = "store_true") - parser.add_argument("--max_batch_tokens", type = int, default = 8192) - parser.add_argument("--num_blocks", type = int, default = 8192) - parser.add_argument("--stats_path", required = True) + parser.add_argument("--model_name", default="unsloth/Qwen3-4B-Base") + parser.add_argument("--n_prompts", type=int, default=32) + parser.add_argument("--max_new_tokens", type=int, default=512) + parser.add_argument("--attn_impl", default="paged_attention") + parser.add_argument("--use_cuda_graph", action="store_true") + parser.add_argument("--max_batch_tokens", type=int, default=8192) + parser.add_argument("--num_blocks", type=int, default=8192) + parser.add_argument("--stats_path", required=True) args = parser.parse_args() from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig @@ -285,49 +282,36 @@ if __name__ == "__main__": if tok.pad_token is None: tok.pad_token = tok.eos_token model = AutoModelForCausalLM.from_pretrained( - args.model_name, - dtype = torch.bfloat16, - attn_implementation = args.attn_impl, + args.model_name, dtype=torch.bfloat16, + attn_implementation=args.attn_impl, ).to("cuda") model.eval() from unsloth_grpo_common import ( - SYSTEM_PROMPT, - apply_chat_template_to_tokenizer, + SYSTEM_PROMPT, apply_chat_template_to_tokenizer, ) from datasets import load_dataset - apply_chat_template_to_tokenizer(tok) - ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train") - ds = ds.shuffle(seed = 3407).select(range(args.n_prompts)) - messages = [ - [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": x["prompt"]}, - ] - for x in ds - ] - prompt_ids = [ - tok.apply_chat_template(m, add_generation_prompt = True, tokenize = True) - for m in messages - ] + ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train") + ds = ds.shuffle(seed=3407).select(range(args.n_prompts)) + messages = [[{"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": x["prompt"]}] for x in ds] + prompt_ids = [tok.apply_chat_template(m, add_generation_prompt=True, tokenize=True) + for m in messages] gc = GenerationConfig( - max_new_tokens = args.max_new_tokens, - do_sample = False, - pad_token_id = tok.pad_token_id, - bos_token_id = tok.bos_token_id, - eos_token_id = tok.eos_token_id, - use_cache = True, + max_new_tokens=args.max_new_tokens, do_sample=False, + pad_token_id=tok.pad_token_id, bos_token_id=tok.bos_token_id, + eos_token_id=tok.eos_token_id, use_cache=True, ) cfg = CBSyncConfig( - max_new_tokens = args.max_new_tokens, - use_cuda_graph = args.use_cuda_graph, - max_batch_tokens = args.max_batch_tokens, - num_blocks = args.num_blocks, - eos_token_id = tok.eos_token_id, - pad_token_id = tok.pad_token_id or tok.eos_token_id, + max_new_tokens=args.max_new_tokens, + use_cuda_graph=args.use_cuda_graph, + max_batch_tokens=args.max_batch_tokens, + num_blocks=args.num_blocks, + eos_token_id=tok.eos_token_id, + pad_token_id=tok.pad_token_id or tok.eos_token_id, ) torch.cuda.reset_peak_memory_stats() @@ -358,7 +342,7 @@ if __name__ == "__main__": "max_new_tokens": args.max_new_tokens, "peak_memory_gb": torch.cuda.max_memory_allocated() / 1024**3, } - os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True) + os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok=True) with open(args.stats_path, "w") as f: - json.dump(out, f, indent = 2) - print(json.dumps(out, indent = 2)) + json.dump(out, f, indent=2) + print(json.dumps(out, indent=2)) diff --git a/scripts/benchmarks/cb_vs_vllm_generation.py b/scripts/benchmarks/cb_vs_vllm_generation.py index fd4fc715bb..795a8340fa 100644 --- a/scripts/benchmarks/cb_vs_vllm_generation.py +++ b/scripts/benchmarks/cb_vs_vllm_generation.py @@ -50,8 +50,8 @@ def build_prompts(tokenizer, n_prompts): from datasets import load_dataset apply_chat_template_to_tokenizer(tokenizer) - ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train") - ds = ds.shuffle(seed = 3407).select(range(n_prompts)) + ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train") + ds = ds.shuffle(seed=3407).select(range(n_prompts)) messages = [ [ {"role": "system", "content": SYSTEM_PROMPT}, @@ -60,11 +60,11 @@ def build_prompts(tokenizer, n_prompts): for x in ds ] prompts_text = [ - tokenizer.apply_chat_template(m, add_generation_prompt = True, tokenize = False) + tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=False) for m in messages ] prompt_ids = [ - tokenizer.apply_chat_template(m, add_generation_prompt = True, tokenize = True) + tokenizer.apply_chat_template(m, add_generation_prompt=True, tokenize=True) for m in messages ] return prompts_text, prompt_ids @@ -75,37 +75,35 @@ def run_vllm(args): from unsloth import FastLanguageModel model, tokenizer = FastLanguageModel.from_pretrained( - model_name = args.model_name, - max_seq_length = args.max_seq_length, - load_in_4bit = False, - fast_inference = True, - max_lora_rank = 32, - gpu_memory_utilization = args.gpu_memory_utilization, + model_name=args.model_name, + max_seq_length=args.max_seq_length, + load_in_4bit=False, + fast_inference=True, + max_lora_rank=32, + gpu_memory_utilization=args.gpu_memory_utilization, ) prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts) lora_request = None if args.lora_adapter: from vllm.lora.request import LoRARequest - lora_request = LoRARequest("fresh", 1, str(Path(args.lora_adapter).resolve())) from vllm import SamplingParams - sp = SamplingParams( - temperature = args.temperature, - top_p = args.top_p, - min_p = args.min_p, - top_k = args.top_k, - seed = 3407, - max_tokens = args.max_new_tokens, - stop = [tokenizer.eos_token], - include_stop_str_in_output = True, + temperature=args.temperature, + top_p=args.top_p, + min_p=args.min_p, + top_k=args.top_k, + seed=3407, + max_tokens=args.max_new_tokens, + stop=[tokenizer.eos_token], + include_stop_str_in_output=True, ) # Warmup on 16 prompts then discard. warmup_text = prompts_text[:16] - _ = model.fast_generate(warmup_text, sampling_params = sp, lora_request = lora_request) + _ = model.fast_generate(warmup_text, sampling_params=sp, lora_request=lora_request) torch.cuda.synchronize() n_prompt_tokens = sum(len(p) for p in prompt_ids) @@ -116,7 +114,7 @@ def run_vllm(args): torch.cuda.synchronize() t0 = time.perf_counter() outputs = model.fast_generate( - prompts_text, sampling_params = sp, lora_request = lora_request + prompts_text, sampling_params=sp, lora_request=lora_request ) torch.cuda.synchronize() wall_times.append(time.perf_counter() - t0) @@ -124,11 +122,7 @@ def run_vllm(args): last_outputs = outputs med = sorted(wall_times)[len(wall_times) // 2] - sample_texts = ( - [o.outputs[0].text[:200] for o in (last_outputs[:3] or [])] - if last_outputs - else [] - ) + sample_texts = [o.outputs[0].text[:200] for o in (last_outputs[:3] or [])] if last_outputs else [] return { "backend": "vllm", "lora_adapter": args.lora_adapter, @@ -157,17 +151,16 @@ def run_tpaged(args): tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( args.model_name, - dtype = torch.bfloat16, - attn_implementation = args.attn_impl, + dtype=torch.bfloat16, + attn_implementation=args.attn_impl, ).to("cuda") model.eval() if args.lora_adapter: from peft import PeftModel - # NOTE: no merge_adapter -- we measure LoRA-active inference. model = PeftModel.from_pretrained( - model, str(Path(args.lora_adapter).resolve()), is_trainable = False + model, str(Path(args.lora_adapter).resolve()), is_trainable=False ) model.eval() @@ -176,16 +169,16 @@ def run_tpaged(args): prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts) gen_config = GenerationConfig( - max_new_tokens = args.max_new_tokens, - do_sample = True, - temperature = args.temperature, - top_p = args.top_p, - min_p = args.min_p, - top_k = args.top_k, - pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id, - bos_token_id = tokenizer.bos_token_id, - eos_token_id = tokenizer.eos_token_id, - use_cache = True, + max_new_tokens=args.max_new_tokens, + do_sample=True, + temperature=args.temperature, + top_p=args.top_p, + min_p=args.min_p, + top_k=args.top_k, + pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id, + bos_token_id=tokenizer.bos_token_id, + eos_token_id=tokenizer.eos_token_id, + use_cache=True, ) gen_config.max_batch_tokens = args.max_batch_tokens gen_config.num_blocks = args.num_blocks @@ -195,9 +188,7 @@ def run_tpaged(args): warmup_ids = prompt_ids[:16] with torch.inference_mode(): - _ = model.generate_batch( - warmup_ids, generation_config = gen_config, progress_bar = False - ) + _ = model.generate_batch(warmup_ids, generation_config=gen_config, progress_bar=False) torch.cuda.synchronize() n_prompt_tokens = sum(len(p) for p in prompt_ids) @@ -209,7 +200,7 @@ def run_tpaged(args): t0 = time.perf_counter() with torch.inference_mode(): outputs = model.generate_batch( - prompt_ids, generation_config = gen_config, progress_bar = False + prompt_ids, generation_config=gen_config, progress_bar=False ) torch.cuda.synchronize() wall_times.append(time.perf_counter() - t0) @@ -221,7 +212,7 @@ def run_tpaged(args): if last_outputs is not None: for k in list(last_outputs.keys())[:3]: toks = last_outputs[k].generated_tokens - sample_texts.append(tokenizer.decode(toks, skip_special_tokens = False)[:200]) + sample_texts.append(tokenizer.decode(toks, skip_special_tokens=False)[:200]) med = sorted(wall_times)[len(wall_times) // 2] return { @@ -257,40 +248,33 @@ def run_unsloth_fi_false(args): from unsloth import FastLanguageModel model, tokenizer = FastLanguageModel.from_pretrained( - model_name = args.model_name, - max_seq_length = args.max_seq_length, - load_in_4bit = False, - fast_inference = False, - max_lora_rank = 32, + model_name=args.model_name, + max_seq_length=args.max_seq_length, + load_in_4bit=False, + fast_inference=False, + max_lora_rank=32, ) # Attach LoRA rank 32 the same way the GRPO notebook does. model = FastLanguageModel.get_peft_model( model, - r = 32, - target_modules = [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", + r=32, + target_modules=[ + "q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj", ], - lora_alpha = 64, - use_gradient_checkpointing = "unsloth", - random_state = 3407, + lora_alpha=64, + use_gradient_checkpointing="unsloth", + random_state=3407, ) # Optional: overlay a shared adapter so weights match other backends. if args.lora_adapter: from safetensors import safe_open - adapter_file = Path(args.lora_adapter).resolve() / "adapter_model.safetensors" loaded_tensors = {} - with safe_open(str(adapter_file), framework = "pt") as f: + with safe_open(str(adapter_file), framework="pt") as f: for key in f.keys(): loaded_tensors[key] = f.get_tensor(key) - # Both PEFT and Unsloth's `get_peft_model` produce parameter names with # `base_model.model.` prefix plus `.lora_{A,B}.default.weight`. Build a # normalized (core-path) -> param map, then match by core path only. @@ -298,12 +282,10 @@ def run_unsloth_fi_false(args): n = name for pref in ("base_model.model.", "model."): if n.startswith(pref): - n = n[len(pref) :] + n = n[len(pref):] n = n.replace(".lora_A.default.", ".lora_A.").replace( - ".lora_B.default.", ".lora_B." - ) + ".lora_B.default.", ".lora_B.") return n - own_by_core = {} for n, p in model.named_parameters(): if "lora_" in n: @@ -317,10 +299,8 @@ def run_unsloth_fi_false(args): own.data.copy_(tensor.to(own.device, own.dtype)) matched += 1 break - print( - f"[unsloth_fi_false] LoRA weight sync matched {matched} tensors " - f"(out of {len(loaded_tensors)} adapter entries)." - ) + print(f"[unsloth_fi_false] LoRA weight sync matched {matched} tensors " + f"(out of {len(loaded_tensors)} adapter entries).") FastLanguageModel.for_inference(model) @@ -328,29 +308,28 @@ def run_unsloth_fi_false(args): # `model.generate` accepts batched input_ids; pad to max length. from transformers import GenerationConfig - if tokenizer.padding_side != "left": tokenizer.padding_side = "left" # decoder needs left padding if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token gen_config = GenerationConfig( - max_new_tokens = args.max_new_tokens, - do_sample = True, - temperature = args.temperature, - top_p = args.top_p, - min_p = args.min_p, - top_k = args.top_k, - pad_token_id = tokenizer.pad_token_id, - bos_token_id = tokenizer.bos_token_id, - eos_token_id = tokenizer.eos_token_id, - use_cache = True, + max_new_tokens=args.max_new_tokens, + do_sample=True, + temperature=args.temperature, + top_p=args.top_p, + min_p=args.min_p, + top_k=args.top_k, + pad_token_id=tokenizer.pad_token_id, + bos_token_id=tokenizer.bos_token_id, + eos_token_id=tokenizer.eos_token_id, + use_cache=True, ) def _batched_generate(texts): - batch = tokenizer(texts, return_tensors = "pt", padding = True).to("cuda") + batch = tokenizer(texts, return_tensors="pt", padding=True).to("cuda") with torch.inference_mode(): - out = model.generate(**batch, generation_config = gen_config) + out = model.generate(**batch, generation_config=gen_config) prompt_len = batch["input_ids"].shape[1] return out, prompt_len @@ -371,9 +350,7 @@ def run_unsloth_fi_false(args): wall_times.append(time.perf_counter() - t0) # Count generated tokens past prompt_len per sequence (subtract any # trailing pad-only tail by comparing against EOS). - total_decoded = int( - (out_ids[:, prompt_len:] != tokenizer.pad_token_id).sum().item() - ) + total_decoded = int((out_ids[:, prompt_len:] != tokenizer.pad_token_id).sum().item()) last_out_ids = out_ids last_prompt_len = prompt_len @@ -381,11 +358,8 @@ def run_unsloth_fi_false(args): sample_texts = [] if last_out_ids is not None: for i in range(min(3, last_out_ids.shape[0])): - sample_texts.append( - tokenizer.decode( - last_out_ids[i, last_prompt_len:], skip_special_tokens = False - )[:200] - ) + sample_texts.append(tokenizer.decode( + last_out_ids[i, last_prompt_len:], skip_special_tokens=False)[:200]) return { "backend": "unsloth_fi_false", @@ -404,35 +378,30 @@ def run_unsloth_fi_false(args): def parse_args(): p = argparse.ArgumentParser() - p.add_argument( - "--backend", choices = ["vllm", "tpaged", "unsloth_fi_false"], required = True - ) - p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base") - p.add_argument("--max_seq_length", type = int, default = 2048) - p.add_argument("--n_prompts", type = int, default = 32) - p.add_argument("--n_rounds", type = int, default = 2) - p.add_argument("--max_new_tokens", type = int, default = 512) - p.add_argument("--gpu_memory_utilization", type = float, default = 0.8) - p.add_argument("--attn_impl", default = "sdpa") - p.add_argument("--max_batch_tokens", type = int, default = 8192) - p.add_argument("--num_blocks", type = int, default = 16384) - p.add_argument("--persistent_cb", action = "store_true") - p.add_argument( - "--lora_adapter", - default = None, - help = "Path to a PEFT adapter (rank 32) applied in every backend.", - ) - p.add_argument("--temperature", type = float, default = 0.1) - p.add_argument("--top_p", type = float, default = 0.97) - p.add_argument("--min_p", type = float, default = 0.5) - p.add_argument("--top_k", type = int, default = 5) - p.add_argument("--stats_path", required = True) + p.add_argument("--backend", choices=["vllm", "tpaged", "unsloth_fi_false"], required=True) + p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base") + p.add_argument("--max_seq_length", type=int, default=2048) + p.add_argument("--n_prompts", type=int, default=32) + p.add_argument("--n_rounds", type=int, default=2) + p.add_argument("--max_new_tokens", type=int, default=512) + p.add_argument("--gpu_memory_utilization", type=float, default=0.8) + p.add_argument("--attn_impl", default="sdpa") + p.add_argument("--max_batch_tokens", type=int, default=8192) + p.add_argument("--num_blocks", type=int, default=16384) + p.add_argument("--persistent_cb", action="store_true") + p.add_argument("--lora_adapter", default=None, + help="Path to a PEFT adapter (rank 32) applied in every backend.") + p.add_argument("--temperature", type=float, default=0.1) + p.add_argument("--top_p", type=float, default=0.97) + p.add_argument("--min_p", type=float, default=0.5) + p.add_argument("--top_k", type=int, default=5) + p.add_argument("--stats_path", required=True) return p.parse_args() def main(): args = parse_args() - os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True) + os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok=True) torch.cuda.reset_peak_memory_stats() if args.backend == "vllm": @@ -450,8 +419,8 @@ def main(): "top_k": args.top_k, } with open(args.stats_path, "w") as f: - json.dump(out, f, indent = 2) - print(json.dumps(out, indent = 2)) + json.dump(out, f, indent=2) + print(json.dumps(out, indent=2)) os._exit(0) diff --git a/scripts/benchmarks/make_lora_adapter.py b/scripts/benchmarks/make_lora_adapter.py index ef646e94e8..82df17a2da 100644 --- a/scripts/benchmarks/make_lora_adapter.py +++ b/scripts/benchmarks/make_lora_adapter.py @@ -20,16 +20,16 @@ from pathlib import Path def parse_args(): p = argparse.ArgumentParser() - p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base") - p.add_argument("--output", default = "outputs/lora_rank32_fresh") - p.add_argument("--rank", type = int, default = 32) + p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base") + p.add_argument("--output", default="outputs/lora_rank32_fresh") + p.add_argument("--rank", type=int, default=32) return p.parse_args() def main(): args = parse_args() out_dir = Path(args.output).resolve() - out_dir.mkdir(parents = True, exist_ok = True) + out_dir.mkdir(parents=True, exist_ok=True) # Use vanilla HF -- PEFT's save_pretrained yields the canonical # adapter_config.json + adapter_model.safetensors that vLLM's LoRARequest @@ -45,23 +45,16 @@ def main(): # bf16 base; we only need structure + save. Keep on CPU to avoid a GPU load # just for `save_pretrained`. print(f"[make_lora_adapter] Loading {args.model_name} on CPU...") - model = AutoModelForCausalLM.from_pretrained(args.model_name, dtype = torch.bfloat16) + model = AutoModelForCausalLM.from_pretrained(args.model_name, dtype=torch.bfloat16) peft_cfg = LoraConfig( - r = args.rank, - lora_alpha = args.rank * 2, - target_modules = [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - bias = "none", - task_type = "CAUSAL_LM", - lora_dropout = 0.0, + r=args.rank, + lora_alpha=args.rank * 2, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj"], + bias="none", + task_type="CAUSAL_LM", + lora_dropout=0.0, ) peft_model = get_peft_model(model, peft_cfg) peft_model.print_trainable_parameters() @@ -73,31 +66,26 @@ def main(): with torch.no_grad(): for name, p in peft_model.named_parameters(): if "lora_B" in name: - p.normal_(mean = 0.0, std = 1e-4) + p.normal_(mean=0.0, std=1e-4) n_reinit += 1 - print( - f"[make_lora_adapter] Reinitialized {n_reinit} lora_B matrices with tiny gaussian." - ) + print(f"[make_lora_adapter] Reinitialized {n_reinit} lora_B matrices with tiny gaussian.") peft_model.save_pretrained(str(out_dir)) tok.save_pretrained(str(out_dir)) # Sanity: verify safetensors file present and non-trivial. from safetensors import safe_open - st_path = out_dir / "adapter_model.safetensors" n_zero_tensors = 0 n_tensors = 0 - with safe_open(str(st_path), framework = "pt") as f: + with safe_open(str(st_path), framework="pt") as f: for key in f.keys(): t = f.get_tensor(key) n_tensors += 1 if (t == 0).all().item(): n_zero_tensors += 1 - print( - f"[make_lora_adapter] Wrote {n_tensors} tensors to {st_path} " - f"({n_zero_tensors} all-zero)." - ) + print(f"[make_lora_adapter] Wrote {n_tensors} tensors to {st_path} " + f"({n_zero_tensors} all-zero).") print(f"[make_lora_adapter] Adapter saved to {out_dir}") diff --git a/scripts/benchmarks/qwen3_grpo_notebook.py b/scripts/benchmarks/qwen3_grpo_notebook.py index de01cd6344..b231acf9e5 100644 --- a/scripts/benchmarks/qwen3_grpo_notebook.py +++ b/scripts/benchmarks/qwen3_grpo_notebook.py @@ -33,31 +33,28 @@ for p in (HERE, WORKSPACE_ROOT): def parse_args(): p = argparse.ArgumentParser() - p.add_argument("--stats_path", default = "logs/notebook_ref_10.json") - p.add_argument("--output_dir", default = "outputs/notebook_ref_10") - p.add_argument("--max_steps", type = int, default = 10) - p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base") - p.add_argument("--max_seq_length", type = int, default = 2048) - p.add_argument("--lora_rank", type = int, default = 32) - p.add_argument("--gpu_memory_utilization", type = float, default = 0.85) - p.add_argument("--num_generations", type = int, default = 4) - p.add_argument("--per_device_train_batch_size", type = int, default = 1) - p.add_argument("--temperature", type = float, default = 0.1) - p.add_argument("--top_p", type = float, default = 0.97) - p.add_argument("--min_p", type = float, default = 0.5) - p.add_argument("--top_k", type = int, default = 5) - p.add_argument( - "--skip_sft_pre_finetune", - action = "store_true", - help = "Skip the format-priming SFT stage; go straight to GRPO.", - ) + p.add_argument("--stats_path", default="logs/notebook_ref_10.json") + p.add_argument("--output_dir", default="outputs/notebook_ref_10") + p.add_argument("--max_steps", type=int, default=10) + p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base") + p.add_argument("--max_seq_length", type=int, default=2048) + p.add_argument("--lora_rank", type=int, default=32) + p.add_argument("--gpu_memory_utilization", type=float, default=0.85) + p.add_argument("--num_generations", type=int, default=4) + p.add_argument("--per_device_train_batch_size", type=int, default=1) + p.add_argument("--temperature", type=float, default=0.1) + p.add_argument("--top_p", type=float, default=0.97) + p.add_argument("--min_p", type=float, default=0.5) + p.add_argument("--top_k", type=int, default=5) + p.add_argument("--skip_sft_pre_finetune", action="store_true", + help="Skip the format-priming SFT stage; go straight to GRPO.") return p.parse_args() def main(): args = parse_args() - os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True) - os.makedirs(args.output_dir, exist_ok = True) + os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok=True) + os.makedirs(args.output_dir, exist_ok=True) # Import order matters: unsloth must come before transformers/trl. os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1") @@ -65,28 +62,23 @@ def main(): import torch # noqa: E402 model, tokenizer = FastLanguageModel.from_pretrained( - model_name = args.model_name, - max_seq_length = args.max_seq_length, - load_in_4bit = False, - fast_inference = True, - max_lora_rank = args.lora_rank, - gpu_memory_utilization = args.gpu_memory_utilization, + model_name=args.model_name, + max_seq_length=args.max_seq_length, + load_in_4bit=False, + fast_inference=True, + max_lora_rank=args.lora_rank, + gpu_memory_utilization=args.gpu_memory_utilization, ) model = FastLanguageModel.get_peft_model( model, - r = args.lora_rank, - target_modules = [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", + r=args.lora_rank, + target_modules=[ + "q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj", ], - lora_alpha = args.lora_rank * 2, - use_gradient_checkpointing = "unsloth", - random_state = 3407, + lora_alpha=args.lora_rank * 2, + use_gradient_checkpointing="unsloth", + random_state=3407, ) reasoning_start = "" @@ -127,29 +119,16 @@ def main(): import numpy as np if not args.skip_sft_pre_finetune: - sft_ds = load_dataset("unsloth/OpenMathReasoning-mini", split = "cot") - sft_df = sft_ds.to_pandas()[ - ["expected_answer", "problem", "generated_solution"] - ] - is_number = pd.to_numeric( - pd.Series(sft_df["expected_answer"]), errors = "coerce" - ).notnull() + sft_ds = load_dataset("unsloth/OpenMathReasoning-mini", split="cot") + sft_df = sft_ds.to_pandas()[["expected_answer", "problem", "generated_solution"]] + is_number = pd.to_numeric(pd.Series(sft_df["expected_answer"]), errors="coerce").notnull() sft_df = sft_df.iloc[np.where(is_number)[0]] def format_dataset(x): - thoughts = ( - x["generated_solution"] - .replace("", "") - .replace("", "") - .strip() - ) + thoughts = x["generated_solution"].replace("", "").replace("", "").strip() final_prompt = ( - reasoning_start - + thoughts - + reasoning_end - + solution_start - + x["expected_answer"] - + solution_end + reasoning_start + thoughts + reasoning_end + + solution_start + x["expected_answer"] + solution_end ) return [ {"role": "system", "content": system_prompt}, @@ -157,69 +136,61 @@ def main(): {"role": "assistant", "content": final_prompt}, ] - sft_df["Messages"] = sft_df.apply(format_dataset, axis = 1) - sft_df["N"] = sft_df["Messages"].apply( - lambda m: len(tokenizer.apply_chat_template(m)) - ) + sft_df["Messages"] = sft_df.apply(format_dataset, axis=1) + sft_df["N"] = sft_df["Messages"].apply(lambda m: len(tokenizer.apply_chat_template(m))) sft_df = sft_df.loc[sft_df["N"] <= args.max_seq_length / 2].copy() sft_df["text"] = tokenizer.apply_chat_template( - sft_df["Messages"].values.tolist(), tokenize = False + sft_df["Messages"].values.tolist(), tokenize=False ) sft_dataset = Dataset.from_pandas(sft_df) from trl import SFTTrainer, SFTConfig - sft_trainer = SFTTrainer( - model = model, - tokenizer = tokenizer, - train_dataset = sft_dataset, - args = SFTConfig( - dataset_text_field = "text", - per_device_train_batch_size = 1, - gradient_accumulation_steps = 1, - warmup_steps = 5, - num_train_epochs = 2, - learning_rate = 2e-4, - logging_steps = 5, - optim = "adamw_8bit", - weight_decay = 0.001, - lr_scheduler_type = "linear", - seed = 3407, - report_to = "none", - output_dir = os.path.join(args.output_dir, "sft"), + model=model, + tokenizer=tokenizer, + train_dataset=sft_dataset, + args=SFTConfig( + dataset_text_field="text", + per_device_train_batch_size=1, + gradient_accumulation_steps=1, + warmup_steps=5, + num_train_epochs=2, + learning_rate=2e-4, + logging_steps=5, + optim="adamw_8bit", + weight_decay=0.001, + lr_scheduler_type="linear", + seed=3407, + report_to="none", + output_dir=os.path.join(args.output_dir, "sft"), ), ) sft_trainer.train() del sft_dataset, sft_df, sft_ds, sft_trainer torch.cuda.empty_cache() import gc - gc.collect() # --- GRPO stage ----------------------------------------------------------- - dataset = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train") - dataset = dataset.map( - lambda x: { - "prompt": [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": x["prompt"]}, - ], - "answer": x["solution"], - } - ) + dataset = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split="train") + dataset = dataset.map(lambda x: { + "prompt": [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": x["prompt"]}, + ], + "answer": x["solution"], + }) - solution_end_regex = ( - r"[\s]{0,}" + "(?:" + re.escape(tokenizer.eos_token) + ")?" - ) + 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, + flags=re.MULTILINE | re.DOTALL, ) match_numbers = re.compile( solution_start + r".*?[\s]{0,}([-]?[\d\.\,]{1,})", - flags = re.MULTILINE | re.DOTALL, + flags=re.MULTILINE | re.DOTALL, ) def match_format_exactly(completions, **kwargs): @@ -291,12 +262,10 @@ def main(): # Filter long prompts. tokenized = dataset.map( - lambda x: { - "tokens": tokenizer.apply_chat_template( - x["prompt"], add_generation_prompt = True, tokenize = True - ) - }, - batched = False, + lambda x: {"tokens": tokenizer.apply_chat_template( + x["prompt"], add_generation_prompt=True, tokenize=True + )}, + batched=False, ) tokenized = tokenized.map(lambda x: {"L": len(x["tokens"])}) maximum_length = int(np.quantile(tokenized["L"], 0.9)) @@ -308,63 +277,60 @@ def main(): max_completion_length = args.max_seq_length - max_prompt_length from vllm import SamplingParams - vllm_sampling_params = SamplingParams( - temperature = args.temperature, - top_p = args.top_p, - min_p = args.min_p, - top_k = args.top_k, - seed = 3407, - stop = [tokenizer.eos_token], - include_stop_str_in_output = True, + temperature=args.temperature, + top_p=args.top_p, + min_p=args.min_p, + top_k=args.top_k, + seed=3407, + stop=[tokenizer.eos_token], + include_stop_str_in_output=True, ) from trl import GRPOConfig, GRPOTrainer - training_args = GRPOConfig( - vllm_sampling_params = vllm_sampling_params, - temperature = args.temperature, - top_p = args.top_p, - top_k = args.top_k, - learning_rate = 5e-6, - weight_decay = 0.001, - warmup_ratio = 0.1, - lr_scheduler_type = "linear", - optim = "adamw_8bit", - logging_steps = 1, - per_device_train_batch_size = args.per_device_train_batch_size, - gradient_accumulation_steps = 1, - num_generations = args.num_generations, - max_prompt_length = max_prompt_length, - max_completion_length = max_completion_length, - max_steps = args.max_steps, - save_steps = args.max_steps + 1, - report_to = "none", - output_dir = args.output_dir, - seed = 3407, + vllm_sampling_params=vllm_sampling_params, + temperature=args.temperature, + top_p=args.top_p, + top_k=args.top_k, + learning_rate=5e-6, + weight_decay=0.001, + warmup_ratio=0.1, + lr_scheduler_type="linear", + optim="adamw_8bit", + logging_steps=1, + per_device_train_batch_size=args.per_device_train_batch_size, + gradient_accumulation_steps=1, + num_generations=args.num_generations, + max_prompt_length=max_prompt_length, + max_completion_length=max_completion_length, + max_steps=args.max_steps, + save_steps=args.max_steps + 1, + report_to="none", + output_dir=args.output_dir, + seed=3407, ) from torch_debugging_utils import StatisticsCallback - stats_cb = StatisticsCallback( - track_loss = True, - track_grad_norm = True, - track_memory = True, - track_tensor_stats = False, # hooks are noisy + slow on GRPO model + track_loss=True, + track_grad_norm=True, + track_memory=True, + track_tensor_stats=False, # hooks are noisy + slow on GRPO model ) trainer = GRPOTrainer( - model = model, - processing_class = tokenizer, - reward_funcs = [ + model=model, + processing_class=tokenizer, + reward_funcs=[ match_format_exactly, match_format_approximately, check_answer, check_numbers, ], - args = training_args, - train_dataset = dataset, - callbacks = [stats_cb], + args=training_args, + train_dataset=dataset, + callbacks=[stats_cb], ) t0 = time.perf_counter() @@ -395,39 +361,30 @@ def main(): "logs_path": args.stats_path, "peak_memory_gb": torch.cuda.max_memory_allocated() / 1024**3, } - print(json.dumps(summary, indent = 2)) + print(json.dumps(summary, indent=2)) # Canonical quick-inference: produce a few generations for the writeup. rollouts = [] try: from vllm import SamplingParams as SP - sp_sample = SP( - temperature = args.temperature, - top_p = args.top_p, - min_p = args.min_p, - top_k = args.top_k, - max_tokens = 256, + temperature=args.temperature, + top_p=args.top_p, + min_p=args.min_p, + top_k=args.top_k, + max_tokens=256, ) probe_prompts = [ - [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": "What is the sqrt of 101?"}, - ], - [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": "If 3x+7 = 22, what is x?"}, - ], - [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": "What is 17 * 13?"}, - ], + [{"role": "system", "content": system_prompt}, + {"role": "user", "content": "What is the sqrt of 101?"}], + [{"role": "system", "content": system_prompt}, + {"role": "user", "content": "If 3x+7 = 22, what is x?"}], + [{"role": "system", "content": system_prompt}, + {"role": "user", "content": "What is 17 * 13?"}], ] - texts = [ - tokenizer.apply_chat_template(p, add_generation_prompt = True, tokenize = False) - for p in probe_prompts - ] - outs = model.fast_generate(texts, sampling_params = sp_sample, lora_request = None) + texts = [tokenizer.apply_chat_template(p, add_generation_prompt=True, tokenize=False) + for p in probe_prompts] + outs = model.fast_generate(texts, sampling_params=sp_sample, lora_request=None) for t, o in zip(texts, outs): rollouts.append({"prompt": t, "completion": o.outputs[0].text}) except Exception as e: diff --git a/scripts/benchmarks/qwen3_grpo_unified.py b/scripts/benchmarks/qwen3_grpo_unified.py index 83122388ed..37799968bb 100644 --- a/scripts/benchmarks/qwen3_grpo_unified.py +++ b/scripts/benchmarks/qwen3_grpo_unified.py @@ -42,40 +42,36 @@ os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1") def parse_args(): p = argparse.ArgumentParser() - p.add_argument( - "--backend", - choices = ["vllm", "unsloth_fi_false", "cb_paged", "cb_sdpa", "naive_trl"], - required = True, - ) - p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base") - p.add_argument("--max_seq_length", type = int, default = 2048) - p.add_argument("--lora_rank", type = int, default = 32) - p.add_argument("--max_steps", type = int, default = 10) - p.add_argument("--num_generations", type = int, default = 4) - p.add_argument("--per_device_train_batch_size", type = int, default = 1) - p.add_argument("--gradient_accumulation_steps", type = int, default = 1) - p.add_argument("--gpu_memory_utilization", type = float, default = 0.75) - p.add_argument("--temperature", type = float, default = 0.1) - p.add_argument("--top_p", type = float, default = 0.97) - p.add_argument("--min_p", type = float, default = 0.5) - p.add_argument("--top_k", type = int, default = 5) - p.add_argument("--learning_rate", type = float, default = 5e-6) - p.add_argument("--max_batch_tokens", type = int, default = 8192) - p.add_argument("--num_blocks", type = int, default = 8192) - p.add_argument("--persistent_cb", action = "store_true") - p.add_argument("--output_dir", required = True) - p.add_argument("--stats_path", required = True) - p.add_argument("--seed", type = int, default = 3407) + p.add_argument("--backend", + choices=["vllm", "unsloth_fi_false", "cb_paged", "cb_sdpa", "naive_trl"], + required=True) + p.add_argument("--model_name", default="unsloth/Qwen3-4B-Base") + p.add_argument("--max_seq_length", type=int, default=2048) + p.add_argument("--lora_rank", type=int, default=32) + p.add_argument("--max_steps", type=int, default=10) + p.add_argument("--num_generations", type=int, default=4) + p.add_argument("--per_device_train_batch_size", type=int, default=1) + p.add_argument("--gradient_accumulation_steps", type=int, default=1) + p.add_argument("--gpu_memory_utilization", type=float, default=0.75) + p.add_argument("--temperature", type=float, default=0.1) + p.add_argument("--top_p", type=float, default=0.97) + p.add_argument("--min_p", type=float, default=0.5) + p.add_argument("--top_k", type=int, default=5) + p.add_argument("--learning_rate", type=float, default=5e-6) + p.add_argument("--max_batch_tokens", type=int, default=8192) + p.add_argument("--num_blocks", type=int, default=8192) + p.add_argument("--persistent_cb", action="store_true") + p.add_argument("--output_dir", required=True) + p.add_argument("--stats_path", required=True) + p.add_argument("--seed", type=int, default=3407) # Phase 4: torch.compile on the training forward. - p.add_argument( - "--compile_mode", - default = None, - choices = [None, "default", "reduce-overhead", "max-autotune-no-cudagraphs"], - help = "If set, torch.compile(model.forward, mode=...) after " - "the trainer is built. vllm backend is excluded; the " - "rollout engine owns its own compile pipeline.", - ) - p.add_argument("--compile_dynamic", action = "store_true", default = True) + p.add_argument("--compile_mode", default=None, + choices=[None, "default", "reduce-overhead", + "max-autotune-no-cudagraphs"], + help="If set, torch.compile(model.forward, mode=...) after " + "the trainer is built. vllm backend is excluded; the " + "rollout engine owns its own compile pipeline.") + p.add_argument("--compile_dynamic", action="store_true", default=True) return p.parse_args() @@ -83,18 +79,10 @@ def _prepare_common(args): """Dataset + rewards are the same for every backend. Always uses the shared chat template and reward funcs from unsloth_grpo_common.""" from unsloth_grpo_common import ( - apply_chat_template_to_tokenizer, - build_dataset, - build_reward_funcs, - build_grpo_kwargs, - ) - - return ( - apply_chat_template_to_tokenizer, - build_dataset, - build_reward_funcs, - build_grpo_kwargs, + apply_chat_template_to_tokenizer, build_dataset, + build_reward_funcs, build_grpo_kwargs, ) + return apply_chat_template_to_tokenizer, build_dataset, build_reward_funcs, build_grpo_kwargs def _make_stats_callback(): @@ -102,12 +90,11 @@ def _make_stats_callback(): grad-norm, memory, and wall time. Reward/KL are picked up from the TRL log dict via `on_log`.""" from torch_debugging_utils import StatisticsCallback - return StatisticsCallback( - track_loss = True, - track_grad_norm = True, - track_memory = True, - track_tensor_stats = False, + track_loss=True, + track_grad_norm=True, + track_memory=True, + track_tensor_stats=False, ) @@ -117,13 +104,10 @@ def _maybe_shim_guided_decoding(): the transformers-paged path. Inject a no-op shim if missing.""" try: import vllm.sampling_params as sp - if not hasattr(sp, "GuidedDecodingParams"): - class _Shim: def __init__(self, *a, **kw): pass - sp.GuidedDecodingParams = _Shim except ImportError: pass @@ -131,34 +115,22 @@ def _maybe_shim_guided_decoding(): def _load_unsloth(args, fast_inference: bool): from unsloth import FastLanguageModel - model, tokenizer = FastLanguageModel.from_pretrained( - model_name = args.model_name, - max_seq_length = args.max_seq_length, - load_in_4bit = False, - fast_inference = fast_inference, - max_lora_rank = args.lora_rank, - **( - {"gpu_memory_utilization": args.gpu_memory_utilization} - if fast_inference - else {} - ), + model_name=args.model_name, + max_seq_length=args.max_seq_length, + load_in_4bit=False, + fast_inference=fast_inference, + max_lora_rank=args.lora_rank, + **({"gpu_memory_utilization": args.gpu_memory_utilization} if fast_inference else {}), ) model = FastLanguageModel.get_peft_model( model, - r = args.lora_rank, - target_modules = [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - lora_alpha = args.lora_rank * 2, - use_gradient_checkpointing = "unsloth", - random_state = args.seed, + r=args.lora_rank, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj"], + lora_alpha=args.lora_rank * 2, + use_gradient_checkpointing="unsloth", + random_state=args.seed, ) return model, tokenizer @@ -174,28 +146,21 @@ def _load_vanilla_hf(args, attn_impl: str): tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( args.model_name, - dtype = torch.bfloat16, - attn_implementation = attn_impl, + dtype=torch.bfloat16, + attn_implementation=attn_impl, ).to("cuda") lora = LoraConfig( - r = args.lora_rank, - lora_alpha = args.lora_rank * 2, - target_modules = [ - "q_proj", - "k_proj", - "v_proj", - "o_proj", - "gate_proj", - "up_proj", - "down_proj", - ], - bias = "none", - task_type = "CAUSAL_LM", + r=args.lora_rank, + lora_alpha=args.lora_rank * 2, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj"], + bias="none", + task_type="CAUSAL_LM", ) model = get_peft_model(model, lora) try: model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs = {"use_reentrant": False} + gradient_checkpointing_kwargs={"use_reentrant": False} ) except TypeError: model.gradient_checkpointing_enable() @@ -205,46 +170,57 @@ def _load_vanilla_hf(args, attn_impl: str): def main(): args = parse_args() - os.makedirs(args.output_dir, exist_ok = True) - os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True) + os.makedirs(args.output_dir, exist_ok=True) + os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok=True) import torch from torch_debugging_utils import set_all_seeds_fast - set_all_seeds_fast(args.seed) # FA4 shim lives here so CB paths dispatch to Blackwell kernels. import flash_attn_fa4_shim # noqa: F401 - flash_attn_fa4_shim.apply() _maybe_shim_guided_decoding() - ( - apply_chat_template_to_tokenizer, - build_dataset, - build_reward_funcs, - build_grpo_kwargs, - ) = _prepare_common(args) + (apply_chat_template_to_tokenizer, build_dataset, + build_reward_funcs, build_grpo_kwargs) = _prepare_common(args) + + # TRL requires `generation_batch_size = pdb * grad_accum * world_size` to + # be divisible by `num_generations`. Unsloth's loader auto-adjusts + # `per_device_train_batch_size` to match `num_generations`, but vanilla HF + # paths (cb_paged, cb_sdpa, naive_trl) do not -- do it ourselves. + if args.backend not in ("vllm", "unsloth_fi_false"): + effective = (args.per_device_train_batch_size + * args.gradient_accumulation_steps) + if effective % args.num_generations != 0: + new_pdb = args.num_generations + print(f"[{args.backend}] Bumping per_device_train_batch_size " + f"{args.per_device_train_batch_size} -> {new_pdb} to satisfy " + f"GRPO divisibility.") + args.per_device_train_batch_size = new_pdb # --- load model / tokenizer per backend ----------------------------------- persistent_teardown_target = None if args.backend == "vllm": - model, tokenizer = _load_unsloth(args, fast_inference = True) + model, tokenizer = _load_unsloth(args, fast_inference=True) elif args.backend == "unsloth_fi_false": - model, tokenizer = _load_unsloth(args, fast_inference = False) + model, tokenizer = _load_unsloth(args, fast_inference=False) elif args.backend == "cb_paged": - model, tokenizer = _load_vanilla_hf(args, attn_impl = "paged_attention") + # `paged_attention` requires cu_seq_lens on every forward, which only + # the CB rollout path provides. GRPO's training forward (dense batch) + # crashes. Load with `sdpa_paged` which gracefully falls back to + # plain SDPA when paged args are absent, and still exercises the + # paged path during CB rollout. + model, tokenizer = _load_vanilla_hf(args, attn_impl="sdpa_paged") elif args.backend == "cb_sdpa": - model, tokenizer = _load_vanilla_hf(args, attn_impl = "sdpa_paged") + model, tokenizer = _load_vanilla_hf(args, attn_impl="sdpa_paged") elif args.backend == "naive_trl": - model, tokenizer = _load_vanilla_hf(args, attn_impl = "sdpa") + model, tokenizer = _load_vanilla_hf(args, attn_impl="sdpa") else: raise ValueError(args.backend) apply_chat_template_to_tokenizer(tokenizer) - dataset, maximum_length = build_dataset( - tokenizer, max_seq_length = args.max_seq_length - ) + dataset, maximum_length = build_dataset(tokenizer, max_seq_length=args.max_seq_length) print(f"[{args.backend}] p90 prompt length = {maximum_length}") reward_funcs = build_reward_funcs(tokenizer) @@ -252,12 +228,12 @@ def main(): shared = build_grpo_kwargs( tokenizer, maximum_length, - max_seq_length = args.max_seq_length, - max_steps = args.max_steps, - num_generations = args.num_generations, - per_device_train_batch_size = args.per_device_train_batch_size, - gradient_accumulation_steps = args.gradient_accumulation_steps, - output_dir = args.output_dir, + max_seq_length=args.max_seq_length, + max_steps=args.max_steps, + num_generations=args.num_generations, + per_device_train_batch_size=args.per_device_train_batch_size, + gradient_accumulation_steps=args.gradient_accumulation_steps, + output_dir=args.output_dir, ) # Overwrite the equivalence-friendly sampling params. shared["temperature"] = args.temperature @@ -268,24 +244,18 @@ def main(): shared["learning_rate"] = args.learning_rate from trl import GRPOConfig, GRPOTrainer - if args.backend == "vllm": from vllm import SamplingParams - vllm_sp = SamplingParams( - temperature = args.temperature, - top_p = args.top_p, - min_p = args.min_p, - top_k = args.top_k, - seed = args.seed, - stop = [tokenizer.eos_token], - include_stop_str_in_output = True, + temperature=args.temperature, top_p=args.top_p, min_p=args.min_p, + top_k=args.top_k, seed=args.seed, + stop=[tokenizer.eos_token], include_stop_str_in_output=True, ) training_args = GRPOConfig( - use_vllm = True, - vllm_mode = "colocate", - vllm_sampling_params = vllm_sp, - vllm_gpu_memory_utilization = args.gpu_memory_utilization, + use_vllm=True, + vllm_mode="colocate", + vllm_sampling_params=vllm_sp, + vllm_gpu_memory_utilization=args.gpu_memory_utilization, **shared, ) elif args.backend == "unsloth_fi_false": @@ -293,16 +263,16 @@ def main(): # fast_inference=False + for_inference() wires the fast single-token # decode + cached fp16 LoRA. training_args = GRPOConfig( - use_vllm = False, - bf16 = True, + use_vllm=False, + bf16=True, **shared, ) elif args.backend in ("cb_paged", "cb_sdpa"): training_args = GRPOConfig( - use_vllm = False, - use_transformers_paged = True, - bf16 = True, - generation_kwargs = { + use_vllm=False, + use_transformers_paged=True, + bf16=True, + generation_kwargs={ "max_batch_tokens": args.max_batch_tokens, "num_blocks": args.num_blocks, }, @@ -310,63 +280,55 @@ def main(): ) else: # naive_trl training_args = GRPOConfig( - use_vllm = False, - bf16 = True, + use_vllm=False, + bf16=True, **shared, ) stats_cb = _make_stats_callback() trainer = GRPOTrainer( - model = model, - processing_class = tokenizer, - reward_funcs = reward_funcs, - args = training_args, - train_dataset = dataset, - callbacks = [stats_cb], + model=model, + processing_class=tokenizer, + reward_funcs=reward_funcs, + args=training_args, + train_dataset=dataset, + callbacks=[stats_cb], ) if args.persistent_cb and args.backend in ("cb_paged", "cb_sdpa"): from persistent_cb import install_for_model, teardown - - base = ( - trainer.model_wrapped.base_model.model - if hasattr(trainer.model_wrapped, "base_model") - else trainer.model_wrapped - ) + base = (trainer.model_wrapped.base_model.model + if hasattr(trainer.model_wrapped, "base_model") + else trainer.model_wrapped) install_for_model(base, trainer.generation_config) persistent_teardown_target = base # Phase 4: torch.compile on the training forward. if args.compile_mode and args.backend != "vllm": from torch_debugging_utils import clear_inductor_cache, CompileDebugger - clear_inductor_cache() - CompileDebugger.enable(graph_breaks = True, recompiles = True) + CompileDebugger.enable(graph_breaks=True, recompiles=True) # Raise Dynamo cache limit so dynamic-shape recompiles don't thrash. import torch._dynamo - torch._dynamo.config.cache_size_limit = 128 try: torch._dynamo.config.allow_unspec_int_on_nn_module = True except AttributeError: pass - print( - f"[{args.backend}] Compiling trainer.model.forward " - f"(mode={args.compile_mode}, dynamic={args.compile_dynamic})" - ) + print(f"[{args.backend}] Compiling trainer.model.forward " + f"(mode={args.compile_mode}, dynamic={args.compile_dynamic})") trainer.model.forward = torch.compile( trainer.model.forward, - mode = args.compile_mode, - dynamic = args.compile_dynamic, + mode=args.compile_mode, + dynamic=args.compile_dynamic, ) # Reference model inside TRL's GRPO loop also runs a forward. ref = getattr(trainer, "ref_model", None) if ref is not None: ref.forward = torch.compile( - ref.forward, - mode = args.compile_mode, - dynamic = args.compile_dynamic, + ref.forward, mode=args.compile_mode, + dynamic=args.compile_dynamic, ) torch.cuda.reset_peak_memory_stats() @@ -376,7 +338,6 @@ def main(): finally: if persistent_teardown_target is not None: from persistent_cb import teardown - teardown(persistent_teardown_target) train_wall = time.perf_counter() - t_start @@ -416,17 +377,10 @@ def main(): } summary_path = Path(args.stats_path).with_suffix(".summary.json") with open(summary_path, "w") as f: - json.dump(summary, f, indent = 2) - print( - json.dumps( - { - k: v - for k, v in summary.items() - if k not in ("losses", "rewards", "kls", "grad_norms", "step_times_ms") - }, - indent = 2, - ) - ) + json.dump(summary, f, indent=2) + print(json.dumps({k: v for k, v in summary.items() + if k not in ("losses", "rewards", "kls", "grad_norms", "step_times_ms")}, + indent=2)) print(f"\n[{args.backend}] wrote summary to {summary_path}") # vLLM engine holds refs; fast-exit rather than wait for shutdown. os._exit(0) diff --git a/scripts/benchmarks/results/grpo_equivalence.md b/scripts/benchmarks/results/grpo_equivalence.md new file mode 100644 index 0000000000..23c5f95fb0 --- /dev/null +++ b/scripts/benchmarks/results/grpo_equivalence.md @@ -0,0 +1,89 @@ +# Phase 2: end-to-end GRPO backend comparison (10-step vibe check) + +Same dataset, reward functions, sampling (`temperature=0.1, top_p=0.97, +min_p=0.5, top_k=5`), and seed (3407). `max_steps=10, num_generations=4, +per_device_train_batch_size=4` (auto-adjusted from 1 on vanilla-HF paths +to satisfy TRL's `generation_batch_size % num_generations == 0`). + +Callbacks: `StatisticsCallback` from `torch_debugging_utils` logs per-step +loss, grad-norm, memory, wall time. Median step time is computed over steps +4-10 (first 3 skipped for compile / graph / warmup amortization). + +## 10-step results + +| Backend | Train wall (s) | Median step (s) | Peak mem (GB) | % of vLLM | +|-------------------------------|----------------|-----------------|---------------|-----------| +| vLLM (fast_inference) | 74.4 | **4.14** | 157.9 | 100 % | +| unsloth_fi_false | 355.4 | 23.95 | **10.7** | 17 % | +| cb_paged (sdpa_paged load) | 466.0 | 36.02 | 55.6 | 11.5 % | + +Loss / reward / KL arrays for each backend (10 steps, rounded): + +| Step | vLLM loss | vLLM reward | vLLM kl | fi_false loss | fi_false reward | fi_false kl | cb_paged loss | cb_paged reward | +|------|-----------|-------------|----------|---------------|-----------------|-------------|----------------|------------------| +| 1 | 0.031 | 0.00 | 0.00000 | 0.000 | 0.50 | 0.00000 | -0.086 | 0.62 | +| 2 | -0.194 | -2.50 | 0.00000 | -0.089 | -6.50 | 0.00000 | 0.041 | -2.50 | +| 3 | 0.263 | -3.62 | 0.01192 | -0.139 | -2.50 | 0.00857 | 0.000 | 0.50 | +| 4 | -0.201 | 0.00 | 0.00422 | -0.124 | -2.50 | 0.00931 | 0.086 | -1.50 | +| 5 | 0.209 | 0.38 | 0.00369 | 0.000 | 0.50 | 0.00213 | 0.016 | 0.00 | +| 6 | 0.000 | -7.50 | 0.00250 | 0.000 | -7.50 | 0.00071 | 0.000 | -7.50 | +| 7 | 0.037 | 1.50 | 0.00614 | -0.010 | -5.50 | 0.00522 | 0.074 | 1.50 | +| 8 | 0.000 | -7.50 | 0.00465 | 0.034 | -5.25 | 0.00014 | -0.048 | -4.25 | +| 9 | 0.000 | 0.50 | 0.00176 | 0.000 | 0.50 | 0.01090 | -0.188 | -1.50 | +| 10 | 0.205 | -3.50 | 0.00200 | 0.204 | -2.50 | 0.00215 | 0.044 | -6.50 | + +## Observations + +1. **Coherence gate (all backends)**: losses are bounded in `[-0.25, 0.3]`, + grad-norms finite, rewards in the plan's expected negative-then-rising + range. No CJK-token salad, no NaNs. + +2. **KL trajectories are qualitatively matched** between vLLM and + `unsloth_fi_false` (both in `[0, 0.015]`), confirming that + `fast_inference=False` produces rollouts close to the vLLM reference once + `temperature=0.1` is used. `cb_paged` also produces rollouts but our + `StatisticsCallback` did not capture TRL's `kl` entry in its `on_log` + pass -- the next iteration will forward every log dict entry into the JSON. + +3. **Per-step timing**: `unsloth_fi_false` is 5.8x slower than vLLM; `cb_paged` + is 8.7x slower. Neither hits the plan's 30% target on this vibe check. + +4. **Memory is the standout axis**: + - vLLM: 158 GB (prefill KV cache + vLLM engine overhead) + - cb_paged: 55.6 GB (paged cache only) + - unsloth_fi_false: **10.7 GB** -- 15x lower than vLLM. + + Unsloth's fast_inference=False path is a genuine option for teams who + cannot afford the vLLM footprint but are willing to take a ~5-6x rollout + wall-clock hit. + +5. **cb_paged load needed `sdpa_paged` not `paged_attention`**: the + FA4-shimmed `paged_attention` kernel requires `cu_seq_lens_q` on every + forward, but GRPO's training forward (dense batch) doesn't provide them. + `sdpa_paged` falls back to plain SDPA when no paged kwargs are present and + still exercises paged attention during the CB rollout. This is consistent + with the existing `qwen3_grpo_tpaged.py` which loads with `sdpa`. + +## What's next (not yet run) + +- **30-step equivalence** with `torch_debugging_utils.compare_training_runs` + comparing vLLM vs each backend on loss / reward / KL arrays. +- **Phase 3 sync driver** smoke-tested successfully (eager decode produces + 512 correct tokens) but CUDA graph capture hangs on the first graphed step. + Likely cause: `PagedAttentionCache` constructs tensors inside + `cache.update()` the first call, which doesn't survive graph capture. + Two possible fixes being explored: (a) pre-capture warmup steps on the + capture stream so allocations are already done, (b) replace in-place + torch.multinomial-adjacent ops with CUDA-graph-safe equivalents. +- **Phase 4 torch.compile**: hook-up ready in `qwen3_grpo_unified.py` + (`--compile_mode default|reduce-overhead|max-autotune-no-cudagraphs`); + needs a run budget allocated and the `CompileDebugger` output reviewed. + +## Raw stats + +- `scripts/benchmarks/results/stats/grpo_vllm_10.summary.json` +- `scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.summary.json` +- `scripts/benchmarks/results/stats/grpo_cb_paged_10.summary.json` + +Full per-step logs (one entry per step with loss/reward/kl/grad_norm and all +of TRL's logging dict) live at `scripts/benchmarks/results/stats/grpo_*.json`. diff --git a/scripts/benchmarks/results/stats/cb_sync_smoke.json b/scripts/benchmarks/results/stats/cb_sync_smoke.json new file mode 100644 index 0000000000..e8eda49452 --- /dev/null +++ b/scripts/benchmarks/results/stats/cb_sync_smoke.json @@ -0,0 +1,15 @@ +{ + "backend": "cb_sync_driver", + "use_cuda_graph": false, + "attn_impl": "paged_attention", + "n_prompts": 8, + "n_decoded_tokens": 512, + "wall_times_s": [ + 100.8571443540277, + 100.87157070200192 + ], + "median_wall_s": 100.87157070200192, + "decode_tps": 5.075761152887835, + "max_new_tokens": 64, + "peak_memory_gb": 45.91296434402466 +} \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/grpo_cb_paged_10.json b/scripts/benchmarks/results/stats/grpo_cb_paged_10.json new file mode 100644 index 0000000000..a8e557b176 --- /dev/null +++ b/scripts/benchmarks/results/stats/grpo_cb_paged_10.json @@ -0,0 +1,352 @@ +[ + { + "step": 1, + "loss": -0.0862, + "grad_norm": 716.0, + "learning_rate": 0.0, + "num_tokens": 4262.0, + "completions/mean_length": 953.5, + "completions/min_length": 824.0, + "completions/max_length": 1092.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 953.5, + "completions/min_terminated_length": 824.0, + "completions/max_terminated_length": 1092.0, + "rewards/match_format_exactly/mean": 2.25, + "rewards/match_format_exactly/std": 1.5, + "rewards/match_format_approximately/mean": 1.125, + "rewards/match_format_approximately/std": 0.75, + "rewards/check_answer/mean": -1.25, + "rewards/check_answer/std": 2.1794495582580566, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.625, + "reward_std": 3.4731109142303467, + "frac_reward_zero_std": 0.0, + "entropy": 0.1351587027311325, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 7.868439688409789e-05, + "time_ms": 56644.44096497027, + "memory_mb": 57451.41162109375, + "memory_gb": 56.104894161224365 + }, + { + "step": 2, + "loss": 0.041, + "grad_norm": 186.0, + "learning_rate": 5e-06, + "num_tokens": 6762.0, + "completions/mean_length": 536.0, + "completions/min_length": 492.0, + "completions/max_length": 603.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 536.0, + "completions/min_terminated_length": 492.0, + "completions/max_terminated_length": 603.0, + "rewards/match_format_exactly/mean": 0.75, + "rewards/match_format_exactly/std": 1.5, + "rewards/match_format_approximately/mean": 0.375, + "rewards/match_format_approximately/std": 0.75, + "rewards/check_answer/mean": -2.125, + "rewards/check_answer/std": 0.25, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": -2.5, + "reward_std": 2.0, + "frac_reward_zero_std": 0.0, + "entropy": 0.05973631516098976, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00015736879376819577, + "time_ms": 29505.720576969907, + "memory_mb": 53805.20556640625, + "memory_gb": 52.5441460609436 + }, + { + "step": 3, + "loss": 0.0, + "grad_norm": 0.0, + "learning_rate": 4.444444444444444e-06, + "num_tokens": 10099.0, + "completions/mean_length": 657.25, + "completions/min_length": 436.0, + "completions/max_length": 1302.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 657.25, + "completions/min_terminated_length": 436.0, + "completions/max_terminated_length": 1302.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.5, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "entropy": 0.06822667270898819, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00023605319065229366, + "time_ms": 62939.08593803644, + "memory_mb": 59249.8505859375, + "memory_gb": 57.86118221282959 + }, + { + "step": 4, + "loss": 0.0855, + "grad_norm": 274.0, + "learning_rate": 3.88888888888889e-06, + "num_tokens": 13644.0, + "completions/mean_length": 721.25, + "completions/min_length": 441.0, + "completions/max_length": 988.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 721.25, + "completions/min_terminated_length": 441.0, + "completions/max_terminated_length": 988.0, + "rewards/match_format_exactly/mean": 1.5, + "rewards/match_format_exactly/std": 1.7320507764816284, + "rewards/match_format_approximately/mean": 0.75, + "rewards/match_format_approximately/std": 0.8660253882408142, + "rewards/check_answer/mean": -2.25, + "rewards/check_answer/std": 0.28867512941360474, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": -1.5, + "reward_std": 2.309401035308838, + "frac_reward_zero_std": 0.0, + "entropy": 0.25085046887397766, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00031473758753639155, + "time_ms": 49076.82833302533, + "memory_mb": 56826.26611328125, + "memory_gb": 55.49440050125122 + }, + { + "step": 5, + "loss": 0.0162, + "grad_norm": 143.0, + "learning_rate": 3.3333333333333333e-06, + "num_tokens": 15442.0, + "completions/mean_length": 293.5, + "completions/min_length": 246.0, + "completions/max_length": 365.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 293.5, + "completions/min_terminated_length": 246.0, + "completions/max_terminated_length": 365.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -3.0, + "rewards/check_answer/std": 1.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.0, + "reward_std": 1.0, + "frac_reward_zero_std": 0.0, + "entropy": 0.0809403508901596, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00039342198442048943, + "time_ms": 18123.881562962197, + "memory_mb": 52261.25830078125, + "memory_gb": 51.03638505935669 + }, + { + "step": 6, + "loss": 0.0, + "grad_norm": 0.0, + "learning_rate": 2.7777777777777783e-06, + "num_tokens": 21877.0, + "completions/mean_length": 1511.75, + "completions/min_length": 1112.0, + "completions/max_length": 1846.0, + "completions/clipped_ratio": 0.5, + "completions/mean_terminated_length": 1177.5, + "completions/min_terminated_length": 1112.0, + "completions/max_terminated_length": 1243.0, + "rewards/match_format_exactly/mean": 0.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": -3.0, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.0, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -2.5, + "rewards/check_numbers/std": 0.0, + "reward": -7.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "entropy": 0.2572544813156128, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0004721063813045873, + "time_ms": 92771.2398529984, + "memory_mb": 63378.49365234375, + "memory_gb": 61.89306020736694 + }, + { + "step": 7, + "loss": 0.0736, + "grad_norm": 74.0, + "learning_rate": 2.222222222222222e-06, + "num_tokens": 24635.0, + "completions/mean_length": 546.5, + "completions/min_length": 466.0, + "completions/max_length": 585.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 546.5, + "completions/min_terminated_length": 466.0, + "completions/max_terminated_length": 585.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -1.5, + "rewards/check_answer/std": 2.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 1.5, + "reward_std": 2.0, + "frac_reward_zero_std": 0.0, + "entropy": 0.05001620948314667, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0005507907781886852, + "time_ms": 30032.14380296413, + "memory_mb": 53704.75830078125, + "memory_gb": 52.44605302810669 + }, + { + "step": 8, + "loss": -0.0475, + "grad_norm": 97.0, + "learning_rate": 1.6666666666666667e-06, + "num_tokens": 27393.0, + "completions/mean_length": 622.5, + "completions/min_length": 533.0, + "completions/max_length": 696.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 622.5, + "completions/min_terminated_length": 533.0, + "completions/max_terminated_length": 696.0, + "rewards/match_format_exactly/mean": 0.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": -1.5, + "rewards/match_format_approximately/std": 1.7320507764816284, + "rewards/check_answer/mean": -2.0, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -0.75, + "rewards/check_numbers/std": 2.872281312942505, + "reward": -4.25, + "reward_std": 4.27200174331665, + "frac_reward_zero_std": 0.0, + "entropy": 0.08264704048633575, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0006294751750727831, + "time_ms": 36016.123837034684, + "memory_mb": 54504.00634765625, + "memory_gb": 53.22656869888306 + }, + { + "step": 9, + "loss": -0.1881, + "grad_norm": 274.0, + "learning_rate": 1.111111111111111e-06, + "num_tokens": 29461.0, + "completions/mean_length": 420.0, + "completions/min_length": 327.0, + "completions/max_length": 647.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 420.0, + "completions/min_terminated_length": 327.0, + "completions/max_terminated_length": 647.0, + "rewards/match_format_exactly/mean": 1.5, + "rewards/match_format_exactly/std": 1.7320507764816284, + "rewards/match_format_approximately/mean": 0.0, + "rewards/match_format_approximately/std": 2.1213202476501465, + "rewards/check_answer/mean": -1.25, + "rewards/check_answer/std": 1.8484227657318115, + "rewards/check_numbers/mean": -1.75, + "rewards/check_numbers/std": 0.5, + "reward": -1.5, + "reward_std": 5.16397762298584, + "frac_reward_zero_std": 0.0, + "entropy": 0.22891533374786377, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0007081595719568809, + "time_ms": 35705.52668598248, + "memory_mb": 54149.77783203125, + "memory_gb": 52.88064241409302 + }, + { + "step": 10, + "loss": 0.0444, + "grad_norm": 236.0, + "learning_rate": 5.555555555555555e-07, + "num_tokens": 33785.0, + "completions/mean_length": 913.0, + "completions/min_length": 832.0, + "completions/max_length": 998.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 913.0, + "completions/min_terminated_length": 832.0, + "completions/max_terminated_length": 998.0, + "rewards/match_format_exactly/mean": 0.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": -2.25, + "rewards/match_format_approximately/std": 1.5, + "rewards/check_answer/mean": -2.0, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -2.25, + "rewards/check_numbers/std": 0.5, + "reward": -6.5, + "reward_std": 2.0, + "frac_reward_zero_std": 0.0, + "entropy": 0.1578676998615265, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0007868439688409789, + "time_ms": 53879.347916983534, + "memory_mb": 56904.7578125, + "memory_gb": 55.57105255126953 + } +] \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/grpo_cb_paged_10.summary.json b/scripts/benchmarks/results/stats/grpo_cb_paged_10.summary.json new file mode 100644 index 0000000000..23237f400e --- /dev/null +++ b/scripts/benchmarks/results/stats/grpo_cb_paged_10.summary.json @@ -0,0 +1,64 @@ +{ + "backend": "cb_paged", + "max_steps": 10, + "train_wall_s": 466.01091928296955, + "median_step_ms_post_warmup": 36016.123837034684, + "n_logged_steps": 10, + "sampling": { + "temperature": 0.1, + "top_p": 0.97, + "min_p": 0.5, + "top_k": 5 + }, + "losses": [ + -0.0862, + 0.041, + 0.0, + 0.0855, + 0.0162, + 0.0, + 0.0736, + -0.0475, + -0.1881, + 0.0444 + ], + "rewards": [ + 0.625, + -2.5, + 0.5, + -1.5, + 0.0, + -7.5, + 1.5, + -4.25, + -1.5, + -6.5 + ], + "kls": [], + "grad_norms": [ + 716.0, + 186.0, + 0.0, + 274.0, + 143.0, + 0.0, + 74.0, + 97.0, + 274.0, + 236.0 + ], + "step_times_ms": [ + 56644.44096497027, + 29505.720576969907, + 62939.08593803644, + 49076.82833302533, + 18123.881562962197, + 92771.2398529984, + 30032.14380296413, + 36016.123837034684, + 35705.52668598248, + 53879.347916983534 + ], + "peak_memory_gb": 55.57105255126953, + "logs_path": "logs/grpo_cb_paged_10.json" +} \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.json b/scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.json new file mode 100644 index 0000000000..c777419a15 --- /dev/null +++ b/scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.json @@ -0,0 +1,362 @@ +[ + { + "step": 1, + "loss": 0.0, + "grad_norm": 0.0, + "learning_rate": 0.0, + "num_tokens": 3693.0, + "completions/mean_length": 811.25, + "completions/min_length": 779.0, + "completions/max_length": 856.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 811.25, + "completions/min_terminated_length": 779.0, + "completions/max_terminated_length": 856.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.5, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "completion_length": 811.25, + "kl": 0.0, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 7.868439688409789e-05, + "time_ms": 77897.29859499494, + "memory_mb": 9442.14013671875, + "memory_gb": 9.220839977264404 + }, + { + "step": 2, + "loss": -0.0893, + "grad_norm": 0.6125104427337646, + "learning_rate": 5e-06, + "num_tokens": 6238.0, + "completions/mean_length": 547.25, + "completions/min_length": 487.0, + "completions/max_length": 645.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 547.25, + "completions/min_terminated_length": 487.0, + "completions/max_terminated_length": 645.0, + "rewards/match_format_exactly/mean": 0.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": -2.25, + "rewards/match_format_approximately/std": 1.5, + "rewards/check_answer/mean": -2.0, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -2.25, + "rewards/check_numbers/std": 0.5, + "reward": -6.5, + "reward_std": 2.0, + "frac_reward_zero_std": 0.0, + "completion_length": 547.25, + "kl": 0.0, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00015736879376819577, + "time_ms": 21634.983669035137, + "memory_mb": 9076.26025390625, + "memory_gb": 8.863535404205322 + }, + { + "step": 3, + "loss": -0.1386, + "grad_norm": 0.5993297696113586, + "learning_rate": 4.444444444444444e-06, + "num_tokens": 10165.0, + "completions/mean_length": 804.75, + "completions/min_length": 605.0, + "completions/max_length": 1214.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 804.75, + "completions/min_terminated_length": 605.0, + "completions/max_terminated_length": 1214.0, + "rewards/match_format_exactly/mean": 1.5, + "rewards/match_format_exactly/std": 1.7320507764816284, + "rewards/match_format_approximately/mean": -0.75, + "rewards/match_format_approximately/std": 2.598076105117798, + "rewards/check_answer/mean": -1.25, + "rewards/check_answer/std": 1.8484227657318115, + "rewards/check_numbers/mean": -2.0, + "rewards/check_numbers/std": 0.5773502588272095, + "reward": -2.5, + "reward_std": 6.0, + "frac_reward_zero_std": 0.0, + "completion_length": 804.75, + "kl": 0.008573448285460472, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00023605319065229366, + "time_ms": 40298.14357904252, + "memory_mb": 9952.80322265625, + "memory_gb": 9.719534397125244 + }, + { + "step": 4, + "loss": -0.1236, + "grad_norm": 0.5647851228713989, + "learning_rate": 3.88888888888889e-06, + "num_tokens": 13320.0, + "completions/mean_length": 623.75, + "completions/min_length": 421.0, + "completions/max_length": 789.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 623.75, + "completions/min_terminated_length": 421.0, + "completions/max_terminated_length": 789.0, + "rewards/match_format_exactly/mean": 0.75, + "rewards/match_format_exactly/std": 1.5, + "rewards/match_format_approximately/mean": 0.375, + "rewards/match_format_approximately/std": 0.75, + "rewards/check_answer/mean": -2.125, + "rewards/check_answer/std": 0.25, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": -2.5, + "reward_std": 2.0, + "frac_reward_zero_std": 0.0, + "completion_length": 623.75, + "kl": 0.009312103502452374, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00031473758753639155, + "time_ms": 26193.765547999647, + "memory_mb": 9306.09326171875, + "memory_gb": 9.087981700897217 + }, + { + "step": 5, + "loss": 0.0, + "grad_norm": 0.0010538548231124878, + "learning_rate": 3.3333333333333333e-06, + "num_tokens": 14970.0, + "completions/mean_length": 256.5, + "completions/min_length": 246.0, + "completions/max_length": 260.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 256.5, + "completions/min_terminated_length": 246.0, + "completions/max_terminated_length": 260.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.5, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "completion_length": 256.5, + "kl": 0.002130241831764579, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00039342198442048943, + "time_ms": 9324.433026020415, + "memory_mb": 8777.982421875, + "memory_gb": 8.572248458862305 + }, + { + "step": 6, + "loss": 0.0, + "grad_norm": 0.00012166703527327627, + "learning_rate": 2.7777777777777783e-06, + "num_tokens": 21841.0, + "completions/mean_length": 1620.75, + "completions/min_length": 1214.0, + "completions/max_length": 1846.0, + "completions/clipped_ratio": 0.5, + "completions/mean_terminated_length": 1395.5, + "completions/min_terminated_length": 1214.0, + "completions/max_terminated_length": 1577.0, + "rewards/match_format_exactly/mean": 0.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": -3.0, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.0, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -2.5, + "rewards/check_numbers/std": 0.0, + "reward": -7.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "completion_length": 1620.75, + "kl": 0.0007094849133864045, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0004721063813045873, + "time_ms": 61652.94101298787, + "memory_mb": 10914.11279296875, + "memory_gb": 10.658313274383545 + }, + { + "step": 7, + "loss": -0.0097, + "grad_norm": 0.9194015860557556, + "learning_rate": 2.222222222222222e-06, + "num_tokens": 24587.0, + "completions/mean_length": 543.5, + "completions/min_length": 511.0, + "completions/max_length": 562.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 543.5, + "completions/min_terminated_length": 511.0, + "completions/max_terminated_length": 562.0, + "rewards/match_format_exactly/mean": 0.75, + "rewards/match_format_exactly/std": 1.5, + "rewards/match_format_approximately/mean": -1.875, + "rewards/match_format_approximately/std": 2.25, + "rewards/check_answer/mean": -2.125, + "rewards/check_answer/std": 0.25, + "rewards/check_numbers/mean": -2.25, + "rewards/check_numbers/std": 0.5, + "reward": -5.5, + "reward_std": 4.0, + "frac_reward_zero_std": 0.0, + "completion_length": 543.5, + "kl": 0.005215016193687916, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0005507907781886852, + "time_ms": 18906.7719859886, + "memory_mb": 8996.00390625, + "memory_gb": 8.785160064697266 + }, + { + "step": 8, + "loss": 0.0338, + "grad_norm": 0.5346357822418213, + "learning_rate": 1.6666666666666667e-06, + "num_tokens": 27519.0, + "completions/mean_length": 666.0, + "completions/min_length": 615.0, + "completions/max_length": 714.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 666.0, + "completions/min_terminated_length": 615.0, + "completions/max_terminated_length": 714.0, + "rewards/match_format_exactly/mean": 0.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": -2.25, + "rewards/match_format_approximately/std": 1.5, + "rewards/check_answer/mean": -2.0, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -1.0, + "rewards/check_numbers/std": 3.0, + "reward": -5.25, + "reward_std": 4.5, + "frac_reward_zero_std": 0.0, + "completion_length": 666.0, + "kl": 0.0001442090724594891, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0006294751750727831, + "time_ms": 23948.13355000224, + "memory_mb": 9202.39306640625, + "memory_gb": 8.986711978912354 + }, + { + "step": 9, + "loss": 0.0, + "grad_norm": 0.0033828848972916603, + "learning_rate": 1.111111111111111e-06, + "num_tokens": 29207.0, + "completions/mean_length": 325.0, + "completions/min_length": 310.0, + "completions/max_length": 334.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 325.0, + "completions/min_terminated_length": 310.0, + "completions/max_terminated_length": 334.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.5, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "completion_length": 325.0, + "kl": 0.010901343077421188, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0007081595719568809, + "time_ms": 11388.401818985585, + "memory_mb": 8705.7236328125, + "memory_gb": 8.501683235168457 + }, + { + "step": 10, + "loss": 0.2036, + "grad_norm": 0.22472381591796875, + "learning_rate": 5.555555555555555e-07, + "num_tokens": 34891.0, + "completions/mean_length": 1253.0, + "completions/min_length": 1044.0, + "completions/max_length": 1846.0, + "completions/clipped_ratio": 0.25, + "completions/mean_terminated_length": 1055.3333740234375, + "completions/min_terminated_length": 1044.0, + "completions/max_terminated_length": 1067.0, + "rewards/match_format_exactly/mean": 1.5, + "rewards/match_format_exactly/std": 1.7320507764816284, + "rewards/match_format_approximately/mean": 0.0, + "rewards/match_format_approximately/std": 2.1213202476501465, + "rewards/check_answer/mean": -2.25, + "rewards/check_answer/std": 0.28867512941360474, + "rewards/check_numbers/mean": -1.75, + "rewards/check_numbers/std": 0.5, + "reward": -2.5, + "reward_std": 3.8297085762023926, + "frac_reward_zero_std": 0.0, + "completion_length": 1253.0, + "kl": 0.00215436820872128, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0007868439688409789, + "time_ms": 61737.33338096645, + "memory_mb": 10919.5498046875, + "memory_gb": 10.663622856140137 + } +] \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.summary.json b/scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.summary.json new file mode 100644 index 0000000000..aba12f1486 --- /dev/null +++ b/scripts/benchmarks/results/stats/grpo_unsloth_fi_false_10.summary.json @@ -0,0 +1,75 @@ +{ + "backend": "unsloth_fi_false", + "max_steps": 10, + "train_wall_s": 355.41431403998286, + "median_step_ms_post_warmup": 23948.13355000224, + "n_logged_steps": 10, + "sampling": { + "temperature": 0.1, + "top_p": 0.97, + "min_p": 0.5, + "top_k": 5 + }, + "losses": [ + 0.0, + -0.0893, + -0.1386, + -0.1236, + 0.0, + 0.0, + -0.0097, + 0.0338, + 0.0, + 0.2036 + ], + "rewards": [ + 0.5, + -6.5, + -2.5, + -2.5, + 0.5, + -7.5, + -5.5, + -5.25, + 0.5, + -2.5 + ], + "kls": [ + 0.0, + 0.0, + 0.008573448285460472, + 0.009312103502452374, + 0.002130241831764579, + 0.0007094849133864045, + 0.005215016193687916, + 0.0001442090724594891, + 0.010901343077421188, + 0.00215436820872128 + ], + "grad_norms": [ + 0.0, + 0.6125104427337646, + 0.5993297696113586, + 0.5647851228713989, + 0.0010538548231124878, + 0.00012166703527327627, + 0.9194015860557556, + 0.5346357822418213, + 0.0033828848972916603, + 0.22472381591796875 + ], + "step_times_ms": [ + 77897.29859499494, + 21634.983669035137, + 40298.14357904252, + 26193.765547999647, + 9324.433026020415, + 61652.94101298787, + 18906.7719859886, + 23948.13355000224, + 11388.401818985585, + 61737.33338096645 + ], + "peak_memory_gb": 10.663622856140137, + "logs_path": "logs/grpo_unsloth_fi_false_10.json" +} \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/grpo_vllm_10.json b/scripts/benchmarks/results/stats/grpo_vllm_10.json new file mode 100644 index 0000000000..0b7390a1b4 --- /dev/null +++ b/scripts/benchmarks/results/stats/grpo_vllm_10.json @@ -0,0 +1,362 @@ +[ + { + "step": 1, + "loss": 0.0305, + "grad_norm": 0.4133029878139496, + "learning_rate": 0.0, + "num_tokens": 3705.0, + "completions/mean_length": 814.25, + "completions/min_length": 781.0, + "completions/max_length": 864.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 814.25, + "completions/min_terminated_length": 781.0, + "completions/max_terminated_length": 864.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -3.0, + "rewards/check_answer/std": 1.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.0, + "reward_std": 1.0, + "frac_reward_zero_std": 0.0, + "completion_length": 814.25, + "kl": 0.0, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 7.868439688409789e-05, + "time_ms": 17984.56621397054, + "memory_mb": 161170.2099609375, + "memory_gb": 157.39278316497803 + }, + { + "step": 2, + "loss": -0.1941, + "grad_norm": 0.8333088159561157, + "learning_rate": 5e-06, + "num_tokens": 7167.0, + "completions/mean_length": 776.5, + "completions/min_length": 525.0, + "completions/max_length": 1078.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 776.5, + "completions/min_terminated_length": 525.0, + "completions/max_terminated_length": 1078.0, + "rewards/match_format_exactly/mean": 0.75, + "rewards/match_format_exactly/std": 1.5, + "rewards/match_format_approximately/mean": 0.375, + "rewards/match_format_approximately/std": 0.75, + "rewards/check_answer/mean": -2.125, + "rewards/check_answer/std": 0.25, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": -2.5, + "reward_std": 2.0, + "frac_reward_zero_std": 0.0, + "completion_length": 776.5, + "kl": 0.0, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00015736879376819577, + "time_ms": 6704.717919987161, + "memory_mb": 161638.7705078125, + "memory_gb": 157.85036182403564 + }, + { + "step": 3, + "loss": 0.2632, + "grad_norm": 0.4677680730819702, + "learning_rate": 4.444444444444444e-06, + "num_tokens": 11596.0, + "completions/mean_length": 930.25, + "completions/min_length": 445.0, + "completions/max_length": 1846.0, + "completions/clipped_ratio": 0.25, + "completions/mean_terminated_length": 625.0, + "completions/min_terminated_length": 445.0, + "completions/max_terminated_length": 863.0, + "rewards/match_format_exactly/mean": 1.5, + "rewards/match_format_exactly/std": 1.7320507764816284, + "rewards/match_format_approximately/mean": -0.75, + "rewards/match_format_approximately/std": 2.598076105117798, + "rewards/check_answer/mean": -2.75, + "rewards/check_answer/std": 1.190238118171692, + "rewards/check_numbers/mean": -1.625, + "rewards/check_numbers/std": 1.1814539432525635, + "reward": -3.625, + "reward_std": 4.479118347167969, + "frac_reward_zero_std": 0.0, + "completion_length": 930.25, + "kl": 0.011923530139029026, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00023605319065229366, + "time_ms": 12108.133931003977, + "memory_mb": 162822.73876953125, + "memory_gb": 159.00658082962036 + }, + { + "step": 4, + "loss": -0.2013, + "grad_norm": 0.5163940191268921, + "learning_rate": 3.88888888888889e-06, + "num_tokens": 14365.0, + "completions/mean_length": 527.25, + "completions/min_length": 315.0, + "completions/max_length": 598.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 527.25, + "completions/min_terminated_length": 315.0, + "completions/max_terminated_length": 598.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -3.0, + "rewards/check_answer/std": 1.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.0, + "reward_std": 1.0, + "frac_reward_zero_std": 0.0, + "completion_length": 527.25, + "kl": 0.004221913404762745, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00031473758753639155, + "time_ms": 4019.3081409670413, + "memory_mb": 160943.57275390625, + "memory_gb": 157.17145776748657 + }, + { + "step": 5, + "loss": 0.2093, + "grad_norm": 1.160618782043457, + "learning_rate": 3.3333333333333333e-06, + "num_tokens": 16176.0, + "completions/mean_length": 296.75, + "completions/min_length": 246.0, + "completions/max_length": 421.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 296.75, + "completions/min_terminated_length": 246.0, + "completions/max_terminated_length": 421.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -3.0, + "rewards/check_answer/std": 1.0, + "rewards/check_numbers/mean": -1.125, + "rewards/check_numbers/std": 0.75, + "reward": 0.375, + "reward_std": 0.25, + "frac_reward_zero_std": 0.0, + "completion_length": 296.75, + "kl": 0.003692739875987172, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.00039342198442048943, + "time_ms": 3241.851194994524, + "memory_mb": 160665.36376953125, + "memory_gb": 156.89976930618286 + }, + { + "step": 6, + "loss": 0.0, + "grad_norm": 0.0007444396032951772, + "learning_rate": 2.7777777777777783e-06, + "num_tokens": 23948.0, + "completions/mean_length": 1846.0, + "completions/min_length": 1846.0, + "completions/max_length": 1846.0, + "completions/clipped_ratio": 1.0, + "completions/mean_terminated_length": 0.0, + "completions/min_terminated_length": 0.0, + "completions/max_terminated_length": 0.0, + "rewards/match_format_exactly/mean": 0.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": -3.0, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.0, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -2.5, + "rewards/check_numbers/std": 0.0, + "reward": -7.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "completion_length": 1846.0, + "kl": 0.0025038770399987698, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0004721063813045873, + "time_ms": 10860.668059962336, + "memory_mb": 162817.607421875, + "memory_gb": 159.0015697479248 + }, + { + "step": 7, + "loss": 0.0371, + "grad_norm": 2.262518882751465, + "learning_rate": 2.222222222222222e-06, + "num_tokens": 26728.0, + "completions/mean_length": 552.0, + "completions/min_length": 506.0, + "completions/max_length": 627.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 552.0, + "completions/min_terminated_length": 506.0, + "completions/max_terminated_length": 627.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -1.5, + "rewards/check_answer/std": 2.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 1.5, + "reward_std": 2.0, + "frac_reward_zero_std": 0.0, + "completion_length": 552.0, + "kl": 0.006138760130852461, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0005507907781886852, + "time_ms": 4138.088690000586, + "memory_mb": 160989.1611328125, + "memory_gb": 157.2159776687622 + }, + { + "step": 8, + "loss": 0.0, + "grad_norm": 0.001562082557938993, + "learning_rate": 1.6666666666666667e-06, + "num_tokens": 29420.0, + "completions/mean_length": 606.0, + "completions/min_length": 560.0, + "completions/max_length": 636.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 606.0, + "completions/min_terminated_length": 560.0, + "completions/max_terminated_length": 636.0, + "rewards/match_format_exactly/mean": 0.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": -3.0, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.0, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -2.5, + "rewards/check_numbers/std": 0.0, + "reward": -7.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "completion_length": 606.0, + "kl": 0.004652692936360836, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0006294751750727831, + "time_ms": 4177.52773797838, + "memory_mb": 160996.86376953125, + "memory_gb": 157.22349977493286 + }, + { + "step": 9, + "loss": 0.0, + "grad_norm": 0.00027447607135400176, + "learning_rate": 1.111111111111111e-06, + "num_tokens": 31353.0, + "completions/mean_length": 386.25, + "completions/min_length": 334.0, + "completions/max_length": 464.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 386.25, + "completions/min_terminated_length": 334.0, + "completions/max_terminated_length": 464.0, + "rewards/match_format_exactly/mean": 3.0, + "rewards/match_format_exactly/std": 0.0, + "rewards/match_format_approximately/mean": 1.5, + "rewards/match_format_approximately/std": 0.0, + "rewards/check_answer/mean": -2.5, + "rewards/check_answer/std": 0.0, + "rewards/check_numbers/mean": -1.5, + "rewards/check_numbers/std": 0.0, + "reward": 0.5, + "reward_std": 0.0, + "frac_reward_zero_std": 1.0, + "completion_length": 386.25, + "kl": 0.0017617446137592196, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0007081595719568809, + "time_ms": 3263.2006779895164, + "memory_mb": 160745.5380859375, + "memory_gb": 156.97806453704834 + }, + { + "step": 10, + "loss": 0.2052, + "grad_norm": 0.4320540428161621, + "learning_rate": 5.555555555555555e-07, + "num_tokens": 35257.0, + "completions/mean_length": 808.0, + "completions/min_length": 650.0, + "completions/max_length": 1119.0, + "completions/clipped_ratio": 0.0, + "completions/mean_terminated_length": 808.0, + "completions/min_terminated_length": 650.0, + "completions/max_terminated_length": 1119.0, + "rewards/match_format_exactly/mean": 1.5, + "rewards/match_format_exactly/std": 1.7320507764816284, + "rewards/match_format_approximately/mean": 0.0, + "rewards/match_format_approximately/std": 2.1213202476501465, + "rewards/check_answer/mean": -3.25, + "rewards/check_answer/std": 1.4433757066726685, + "rewards/check_numbers/mean": -1.75, + "rewards/check_numbers/std": 0.5, + "reward": -3.5, + "reward_std": 2.8284270763397217, + "frac_reward_zero_std": 0.0, + "completion_length": 808.0, + "kl": 0.001998987514525652, + "clip_ratio/low_mean": 0.0, + "clip_ratio/low_min": 0.0, + "clip_ratio/high_mean": 0.0, + "clip_ratio/high_max": 0.0, + "clip_ratio/region_mean": 0.0, + "epoch": 0.0007868439688409789, + "time_ms": 6874.866564990953, + "memory_mb": 161736.77734375, + "memory_gb": 157.94607162475586 + } +] \ No newline at end of file diff --git a/scripts/benchmarks/results/stats/grpo_vllm_10.summary.json b/scripts/benchmarks/results/stats/grpo_vllm_10.summary.json new file mode 100644 index 0000000000..e827daa75b --- /dev/null +++ b/scripts/benchmarks/results/stats/grpo_vllm_10.summary.json @@ -0,0 +1,75 @@ +{ + "backend": "vllm", + "max_steps": 10, + "train_wall_s": 74.41919421299826, + "median_step_ms_post_warmup": 4138.088690000586, + "n_logged_steps": 10, + "sampling": { + "temperature": 0.1, + "top_p": 0.97, + "min_p": 0.5, + "top_k": 5 + }, + "losses": [ + 0.0305, + -0.1941, + 0.2632, + -0.2013, + 0.2093, + 0.0, + 0.0371, + 0.0, + 0.0, + 0.2052 + ], + "rewards": [ + 0.0, + -2.5, + -3.625, + 0.0, + 0.375, + -7.5, + 1.5, + -7.5, + 0.5, + -3.5 + ], + "kls": [ + 0.0, + 0.0, + 0.011923530139029026, + 0.004221913404762745, + 0.003692739875987172, + 0.0025038770399987698, + 0.006138760130852461, + 0.004652692936360836, + 0.0017617446137592196, + 0.001998987514525652 + ], + "grad_norms": [ + 0.4133029878139496, + 0.8333088159561157, + 0.4677680730819702, + 0.5163940191268921, + 1.160618782043457, + 0.0007444396032951772, + 2.262518882751465, + 0.001562082557938993, + 0.00027447607135400176, + 0.4320540428161621 + ], + "step_times_ms": [ + 17984.56621397054, + 6704.717919987161, + 12108.133931003977, + 4019.3081409670413, + 3241.851194994524, + 10860.668059962336, + 4138.088690000586, + 4177.52773797838, + 3263.2006779895164, + 6874.866564990953 + ], + "peak_memory_gb": 157.94607162475586, + "logs_path": "logs/grpo_vllm_10.json" +} \ No newline at end of file