From 29c477ffa3236adcdb54784577552ecdc944b230 Mon Sep 17 00:00:00 2001 From: yuxz Date: Fri, 21 Nov 2025 18:00:48 +0800 Subject: [PATCH] fix: add LoRA support and batch generation to gsm8k eval script --- eval/gsm8k.py | 132 +++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 110 insertions(+), 22 deletions(-) diff --git a/eval/gsm8k.py b/eval/gsm8k.py index 1bd4d32..2244ece 100644 --- a/eval/gsm8k.py +++ b/eval/gsm8k.py @@ -10,6 +10,7 @@ ) from vllm import LLM, SamplingParams +from vllm.lora.request import LoRARequest from datasets import load_dataset, Dataset from peft import PeftModel, PeftConfig @@ -186,7 +187,7 @@ def load_nvfp4(model_name_or_path): return llm, tokenizer -def generate(llm, tokenizer, input_text, generate_kwargs): +def generate(llm, tokenizer, input_text, generate_kwargs, args): text = tokenizer.apply_chat_template( input_text, tokenize=False, @@ -195,9 +196,84 @@ def generate(llm, tokenizer, input_text, generate_kwargs): # max_tokens is for the maximum length for generation. sampling_params = SamplingParams(temperature=0.6, top_p=0.95, max_tokens=1024) - outputs = llm.generate(text, sampling_params,use_tqdm=False) + + # Check if there are LoRA adapters loaded in vLLM + if args.load == "nvfp4": + lora_requests = None + try: + # Check if LoRA is enabled by trying to access lora_manager + model_runner = llm.llm_engine.model_executor.driver_worker.model_runner + if hasattr(model_runner, 'lora_manager') and model_runner.lora_manager is not None: + lora_int_ids = list(llm.llm_engine.list_loras()) + if len(lora_int_ids) > 0: + # Use the first LoRA adapter + lora_int_id = lora_int_ids[0] + lora_requests = LoRARequest( + lora_name=f"{lora_int_id}", + lora_int_id=lora_int_id, + lora_path="lora_stub_path" + ) + except (AttributeError, Exception): + # LoRA is not enabled, continue without LoRA + pass + + # Generate with or without LoRA + outputs = llm.generate(text, sampling_params, lora_request=lora_requests, use_tqdm=False) + else: + outputs = llm.generate(text, sampling_params, use_tqdm=False) return outputs[0].outputs[0].text +def generate_batch(llm, tokenizer, input_texts, generate_kwargs, args): + """Batch generate for multiple prompts""" + texts = [] + for input_text in input_texts: + text = tokenizer.apply_chat_template( + input_text, + tokenize=False, + add_generation_prompt=True + ) + texts.append(text) + + # max_tokens is for the maximum length for generation. + if "checkpoint" in args.model or "ckpt" in args.model: + max_tokens = generate_kwargs.get("max_new_tokens", 4096) + else: + max_tokens = generate_kwargs.get("max_new_tokens", 2048) + + sampling_params = SamplingParams( + temperature=generate_kwargs.get("temperature", 0.8), + top_p=generate_kwargs.get("top_p", 0.95), + max_tokens=max_tokens + ) + + # Check if there are LoRA adapters loaded in vLLM + if args.load == "nvfp4": + lora_requests = None + try: + # Check if LoRA is enabled by trying to access lora_manager + model_runner = llm.llm_engine.model_executor.driver_worker.model_runner + if hasattr(model_runner, 'lora_manager') and model_runner.lora_manager is not None: + lora_int_ids = list(llm.llm_engine.list_loras()) + if len(lora_int_ids) > 0: + # Use the first LoRA adapter + lora_int_id = lora_int_ids[0] + lora_requests = LoRARequest( + lora_name=f"{lora_int_id}", + lora_int_id=lora_int_id, + lora_path="lora_stub_path" + ) + except (AttributeError, Exception): + # LoRA is not enabled, continue without LoRA + pass + + # Generate with or without LoRA + outputs = llm.generate(texts, sampling_params, lora_request=lora_requests, use_tqdm=True) + else: + outputs = llm.generate(texts, sampling_params, use_tqdm=True) + + # Extract texts from outputs + return [output.outputs[0].text for output in outputs] + def extract_xml_answer(text: str) -> str: answer = text.split("")[-1] answer = answer.split("")[0] @@ -207,7 +283,7 @@ def correctness_reward_func(prompts, response, answer, **kwargs) -> list[float]: q = prompts extracted_responses = extract_xml_answer(response) print('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer}", f"\nResponse:\n{response}", f"\nExtracted:\n{extracted_responses}") - return [1.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)] + return [1.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)] or [0.0] def accuracy_reward(prompts, completions, answer): """Reward function that checks if the completion is the same as the ground truth.""" @@ -303,30 +379,42 @@ def main(): else: llm, tokenizer = load(args.model) - answers = [] - for sample in tqdm(data): - input_text = sample["prompt"] - prompt = input_text[1]['content'] - if "checkpoint" in args.model or "ckpt" in args.model: - generate_kwargs = dict(max_new_tokens=4096, top_p=0.95, temperature=0.8) + # Prepare all prompts and answers for batch processing + input_texts = [] + prompts = [] + gold_answers = [] + + for sample in data: + input_texts.append(sample["prompt"]) + if len(sample["prompt"]) > 1: + prompts.append(sample["prompt"][1]['content']) else: - generate_kwargs = dict(max_new_tokens=2048, top_p=0.95, temperature=0.8) - response = generate(llm, tokenizer, input_text, generate_kwargs) - - gold_answer = sample["answer"] - + prompts.append(sample["prompt"][0]['content']) + gold_answers.append(sample["answer"]) + + # Prepare generate_kwargs + if "checkpoint" in args.model or "ckpt" in args.model: + generate_kwargs = dict(max_new_tokens=4096, top_p=0.95, temperature=0.8) + else: + generate_kwargs = dict(max_new_tokens=2048, top_p=0.95, temperature=0.8) + + # Batch generate all responses + print(f"Generating responses for {len(input_texts)} samples in batch...") + responses = generate_batch(llm, tokenizer, input_texts, generate_kwargs, args) + + # Process all results + answers = [] + for prompt, response, gold_answer in tqdm(zip(prompts, responses, gold_answers), total=len(prompts), desc="Evaluating"): if args.data == "gsm8k": answers.append(correctness_reward_func(prompt, response, gold_answer)[0]) else: answers.append(math_verified(prompt, response, gold_answer)[0]) - - - - print( - f"Num of total question: {len(answers)}, " - f"Correct num: {sum(answers)}, " - f"Accuracy: {float(sum(answers))/len(answers)}." - ) + + print( + f"Num of total question: {len(answers)}, " + f"Correct num: {sum(answers)}, " + f"Accuracy: {float(sum(answers))/len(answers)}." + ) os.makedirs(args.output_dir, exist_ok=True)