From 4b46889498f16b44d7fef707412182dc6f8e3061 Mon Sep 17 00:00:00 2001 From: JoshuaL3000 Date: Tue, 14 Jul 2026 21:30:24 -0700 Subject: [PATCH 1/4] test: Add e2e test for fast_inference vLLM path --- tests/fast_inference/test_fast_inference.py | 139 ++++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 tests/fast_inference/test_fast_inference.py diff --git a/tests/fast_inference/test_fast_inference.py b/tests/fast_inference/test_fast_inference.py new file mode 100644 index 0000000000..bbd96d0d7e --- /dev/null +++ b/tests/fast_inference/test_fast_inference.py @@ -0,0 +1,139 @@ +"""GRPO training test for the `fast_inference=True` vLLM rollout path. + +Usage: + python tests/fast_inference/test_fast_inference.py +""" + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from unsloth import FastLanguageModel +import torch +from datasets import Dataset +from trl import GRPOConfig, GRPOTrainer + +from tests.utils import header_footer_context + + +MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" +MAX_SEQ_LENGTH = 512 +MAX_COMPLETION_LENGTH = 256 +LOAD_IN_4BIT = False +GPU_MEMORY_UTILIZATION = 0.5 +LORA_RANK = 16 +NUM_GENERATIONS = 2 +MAX_STEPS = 3 +MAX_NEW_TOKENS = 32 +TEMPERATURE = 0.8 + +SYSTEM_PROMPT = "Respond concisely." +QUESTIONS = [ + "What is the capital of France?", + "Name a primary color.", + "What is 2 + 2?", + "Write the word 'hello'.", +] +PROMPTS = [ + [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": q}, + ] + for q in QUESTIONS +] + + +def length_reward_func(completions, **kwargs) -> list[float]: + """Simple reward: prefer non-empty completions.""" + responses = [completion[0]["content"] for completion in completions] + return [1.0 if len(r) > 0 else 0.0 for r in responses] + + +def get_unsloth_peft_model( + model, + lora_rank: int, + target_modules: list[str] = "all-linear", + use_gradient_checkpointing: str = False, + random_state: int = 42, +): + return FastLanguageModel.get_peft_model( + model, + r = lora_rank, + target_modules = target_modules, + lora_alpha = lora_rank, + use_gradient_checkpointing = use_gradient_checkpointing, + random_state = random_state, + ) + + +if __name__ == "__main__": + target_modules = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + + with header_footer_context("Load model (fast_inference=True)"): + model, tokenizer = FastLanguageModel.from_pretrained( + model_name=MODEL_NAME, + max_seq_length=MAX_SEQ_LENGTH, + dtype=torch.bfloat16, + load_in_4bit=LOAD_IN_4BIT, + fast_inference=True, + gpu_memory_utilization=GPU_MEMORY_UTILIZATION, + ) + + assert hasattr(model, "vllm_engine"), "model.vllm_engine not set — fast_inference=True failed" + print("vllm_engine attached") + + model = get_unsloth_peft_model( + model, + lora_rank=LORA_RANK, + target_modules=target_modules, + use_gradient_checkpointing=False, + random_state=42, + ) + + prompt = tokenizer.apply_chat_template( + PROMPTS[0], tokenize=False, add_generation_prompt=True + ) + with header_footer_context("Test Prompt"): + print(prompt) + + dataset = Dataset.from_dict({"prompt": PROMPTS}) + + with header_footer_context("GRPO config and trainer"): + training_args = GRPOConfig( + learning_rate=5e-6, + optim="adamw_torch", + temperature=TEMPERATURE, + top_k=50, + logging_steps=1, + per_device_train_batch_size=1, + gradient_accumulation_steps=NUM_GENERATIONS, + num_generations=NUM_GENERATIONS, + max_completion_length=MAX_COMPLETION_LENGTH, + max_steps=MAX_STEPS, + report_to="none", + ) + print(training_args) + + trainer = GRPOTrainer( + model=model, + processing_class=tokenizer, + reward_funcs=[length_reward_func], + args=training_args, + train_dataset=dataset, + ) + + with header_footer_context("GRPO train (vLLM rollout)"): + trainer_stats = trainer.train() + print(trainer_stats) + + assert trainer_stats is not None, "trainer.train() is None" From ba72958a153667ffd63a7359e943e0a00b6dc561 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 07:43:24 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/fast_inference/test_fast_inference.py | 60 ++++++++++----------- 1 file changed, 29 insertions(+), 31 deletions(-) diff --git a/tests/fast_inference/test_fast_inference.py b/tests/fast_inference/test_fast_inference.py index bbd96d0d7e..d60e41eb10 100644 --- a/tests/fast_inference/test_fast_inference.py +++ b/tests/fast_inference/test_fast_inference.py @@ -78,31 +78,29 @@ def get_unsloth_peft_model( "up_proj", "down_proj", ] - + with header_footer_context("Load model (fast_inference=True)"): model, tokenizer = FastLanguageModel.from_pretrained( - model_name=MODEL_NAME, - max_seq_length=MAX_SEQ_LENGTH, - dtype=torch.bfloat16, - load_in_4bit=LOAD_IN_4BIT, - fast_inference=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, + model_name = MODEL_NAME, + max_seq_length = MAX_SEQ_LENGTH, + dtype = torch.bfloat16, + load_in_4bit = LOAD_IN_4BIT, + fast_inference = True, + gpu_memory_utilization = GPU_MEMORY_UTILIZATION, ) assert hasattr(model, "vllm_engine"), "model.vllm_engine not set — fast_inference=True failed" print("vllm_engine attached") - + model = get_unsloth_peft_model( model, - lora_rank=LORA_RANK, - target_modules=target_modules, - use_gradient_checkpointing=False, - random_state=42, + lora_rank = LORA_RANK, + target_modules = target_modules, + use_gradient_checkpointing = False, + random_state = 42, ) - prompt = tokenizer.apply_chat_template( - PROMPTS[0], tokenize=False, add_generation_prompt=True - ) + prompt = tokenizer.apply_chat_template(PROMPTS[0], tokenize = False, add_generation_prompt = True) with header_footer_context("Test Prompt"): print(prompt) @@ -110,26 +108,26 @@ def get_unsloth_peft_model( with header_footer_context("GRPO config and trainer"): training_args = GRPOConfig( - learning_rate=5e-6, - optim="adamw_torch", - temperature=TEMPERATURE, - top_k=50, - logging_steps=1, - per_device_train_batch_size=1, - gradient_accumulation_steps=NUM_GENERATIONS, - num_generations=NUM_GENERATIONS, - max_completion_length=MAX_COMPLETION_LENGTH, - max_steps=MAX_STEPS, - report_to="none", + learning_rate = 5e-6, + optim = "adamw_torch", + temperature = TEMPERATURE, + top_k = 50, + logging_steps = 1, + per_device_train_batch_size = 1, + gradient_accumulation_steps = NUM_GENERATIONS, + num_generations = NUM_GENERATIONS, + max_completion_length = MAX_COMPLETION_LENGTH, + max_steps = MAX_STEPS, + report_to = "none", ) print(training_args) trainer = GRPOTrainer( - model=model, - processing_class=tokenizer, - reward_funcs=[length_reward_func], - args=training_args, - train_dataset=dataset, + model = model, + processing_class = tokenizer, + reward_funcs = [length_reward_func], + args = training_args, + train_dataset = dataset, ) with header_footer_context("GRPO train (vLLM rollout)"): From 9c4d5923020c7ca6243dbe708f42f60edc373789 Mon Sep 17 00:00:00 2001 From: JoshuaL3000 <112940391+JoshuaL3000@users.noreply.github.com> Date: Wed, 15 Jul 2026 15:45:00 +0800 Subject: [PATCH 3/4] Update tests/fast_inference/test_fast_inference.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/fast_inference/test_fast_inference.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/fast_inference/test_fast_inference.py b/tests/fast_inference/test_fast_inference.py index d60e41eb10..75406e8a9c 100644 --- a/tests/fast_inference/test_fast_inference.py +++ b/tests/fast_inference/test_fast_inference.py @@ -11,6 +11,7 @@ sys.path.insert(0, str(REPO_ROOT)) from unsloth import FastLanguageModel +from unsloth.models._utils import is_bfloat16_supported import torch from datasets import Dataset from trl import GRPOConfig, GRPOTrainer From 92b44bf9ccf795fa4971f11d8157233fba5a9851 Mon Sep 17 00:00:00 2001 From: JoshuaL3000 <112940391+JoshuaL3000@users.noreply.github.com> Date: Wed, 15 Jul 2026 15:45:11 +0800 Subject: [PATCH 4/4] Update tests/fast_inference/test_fast_inference.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/fast_inference/test_fast_inference.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/fast_inference/test_fast_inference.py b/tests/fast_inference/test_fast_inference.py index 75406e8a9c..841a253453 100644 --- a/tests/fast_inference/test_fast_inference.py +++ b/tests/fast_inference/test_fast_inference.py @@ -55,8 +55,8 @@ def length_reward_func(completions, **kwargs) -> list[float]: def get_unsloth_peft_model( model, lora_rank: int, - target_modules: list[str] = "all-linear", - use_gradient_checkpointing: str = False, + target_modules = "all-linear", + use_gradient_checkpointing = False, random_state: int = 42, ): return FastLanguageModel.get_peft_model(