| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for Qwen3 unembed LoRA support. |
| 4 | |
| 5 | This test creates synthetic LoRA weights that include lm_head (output embedding) |
| 6 | to verify that Qwen3 properly supports LoRA on the unembed/lm_head layer. |
| 7 | """ |
| 8 | |
| 9 | import json |
| 10 | import os |
| 11 | import tempfile |
| 12 | |
| 13 | import numpy as np |
| 14 | import torch |
| 15 | from safetensors.torch import save_file |
| 16 | |
| 17 | from vllm import LLM, SamplingParams |
| 18 | from vllm.lora.request import LoRARequest |
| 19 | |
| 20 | MODEL_PATH = "Qwen/Qwen3-0.6B" |
| 21 | HIDDEN_SIZE = 1024 |
| 22 | VOCAB_SIZE = 151936 |
| 23 | |
| 24 | |
| 25 | def create_qwen3_lora_with_lm_head(save_dir: str, rank: int = 8) -> None: |
| 26 | """Create synthetic Qwen3 LoRA weights with lm_head.""" |
| 27 | lora_weights = {} |
| 28 | for module in ["q_proj", "v_proj"]: |
| 29 | lora_A = torch.from_numpy( |
| 30 | np.random.randn(rank, HIDDEN_SIZE).astype(np.float16) * 0.01 |
| 31 | ) |
| 32 | lora_B = torch.zeros(HIDDEN_SIZE, rank, dtype=torch.float16) |
| 33 | key_prefix = f"base_model.model.model.layers.0.self_attn.{module}" |
| 34 | lora_weights[f"{key_prefix}.lora_A.weight"] = lora_A |
| 35 | lora_weights[f"{key_prefix}.lora_B.weight"] = lora_B |
| 36 | |
| 37 | # lm_head LoRA weights |
| 38 | lora_weights["base_model.model.lm_head.lora_A.weight"] = torch.from_numpy( |
| 39 | np.random.randn(rank, HIDDEN_SIZE).astype(np.float16) * 0.01 |
| 40 | ) |
| 41 | lora_weights["base_model.model.lm_head.lora_B.weight"] = torch.zeros( |
| 42 | VOCAB_SIZE, rank, dtype=torch.float16 |
| 43 | ) |
| 44 | |
| 45 | adapter_config = { |
| 46 | "peft_type": "LORA", |
| 47 | "base_model_name_or_path": MODEL_PATH, |
| 48 | "task_type": "CAUSAL_LM", |
| 49 | "inference_mode": True, |
| 50 | "r": rank, |
| 51 | "lora_alpha": rank * 2, |
| 52 | "lora_dropout": 0.0, |
| 53 | "bias": "none", |
| 54 | "target_modules": ["q_proj", "v_proj", "lm_head"], |
| 55 | } |
| 56 | |
| 57 | os.makedirs(save_dir, exist_ok=True) |
| 58 | with open(os.path.join(save_dir, "adapter_config.json"), "w") as f: |
| 59 | json.dump(adapter_config, f) |
| 60 | save_file(lora_weights, os.path.join(save_dir, "adapter_model.safetensors")) |
| 61 | |
| 62 | |
| 63 | def test_qwen3_unembed_lora(): |
| 64 | """Verify Qwen3 can load and generate with LoRA adapters with lm_head.""" |
| 65 | with tempfile.TemporaryDirectory() as tmpdir: |
| 66 | # Initialize engine first (before creating torch tensors) |
| 67 | llm = LLM( |
| 68 | model=MODEL_PATH, |
| 69 | enable_lora=True, |
| 70 | max_loras=4, |
| 71 | max_lora_rank=8, |
| 72 | max_model_len=128, |
| 73 | gpu_memory_utilization=0.8, |
| 74 | enforce_eager=True, |
| 75 | ) |
| 76 | |
| 77 | # Create LoRA weights after engine init |
| 78 | create_qwen3_lora_with_lm_head(tmpdir, rank=8) |
| 79 | |
| 80 | lora_request = LoRARequest("lm_head_lora", 1, tmpdir) |
| 81 | llm.llm_engine.add_lora(lora_request) |
| 82 | |
| 83 | assert 1 in llm.llm_engine.list_loras(), "lm_head LoRA should be loaded" |
| 84 | |
| 85 | # Test generation |
| 86 | sampling_params = SamplingParams(temperature=0, max_tokens=32) |
| 87 | prompts = ["Hello, my name is"] |
| 88 | |
| 89 | # Generate with base model (no LoRA) |
| 90 | base_outputs = llm.generate(prompts, sampling_params, use_tqdm=False) |
| 91 | assert len(base_outputs) == 1 |
| 92 | assert len(base_outputs[0].outputs[0].text) > 0 |
| 93 | |
| 94 | # Generate with lm_head LoRA |
| 95 | lora_outputs = llm.generate( |
| 96 | prompts, sampling_params, lora_request=lora_request, use_tqdm=False |
| 97 | ) |
| 98 | assert len(lora_outputs) == 1 |
| 99 | assert len(lora_outputs[0].outputs[0].text) > 0 |