A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 19:45 UTC
[Bugfix] Fix JambaForCausalLM LoRA (#14370)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Mar 7, 2025 at 14:05 UTC
ddd1ef66ec3f40525b8726885b766985b1a82478
4 files changed
+35
-83
tests/lora/conftest.py
-24
@@ -6,7 +6,6 @@ from typing import TypedDict
6
from unittest.mock import MagicMock, patch
7
8
import pytest
9
-import safetensors
9
import torch
10
import torch.nn as nn
11
from huggingface_hub import snapshot_download
@@ -191,29 +190,6 @@ def mixtral_lora_files_all_target_modules():
190
return snapshot_download(repo_id="dyang415/mixtral-lora-v0")
191
192
194
-@pytest.fixture(scope="session")
195
-def jamba_lora_files():
196
- # some of the adapters have unnecessary weights for serving,
197
- # hence we remove them
198
- def remove_unnecessary_weights(path):
199
- lora_path = f"{adapter_path}/adapter_model.safetensors"
200
- tensors = safetensors.torch.load_file(lora_path)
201
- nonlora_keys = []
202
- for k in list(tensors.keys()):
203
- if "lora" not in k:
204
- nonlora_keys.append(k)
205
- for k in nonlora_keys:
206
- del tensors[k]
207
- safetensors.torch.save_file(tensors, lora_path)
208
-
209
- adapter_path = snapshot_download(
210
- repo_id=
211
- "hf-100/Jamba-1.5-mini-Spellbound-StoryWriter-0.1-6583896-ckpt53-lora")
212
-
213
- remove_unnecessary_weights(adapter_path)
214
- return adapter_path
215
-
216
-
193
@pytest.fixture(scope="session")
194
def gemma_lora_files():
195
return snapshot_download(repo_id="wskwon/gemma-7b-test-lora")
tests/lora/test_jamba.py
deleted
-54
@@ -1,54 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-
3
-import pytest
4
-import torch
5
-
6
-import vllm
7
-from vllm.lora.request import LoRARequest
8
-
9
-MODEL_PATH = "ai21labs/AI21-Jamba-1.5-Mini"
10
-
11
-MAX_TOKENS = 40
12
-
13
-
14
-def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int,
15
- prompts: list[str]) -> list[str]:
16
-
17
- sampling_params = vllm.SamplingParams(temperature=0, max_tokens=MAX_TOKENS)
18
- outputs = llm.generate(
19
- prompts,
20
- sampling_params,
21
- lora_request=LoRARequest(str(lora_id), lora_id, lora_path)
22
- if lora_id else None)
23
- # Print the outputs.
24
- generated_texts: list[str] = []
25
- for output in outputs:
26
- prompt = output.prompt
27
- generated_text = output.outputs[0].text.strip()
28
- generated_texts.append(generated_text)
29
- print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
30
- return generated_texts
31
-
32
-
33
-@pytest.mark.parametrize("tp_size", [4])
34
-def test_jamba_lora(jamba_lora_files, tp_size):
35
- """Original test, the LoRA model has the common target modules, not all"""
36
- if torch.cuda.device_count() < tp_size:
37
- pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}")
38
-
39
- prompts = ["Write a story about a sheep and a goat."]
40
-
41
- llm = vllm.LLM(
42
- MODEL_PATH,
43
- enable_lora=True,
44
- max_num_seqs=16,
45
- max_loras=4,
46
- distributed_executor_backend="ray",
47
- tensor_parallel_size=tp_size,
48
- )
49
-
50
- expected_jamba_output = [
51
- """Once upon a time, in a lush green meadow, there lived a sheep named Clara and a goat named Billy. Clara was a gentle creature, always nibbling on the soft grass and humming""" # noqa: E501
52
- ]
53
- assert do_sample(llm, jamba_lora_files, lora_id=1,
54
- prompts=prompts) == expected_jamba_output
tests/lora/test_layers.py
+3
@@ -632,6 +632,7 @@ def test_linear_replicated(dist_init, num_loras, device, stage,
632
633
id_to_index = get_random_id_to_index(num_loras, max_loras)
634
linear, lora_linear = create_random_linear_replicated_layer()
635
+ assert torch.equal(linear.weight, lora_linear.weight)
636
lora_linear.set_mapping(punica_wrapper)
637
lora_dict, _ = populate_loras(
638
id_to_index,
@@ -757,6 +758,7 @@ def test_linear_parallel(dist_init, num_loras, orientation, fully_shard,
758
759
id_to_index = get_random_id_to_index(num_loras, max_loras)
760
linear, lora_linear = create_random_linear_parallel_layer()
761
+ assert torch.equal(linear.weight, lora_linear.weight)
762
lora_linear.set_mapping(punica_wrapper)
763
lora_dict, _ = populate_loras(
764
id_to_index,
@@ -904,6 +906,7 @@ def test_column_parallel_packed(dist_init, num_loras, repeats, fully_shard,
906
id_to_index = get_random_id_to_index(num_loras, max_loras)
907
908
linear, lora_linear = create_column_parallel_packed_layer()
909
+ assert torch.equal(linear.weight, lora_linear.weight)
910
lora_linear.set_mapping(punica_wrapper)
911
lora_dict, sublora_dict = populate_loras(
912
id_to_index,
vllm/lora/layers.py
+32
-5
@@ -274,6 +274,10 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
274
) -> bool:
275
return type(source_layer) is VocabParallelEmbedding
276
277
+ @property
278
+ def weight(self):
279
+ return self.base_layer.weight
280
+
281
282
class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
283
@@ -409,6 +413,34 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
413
self.output_slices)
414
return output
415
416
+ @property
417
+ def weight(self) -> torch.Tensor:
418
+
419
+ # unquantizedLinear
420
+ if hasattr(self.base_layer, "weight"):
421
+ return self.base_layer.weight
422
+ # Compressed Tensor
423
+ elif hasattr(self.base_layer, "weight_packed"):
424
+ return self.base_layer.weight_packed
425
+ # GPTQ/AWQ
426
+ elif hasattr(self.base_layer, "qweight"):
427
+ return self.base_layer.qweight
428
+ # marlin
429
+ elif hasattr(self.base_layer, "B"):
430
+ return self.base_layer.B
431
+ # HQQ marlin
432
+ elif hasattr(self.base_layer, "W_q"):
433
+ return self.base_layer.W_q
434
+ else:
435
+ raise ValueError(f"Unsupported base layer: {self.base_layer}")
436
+
437
+ @property
438
+ def bias(self) -> Optional[torch.Tensor]:
439
+ if hasattr(self.base_layer, "bias"):
440
+ return self.base_layer.bias
441
+ else:
442
+ return None
443
+
444
445
class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):
446
@@ -902,11 +934,6 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):
934
935
return output, output_bias
936
905
- @property
906
- def weight(self):
907
- return (self.base_layer.weight if hasattr(self.base_layer, "weight")
908
- else self.base_layer.qweight)
909
-
937
@classmethod
938
@_not_fully_sharded_can_replace
939
def can_replace_layer(