@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9875be6431
Fork 0

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 12, 19:45 UTC

[LoRA][2/2]Remove LoRA extra vocab (#28545)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Nov 21, 2025 at 09:46 UTC 9875be6431872b513a8554c518e48ad79eba4656
28 files changed +133 -528
tests/lora/conftest.py
+10
@@ -250,6 +250,16 @@ def olmoe_lora_files():
250 return snapshot_download(repo_id="jeeejeee/olmoe-instruct-text2sql-spider")
251
252
253 +@pytest.fixture(scope="session")
254 +def qwen3_lora_files():
255 + return snapshot_download(repo_id="charent/self_cognition_Alice")
256 +
257 +
258 +@pytest.fixture(scope="session")
259 +def llama32_lora_files():
260 + return snapshot_download(repo_id="jeeejeee/llama32-3b-text2sql-spider")
261 +
262 +
263 @pytest.fixture
264 def reset_default_device():
265 """
tests/lora/test_layers.py
+13 -176
@@ -136,7 +136,6 @@ def populate_loras(
136 id_to_index: list[int | None],
137 layer: BaseLayerWithLoRA,
138 layer_weights: torch.Tensor,
139 - generate_embeddings_tensor: int = 0,
139 repeats: int = 1,
140 ) -> tuple[dict[int, LoRALayerWeights], dict[int, list[LoRALayerWeights]]]:
141 """This method populates the lora layers with lora weights.
@@ -148,8 +147,6 @@ def populate_loras(
147 layer: the LoRAlayer to populate.
148 layer_weights: the PyTorch tensor containing the layer's
149 weights.
151 - generate_embeddings_tensor: whether to generate an
152 - embeddings tensor for each LoRA.
150 repeats: must only be set for column parallel packed
151 layers. Indicates the number of loras to compose
152 together to create a single lora layer.
@@ -171,7 +168,6 @@ def populate_loras(
168 sublora = DummyLoRAManager(layer_weights.device).init_random_lora(
169 module_name=f"fake_{i}",
170 weight=layer_weights,
174 - generate_embeddings_tensor=generate_embeddings_tensor,
171 )
172 sublora.lora_b = sublora.lora_b[
173 (sublora_len * i) : (sublora_len * (i + 1)), :
@@ -185,7 +181,6 @@ def populate_loras(
181 slot_idx,
182 lora_a=lora.lora_a,
183 lora_b=lora.lora_b,
188 - embeddings_tensor=lora.embeddings_tensor,
184 )
185
186 lora_dict[lora_id] = lora
@@ -306,7 +301,6 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:
301 id_to_index,
302 max_loras,
303 vocab_size,
309 - lora_config.lora_extra_vocab_size,
304 )
305
306 lora_result = lora_embedding(torch.cat(inputs))
@@ -344,7 +338,6 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:
338 id_to_index,
339 max_loras,
340 vocab_size,
347 - lora_config.lora_extra_vocab_size,
341 )
342
343 lora_result = lora_embedding(torch.cat(inputs))
@@ -354,149 +347,6 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:
347 torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
348
349
357 -@torch.inference_mode()
358 -# @pytest.mark.skip(
359 -# reason="Fails when loras are in any slot other than the first.")
360 -@pytest.mark.parametrize("num_loras", [1, 2, 4])
361 -@pytest.mark.parametrize("device", DEVICES)
362 -@pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 128000])
363 -@pytest.mark.parametrize("stage", STAGES)
364 -def test_embeddings_with_new_embeddings(
365 - dist_init, num_loras, device, vocab_size, stage
366 -) -> None:
367 - if current_platform.is_cuda_alike():
368 - torch.cuda.set_device(device)
369 -
370 - torch.set_default_device(device)
371 - max_loras = 8
372 - punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)
373 - assert check_punica_wrapper(punica_wrapper)
374 - lora_config = LoRAConfig(
375 - max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16
376 - )
377 -
378 - def create_random_embedding_layer():
379 - embedding = VocabParallelEmbedding(vocab_size, 256)
380 - embedding_data = torch.rand_like(embedding.weight.data)
381 - embedding.weight.data = embedding_data
382 - embedding.weight.data[vocab_size:, :] = 0
383 - expanded_embedding = VocabParallelEmbedding(
384 - vocab_size + lora_config.lora_extra_vocab_size * max_loras,
385 - 256,
386 - org_num_embeddings=vocab_size,
387 - )
388 - expanded_embedding.weight.data[:vocab_size, :] = embedding_data
389 - # We need to deepcopy the embedding as it will be modified
390 - # in place
391 - lora_embedding = VocabParallelEmbeddingWithLoRA(deepcopy(expanded_embedding))
392 - lora_embedding.create_lora_weights(max_loras, lora_config)
393 -
394 - return expanded_embedding, lora_embedding
395 -
396 - for i in range(NUM_RANDOM_SEEDS):
397 - set_random_seed(i)
398 -
399 - id_to_index = get_random_id_to_index(num_loras, max_loras)
400 - expanded_embedding, lora_embedding = create_random_embedding_layer()
401 - lora_dict, _ = populate_loras(
402 - id_to_index,
403 - layer=lora_embedding,
404 - layer_weights=torch.zeros(
405 - (256, vocab_size + lora_config.lora_extra_vocab_size)
406 - ),
407 - generate_embeddings_tensor=256,
408 - )
409 -
410 - lora_embedding.set_mapping(punica_wrapper)
411 - # All embeddings tensors have the same shape.
412 - embeddings_tensors = [
413 - lora_dict[id].embeddings_tensor for id in sorted(lora_dict.keys())
414 - ]
415 - embeddings_tensor_len = embeddings_tensors[0].shape[0]
416 -
417 - # Add empty embeddings_tensors for unoccupied lora slots.
418 - for _ in range(max_loras - len(embeddings_tensors)):
419 - embeddings_tensors.append(torch.zeros(embeddings_tensors[0].shape))
420 -
421 - inputs, index_mapping, prompt_mapping = create_random_inputs(
422 - active_lora_ids=list(lora_dict.keys()),
423 - num_inputs=num_loras * 3,
424 - input_size=(200,),
425 - input_range=(1, vocab_size),
426 - device=device,
427 - )
428 - lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
429 - punica_wrapper.update_metadata(
430 - lora_mapping,
431 - id_to_index,
432 - max_loras,
433 - vocab_size,
434 - lora_config.lora_extra_vocab_size,
435 - )
436 - original_inputs = deepcopy(inputs)
437 -
438 - # Force some of the inputs to be in the extended embeddings range
439 - # to guarantee that their behavior is tested.
440 - for input_, original_input_, lora_id in zip(
441 - inputs, original_inputs, prompt_mapping
442 - ):
443 - embedding_id = lora_id - 1
444 - input_[-1] = vocab_size + (embedding_id * embeddings_tensor_len)
445 - original_input_[-1] = vocab_size
446 - input_[-2] = vocab_size + ((embedding_id + 1) * embeddings_tensor_len - 1)
447 - original_input_[-2] = vocab_size + embeddings_tensor_len - 1
448 -
449 - expanded_embedding.weight[
450 - vocab_size : vocab_size + (embeddings_tensor_len * max_loras)
451 - ] = torch.cat(embeddings_tensors)
452 -
453 - lora_result = lora_embedding(torch.cat(original_inputs))
454 -
455 - expected_results: list[torch.Tensor] = []
456 - for input_, original_input_, lora_id in zip(
457 - inputs, original_inputs, prompt_mapping
458 - ):
459 - lora = lora_dict[lora_id]
460 - result = expanded_embedding(input_)
461 - after_a = F.embedding(
462 - original_input_,
463 - lora.lora_a.T,
464 - )
465 - result += after_a @ lora.lora_b.T
466 - expected_results.append(result)
467 - expected_result = torch.cat(expected_results)
468 -
469 - rtol, atol = TOLERANCES[lora_result.dtype]
470 - torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
471 -
472 - # Check that resetting the lora weights succeeds
473 -
474 - for slot_idx in range(max_loras):
475 - lora_embedding.reset_lora(slot_idx)
476 -
477 - inputs, index_mapping, prompt_mapping = create_random_inputs(
478 - active_lora_ids=[0],
479 - num_inputs=num_loras * 3,
480 - input_size=(200,),
481 - input_range=(1, vocab_size),
482 - device=device,
483 - )
484 - original_inputs = deepcopy(inputs)
485 - lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
486 - punica_wrapper.update_metadata(
487 - lora_mapping,
488 - id_to_index,
489 - max_loras,
490 - vocab_size,
491 - lora_config.lora_extra_vocab_size,
492 - )
493 - lora_result = lora_embedding(torch.cat(original_inputs))
494 - expected_result = expanded_embedding(torch.cat(inputs))
495 -
496 - rtol, atol = TOLERANCES[lora_result.dtype]
497 - torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
498 -
499 -
350 @torch.inference_mode()
351 @pytest.mark.parametrize("num_loras", [1, 2, 4])
352 @pytest.mark.parametrize("device", DEVICES)
@@ -518,16 +368,13 @@ def test_lm_head_logits_processor(
368
369 def _pretest():
370 linear = ParallelLMHead(
521 - vocab_size + lora_config.lora_extra_vocab_size,
522 - 1024,
523 - vocab_size,
371 + num_embeddings=vocab_size,
372 + embedding_dim=1024,
373 params_dtype=torch.float16,
374 )
375 linear.weight.data = torch.rand_like(linear.weight.data)
376 linear.weight.data[:, vocab_size:] = 0
528 - logits_processor = LogitsProcessor(
529 - vocab_size + lora_config.lora_extra_vocab_size, vocab_size
530 - )
377 + logits_processor = LogitsProcessor(vocab_size)
378 lora_logits_processor = LogitsProcessorWithLoRA(
379 logits_processor, 1024, linear.weight.dtype, linear.weight.device, None
380 )
@@ -541,15 +388,12 @@ def test_lm_head_logits_processor(
388 id_to_index = get_random_id_to_index(num_loras, max_loras)
389 linear, logits_processor, lora_logits_processor = _pretest()
390 lora_logits_processor.set_mapping(punica_wrapper)
544 - # NOTE: all the generated loras share the same embeddings tensor.
391 +
392 lora_dict, _ = populate_loras(
393 id_to_index,
394 layer=lora_logits_processor,
395 layer_weights=linear.weight,
549 - generate_embeddings_tensor=1024,
396 )
551 - embeddings_tensor = list(lora_dict.values())[0].embeddings_tensor
552 - embeddings_tensor_len = embeddings_tensor.shape[0]
397
398 inputs, index_mapping, prompt_mapping = create_random_inputs(
399 active_lora_ids=list(lora_dict.keys()),
@@ -565,7 +409,6 @@ def test_lm_head_logits_processor(
409 id_to_index,
410 max_loras,
411 vocab_size,
568 - lora_config.lora_extra_vocab_size,
412 )
413 input_ = torch.rand(20, 1024)
414
@@ -575,23 +418,16 @@ def test_lm_head_logits_processor(
418
419 original_lm_head = deepcopy(linear)
420
578 - linear.weight[
579 - logits_processor.org_vocab_size : logits_processor.org_vocab_size
580 - + embeddings_tensor_len
581 - ] = embeddings_tensor
582 -
583 - logits_processor.org_vocab_size = vocab_size + lora_config.lora_extra_vocab_size
421 expected_results: list[torch.Tensor] = []
422 for input_, lora_id in zip(inputs, prompt_mapping):
423 lora = lora_dict[lora_id]
424 result = logits_processor._get_logits(
425 hidden_states=input_, lm_head=linear, embedding_bias=None
426 )
590 - result[:, vocab_size + embeddings_tensor_len :] = float("-inf")
427 +
428 result += input_ @ lora.lora_a.T @ lora.lora_b.T * lora.scaling
429 expected_results.append(result)
430 expected_result = torch.cat(expected_results)
594 - logits_processor.org_vocab_size = vocab_size
431
432 # Check that resetting the lora weights succeeds
433
@@ -612,7 +448,6 @@ def test_lm_head_logits_processor(
448 id_to_index,
449 max_loras,
450 vocab_size,
615 - lora_config.lora_extra_vocab_size,
451 )
452
453 lora_result = lora_logits_processor._get_logits(
@@ -694,7 +529,6 @@ def test_linear_replicated(
529 id_to_index,
530 max_loras,
531 512,
697 - lora_config.lora_extra_vocab_size,
532 )
533
534 lora_result = lora_linear(torch.cat(inputs))[0]
@@ -726,7 +560,10 @@ def test_linear_replicated(
560 lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
561
562 punica_wrapper.update_metadata(
729 - lora_mapping, id_to_index, max_loras, 512, lora_config.lora_extra_vocab_size
563 + lora_mapping,
564 + id_to_index,
565 + max_loras,
566 + 512,
567 )
568
569 lora_result = lora_linear(torch.cat(inputs))[0]
@@ -817,7 +654,6 @@ def test_linear_parallel(
654 id_to_index,
655 max_loras,
656 512,
820 - lora_config.lora_extra_vocab_size,
657 )
658
659 lora_result = lora_linear(torch.cat(inputs))[0]
@@ -849,7 +685,10 @@ def test_linear_parallel(
685 lora_mapping = LoRAMapping(index_mapping, prompt_mapping, is_prefill=stage)
686
687 punica_wrapper.update_metadata(
852 - lora_mapping, id_to_index, max_loras, 512, lora_config.lora_extra_vocab_size
688 + lora_mapping,
689 + id_to_index,
690 + max_loras,
691 + 512,
692 )
693
694 lora_result = lora_linear(torch.cat(inputs))[0]
@@ -963,7 +802,6 @@ def test_column_parallel_packed(
802 id_to_index,
803 max_loras,
804 512,
966 - lora_config.lora_extra_vocab_size,
805 )
806
807 lora_result = lora_linear(torch.cat(inputs))[0]
@@ -1000,7 +838,6 @@ def test_column_parallel_packed(
838 id_to_index,
839 max_loras,
840 512,
1003 - lora_config.lora_extra_vocab_size,
841 )
842
843 lora_result = lora_linear(torch.cat(inputs))[0]
tests/lora/test_llama_tp.py
+49 -35
@@ -13,17 +13,27 @@ from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
13
14 from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test
15
16 -MODEL_PATH = "meta-llama/Llama-2-7b-hf"
16 +PROMPT_TEMPLATE = """<|eot_id|><|start_header_id|>user<|end_header_id|>
17 +I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
18 +"
19 +##Instruction:
20 +candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
21 +Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
22 +The People_ID of candidate is the foreign key of People_ID of people.
23 +###Input:
24 +{context}
25 +###Response:<|eot_id|><|start_header_id|>assistant<|end_header_id|>
26 +""" # noqa: E501
27
28 EXPECTED_LORA_OUTPUT = [
19 - " SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport' ", # noqa: E501
20 - " SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone' ",
21 - " SELECT one_mora FROM table_name_95 WHERE gloss = 'low tone mora with a gloss of /˩okiru/' [òkìɽɯ́] AND accented_mora = 'low tone mora with a gloss of /˩okiru/' [òkìɽɯ́] ", # noqa: E501
22 - " SELECT sex FROM people WHERE people_id IN (SELECT people_id FROM candidate GROUP BY sex ORDER BY COUNT(people_id) DESC LIMIT 1) ", # noqa: E501
23 - " SELECT pick FROM table_name_60 WHERE former_wnba_team = 'Minnesota Lynx' ",
24 - " SELECT womens_doubles FROM table_28138035_4 WHERE mens_singles = 'Werner Schlager' ", # noqa: E501
29 + "SELECT count(*) FROM candidate",
30 + "SELECT count(*) FROM candidate",
31 + "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
32 + "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1", # noqa: E501
33 ]
34
35 +MODEL_PATH = "meta-llama/Llama-3.2-3B-Instruct"
36 +
37
38 def do_sample(
39 llm: vllm.LLM,
@@ -32,18 +42,19 @@ def do_sample(
42 tensorizer_config_dict: dict | None = None,
43 ) -> list[str]:
44 prompts = [
35 - "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]", # noqa: E501
36 - "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]", # noqa: E501
37 - "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_95 (one_mora VARCHAR, gloss VARCHAR, accented_mora VARCHAR)\n\n question: What is the one mora for a low tone mora with a gloss of /˩okiru/ [òkìɽɯ́]? [/user] [assistant]", # noqa: E501
38 - "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE candidate (people_id VARCHAR, unsure_rate INTEGER); CREATE TABLE people (sex VARCHAR, people_id VARCHAR)\n\n question: which gender got the highest average uncertain ratio. [/user] [assistant]", # noqa: E501
39 - "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_60 (pick INTEGER, former_wnba_team VARCHAR)\n\n question: What pick was a player that previously played for the Minnesota Lynx? [/user] [assistant]", # noqa: E501
40 - "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_28138035_4 (womens_doubles VARCHAR, mens_singles VARCHAR)\n\n question: Name the women's doubles for werner schlager [/user] [assistant]", # noqa: E501
45 + PROMPT_TEMPLATE.format(context="How many candidates are there?"),
46 + PROMPT_TEMPLATE.format(context="Count the number of candidates."),
47 + PROMPT_TEMPLATE.format(
48 + context="Which poll resource provided the most number of candidate information?" # noqa: E501
49 + ),
50 + PROMPT_TEMPLATE.format(
51 + context="Return the poll resource associated with the most candidates."
52 + ),
53 ]
54
55 sampling_params = vllm.SamplingParams(
44 - temperature=0, max_tokens=256, skip_special_tokens=False, stop=["[/assistant]"]
56 + temperature=0, max_tokens=64, stop=["<|im_end|>"]
57 )
46 -
58 if tensorizer_config_dict is not None:
59 outputs = llm.generate(
60 prompts,
@@ -75,13 +86,15 @@ def do_sample(
86 return generated_texts
87
88
78 -def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None = None):
89 +def generate_and_test(
90 + llm, llama32_lora_files, tensorizer_config_dict: dict | None = None
91 +):
92 print("lora adapter created")
93 print("lora 1")
94 assert (
95 do_sample(
96 llm,
84 - sql_lora_files,
97 + llama32_lora_files,
98 tensorizer_config_dict=tensorizer_config_dict,
99 lora_id=1,
100 )
@@ -92,7 +105,7 @@ def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None =
105 assert (
106 do_sample(
107 llm,
95 - sql_lora_files,
108 + llama32_lora_files,
109 tensorizer_config_dict=tensorizer_config_dict,
110 lora_id=2,
111 )
@@ -104,51 +117,52 @@ def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None =
117
118 @create_new_process_for_each_test()
119 @pytest.mark.parametrize("cudagraph_specialize_lora", [True, False])
107 -def test_llama_lora(sql_lora_files, cudagraph_specialize_lora: bool):
120 +def test_llama_lora(llama32_lora_files, cudagraph_specialize_lora: bool):
121 llm = vllm.LLM(
122 MODEL_PATH,
110 - tokenizer=sql_lora_files,
123 enable_lora=True,
124 # also test odd max_num_seqs
113 - max_num_seqs=13,
125 + max_num_seqs=7,
126 + max_model_len=1024,
127 max_loras=4,
128 compilation_config=vllm.config.CompilationConfig(
129 cudagraph_specialize_lora=cudagraph_specialize_lora,
130 ),
131 )
119 - generate_and_test(llm, sql_lora_files)
132 + generate_and_test(llm, llama32_lora_files)
133
134
135 @multi_gpu_test(num_gpus=4)
123 -def test_llama_lora_tp4(sql_lora_files):
136 +def test_llama_lora_tp4(llama32_lora_files):
137 llm = vllm.LLM(
138 MODEL_PATH,
126 - tokenizer=sql_lora_files,
139 enable_lora=True,
128 - max_num_seqs=16,
140 + max_num_seqs=7,
141 + max_model_len=1024,
142 max_loras=4,
143 tensor_parallel_size=4,
144 )
132 - generate_and_test(llm, sql_lora_files)
145 + generate_and_test(llm, llama32_lora_files)
146
147
148 @multi_gpu_test(num_gpus=4)
136 -def test_llama_lora_tp4_fully_sharded_loras(sql_lora_files):
149 +def test_llama_lora_tp4_fully_sharded_loras(llama32_lora_files):
150 llm = vllm.LLM(
151 MODEL_PATH,
139 - tokenizer=sql_lora_files,
152 enable_lora=True,
141 - max_num_seqs=16,
153 + max_num_seqs=8,
154 max_loras=4,
155 + max_model_len=1024,
156 tensor_parallel_size=4,
157 fully_sharded_loras=True,
158 )
146 - generate_and_test(llm, sql_lora_files)
159 + generate_and_test(llm, llama32_lora_files)
160
161
162 @multi_gpu_test(num_gpus=2)
163 def test_tp2_serialize_and_deserialize_lora(
151 - tmp_path, sql_lora_files, sql_lora_huggingface_id
164 + tmp_path,
165 + llama32_lora_files,
166 ):
167 # Run the tensorizing of the LoRA adapter and the model in a subprocess
168 # to guarantee cleanup
@@ -157,7 +171,7 @@ def test_tp2_serialize_and_deserialize_lora(
171 model_name = "model-rank-%03d.tensors"
172
173 model_ref = MODEL_PATH
160 - lora_path = sql_lora_huggingface_id
174 + lora_path = llama32_lora_files
175 suffix = "test"
176 try:
177 result = subprocess.run(
@@ -195,12 +209,12 @@ def test_tp2_serialize_and_deserialize_lora(
209
210 loaded_llm = LLM(
211 model=model_ref,
198 - tokenizer=sql_lora_files,
212 load_format="tensorizer",
213 enable_lora=True,
214 enforce_eager=True,
215 model_loader_extra_config=tensorizer_config,
203 - max_num_seqs=13,
216 + max_num_seqs=7,
217 + max_model_len=1024,
218 tensor_parallel_size=2,
219 max_loras=2,
220 )
@@ -211,7 +225,7 @@ def test_tp2_serialize_and_deserialize_lora(
225 print("lora 1")
226 assert (
227 do_sample(
214 - loaded_llm, sql_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1
228 + loaded_llm, llama32_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1
229 )
230 == EXPECTED_LORA_OUTPUT
231 )
tests/lora/test_lora_functions.py
+2 -2
@@ -13,8 +13,8 @@ from vllm.entrypoints.openai.api_server import (
13 from vllm.lora.request import LoRARequest
14 from vllm.v1.engine.llm_engine import LLMEngine
15
16 -MODEL_PATH = "meta-llama/Llama-2-7b-hf"
17 -LORA_MODULE_PATH = "yard1/llama-2-7b-sql-lora-test"
16 +MODEL_PATH = "Qwen/Qwen3-0.6B"
17 +LORA_MODULE_PATH = "charent/self_cognition_Alice"
18 LORA_RANK = 8
19
20
tests/lora/test_lora_manager.py
+1 -19
@@ -48,9 +48,6 @@ DEFAULT_DTYPE = torch.get_default_dtype()
48 @pytest.mark.parametrize("device", DEVICES)
49 def test_from_lora_tensors(sql_lora_files, device):
50 tensors = load_file(os.path.join(sql_lora_files, "adapter_model.safetensors"))
51 - new_embeddings = load_file(
52 - os.path.join(sql_lora_files, "new_embeddings.safetensors")
53 - )
51
52 peft_helper = PEFTHelper.from_local_dir(
53 sql_lora_files, max_position_embeddings=4096
@@ -60,7 +57,6 @@ def test_from_lora_tensors(sql_lora_files, device):
57 tensors,
58 peft_helper=peft_helper,
59 device=device,
63 - embeddings=new_embeddings,
60 embedding_modules=EMBEDDING_MODULES,
61 embedding_padding_modules=EMBEDDING_PADDING_MODULES,
62 )
@@ -76,18 +72,6 @@ def test_from_lora_tensors(sql_lora_files, device):
72 f"{lora.lora_a.shape=}, {lora.lora_b.shape=}"
73 )
74 assert lora.lora_a.shape[0] == 8
79 - embeddings_module = next(
80 - (k for k in EMBEDDING_MODULES if k in module_name), None
81 - )
82 - if embeddings_module:
83 - assert torch.equal(
84 - lora.embeddings_tensor,
85 - new_embeddings[EMBEDDING_MODULES[embeddings_module]].to(
86 - device=lora.embeddings_tensor.device
87 - ),
88 - )
89 - else:
90 - assert lora.embeddings_tensor is None
75
76
77 def create_lora(
@@ -552,9 +536,7 @@ def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path
536 worker_adapter_manager = WorkerLoRAManager(
537 vllm_config, device, EMBEDDING_MODULES, EMBEDDING_PADDING_MODULES
538 )
555 - worker_adapter_manager.vocab_size = (
556 - dummy_model_gate_up.unpadded_vocab_size - lora_config.lora_extra_vocab_size
557 - )
539 + worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
540 worker_adapter_manager.create_lora_manager(dummy_model_gate_up)
541
542 dummy_lora_files = f"{tmp_path}/lora_adapter"
tests/lora/test_worker.py
+5 -3
@@ -20,11 +20,12 @@ from vllm.lora.models import LoRAMapping
20 from vllm.lora.request import LoRARequest
21 from vllm.v1.worker.gpu_worker import Worker
22
23 +MODEL_PATH = "Qwen/Qwen3-0.6B"
24 NUM_LORAS = 16
25
26
27 @patch.dict(os.environ, {"RANK": "0"})
27 -def test_worker_apply_lora(sql_lora_files):
28 +def test_worker_apply_lora(qwen3_lora_files):
29 def set_active_loras(worker: Worker, lora_requests: list[LoRARequest]):
30 lora_mapping = LoRAMapping([], [])
31
@@ -34,9 +35,10 @@ def test_worker_apply_lora(sql_lora_files):
35
36 vllm_config = VllmConfig(
37 model_config=ModelConfig(
37 - "meta-llama/Llama-2-7b-hf",
38 + MODEL_PATH,
39 seed=0,
40 dtype="float16",
41 + max_model_len=127,
42 enforce_eager=True,
43 ),
44 load_config=LoadConfig(
@@ -73,7 +75,7 @@ def test_worker_apply_lora(sql_lora_files):
75 assert worker.list_loras() == set()
76
77 lora_requests = [
76 - LoRARequest(str(i + 1), i + 1, sql_lora_files) for i in range(NUM_LORAS)
78 + LoRARequest(str(i + 1), i + 1, qwen3_lora_files) for i in range(NUM_LORAS)
79 ]
80
81 set_active_loras(worker, lora_requests)
tests/lora/utils.py
-8
@@ -28,7 +28,6 @@ class DummyLoRAManager:
28 module_name: str,
29 weight: torch.Tensor,
30 rank: int = 8,
31 - generate_embeddings_tensor: int = 0,
31 ):
32 lora = LoRALayerWeights(
33 module_name,
@@ -41,13 +40,6 @@ class DummyLoRAManager:
40 [weight.shape[0], rank], dtype=weight.dtype, device=self._device
41 ),
42 )
44 - if generate_embeddings_tensor:
45 - lora.embeddings_tensor = torch.rand(
46 - 5,
47 - generate_embeddings_tensor,
48 - dtype=weight.dtype,
49 - device=self._device,
50 - )
43 self.set_module_lora(module_name, lora)
44
45 return lora
vllm/config/lora.py
+1 -17
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import hashlib
5 -from typing import TYPE_CHECKING, Any, ClassVar, Literal
5 +from typing import TYPE_CHECKING, Any, Literal
6
7 import torch
8 from pydantic import ConfigDict, Field, model_validator
@@ -11,7 +11,6 @@ from typing_extensions import Self
11
12 from vllm.config.utils import config
13 from vllm.logger import init_logger
14 -from vllm.platforms import current_platform
14
15 if TYPE_CHECKING:
16 from vllm.config import ModelConfig
@@ -46,19 +45,6 @@ class LoRAConfig:
45 `max_loras`."""
46 lora_dtype: torch.dtype | LoRADType = "auto"
47 """Data type for LoRA. If auto, will default to base model dtype."""
49 - lora_extra_vocab_size: LoRAExtraVocabSize = Field(
50 - default=256,
51 - deprecated=(
52 - "`lora_extra_vocab_size` is deprecated and will be removed "
53 - "in v0.12.0. Additional vocabulary support for "
54 - "LoRA adapters is being phased out."
55 - ),
56 - )
57 - """(Deprecated) Maximum size of extra vocabulary that can be present in a
58 - LoRA adapter. Will be removed in v0.12.0."""
59 - lora_vocab_padding_size: ClassVar[int] = (
60 - current_platform.get_lora_vocab_padding_size()
61 - )
48 default_mm_loras: dict[str, str] | None = None
49 """Dictionary mapping specific modalities to LoRA model paths; this field
50 is only applicable to multimodal models and should be leveraged when a
@@ -87,8 +73,6 @@ class LoRAConfig:
73 factors.append(self.max_loras)
74 factors.append(self.fully_sharded_loras)
75 factors.append(self.lora_dtype)
90 - factors.append(self.lora_extra_vocab_size)
91 - factors.append(self.lora_vocab_padding_size)
76
77 hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()
78 return hash_str
vllm/engine/arg_utils.py
-5
@@ -484,7 +484,6 @@ class EngineArgs:
484 fully_sharded_loras: bool = LoRAConfig.fully_sharded_loras
485 max_cpu_loras: int | None = LoRAConfig.max_cpu_loras
486 lora_dtype: str | torch.dtype | None = LoRAConfig.lora_dtype
487 - lora_extra_vocab_size: int = LoRAConfig.lora_extra_vocab_size
487
488 ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight
489 num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override
@@ -1011,9 +1010,6 @@ class EngineArgs:
1010 )
1011 lora_group.add_argument("--max-loras", **lora_kwargs["max_loras"])
1012 lora_group.add_argument("--max-lora-rank", **lora_kwargs["max_lora_rank"])
1014 - lora_group.add_argument(
1015 - "--lora-extra-vocab-size", **lora_kwargs["lora_extra_vocab_size"]
1016 - )
1013 lora_group.add_argument(
1014 "--lora-dtype",
1015 **lora_kwargs["lora_dtype"],
@@ -1680,7 +1676,6 @@ class EngineArgs:
1676 max_loras=self.max_loras,
1677 default_mm_loras=self.default_mm_loras,
1678 fully_sharded_loras=self.fully_sharded_loras,
1683 - lora_extra_vocab_size=self.lora_extra_vocab_size,
1679 lora_dtype=self.lora_dtype,
1680 max_cpu_loras=self.max_cpu_loras
1681 if self.max_cpu_loras and self.max_cpu_loras > 0
vllm/lora/layers/base.py
-1
@@ -44,7 +44,6 @@ class BaseLayerWithLoRA(nn.Module):
44 index: int,
45 lora_a: torch.Tensor,
46 lora_b: torch.Tensor,
47 - embeddings_tensor: torch.Tensor | None,
47 ):
48 """Overwrites lora tensors at index."""
49 ...
vllm/lora/layers/base_linear.py
-1
@@ -96,7 +96,6 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
96 index: int,
97 lora_a: torch.Tensor,
98 lora_b: torch.Tensor,
99 - embeddings_tensor: torch.Tensor | None,
99 ):
100 # Except for QKVParallelLinearWithLoRA and
101 # MergedColumnParallelLinearWithLoRA, all other linear LoRA layers
vllm/lora/layers/column_parallel_linear.py
-1
@@ -248,7 +248,6 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
248 index: int,
249 lora_a: torch.Tensor,
250 lora_b: torch.Tensor,
251 - embeddings_tensor: torch.Tensor | None,
251 ):
252 self.reset_lora(index)
253
vllm/lora/layers/fused_moe.py
-2
@@ -406,8 +406,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
406 index: int,
407 lora_a: torch.Tensor,
408 lora_b: torch.Tensor,
409 - embeddings_tensor: torch.Tensor | None,
410 - bias: torch.Tensor | None = None,
409 ):
410 """Overwrites lora tensors at index."""
411 self.reset_lora(index)
vllm/lora/layers/logits_processor.py
+2 -53
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import math
4
5 import torch
6 import torch.nn as nn
@@ -108,22 +107,13 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
107 (
108 max_loras,
109 1,
111 - # Pad for kernel compatibility
112 - math.ceil(
113 - self.base_layer.vocab_size / lora_config.lora_vocab_padding_size
114 - )
115 - * lora_config.lora_vocab_padding_size,
110 + self.base_layer.vocab_size,
111 lora_config.max_lora_rank,
112 ),
113 dtype=lora_config.lora_dtype,
114 device=self.device,
115 )
121 - self.embeddings_tensors = torch.full(
122 - (max_loras, lora_config.lora_extra_vocab_size, self.hidden_size),
123 - fill_value=float("-inf"),
124 - dtype=self.dtype,
125 - device=self.device,
126 - )
116 +
117 if self.sharded_to_full_mapping is not None:
118 self.sharded_to_full_mapping_gpu = torch.tensor(
119 self.sharded_to_full_mapping, device=self.device, dtype=torch.long
@@ -134,14 +124,12 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
124 def reset_lora(self, index: int):
125 self.lora_a_stacked[index] = 0
126 self.lora_b_stacked[index] = 0
137 - self.embeddings_tensors[index] = float("-inf")
127
128 def set_lora(
129 self,
130 index: int,
131 lora_a: torch.Tensor,
132 lora_b: torch.Tensor,
144 - embeddings_tensor: torch.Tensor | None,
133 ):
134 self.reset_lora(index)
135 self.lora_a_stacked[index, 0, : lora_a.shape[0], : lora_a.shape[1]].copy_(
@@ -150,12 +138,6 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
138 self.lora_b_stacked[index, 0, : lora_b.shape[0], : lora_b.shape[1]].copy_(
139 lora_b, non_blocking=True
140 )
153 - if embeddings_tensor is not None:
154 - self.embeddings_tensors[
155 - index,
156 - : embeddings_tensor.shape[0],
157 - : embeddings_tensor.shape[1],
158 - ] = embeddings_tensor
141
142 def _get_logits(
143 self,
@@ -193,39 +175,6 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
175 # token_id: [0, 1, 2, 3, 4, 5, -1, -1]
176 logits = logits[:, self.sharded_to_full_mapping_gpu]
177
196 - lora_logits = torch.empty(
197 - self.embeddings_tensors.shape[0] + 1,
198 - self.embeddings_tensors.shape[1],
199 - hidden_states.shape[0],
200 - dtype=self.embeddings_tensors.dtype,
201 - device=self.embeddings_tensors.device,
202 - )
203 - torch.matmul(self.embeddings_tensors, hidden_states.T, out=lora_logits[:-1])
204 -
205 - neg_inf, pos_inf = current_platform.get_infinity_values(lora_logits.dtype)
206 -
207 - lora_logits[-1] = neg_inf
208 - lora_logits = lora_logits.mT
209 - indices_padded = self.punica_wrapper.sampler_indices_padded
210 -
211 - if current_platform.is_tpu() or current_platform.is_xpu():
212 - indices_padded = indices_padded[: logits.size(0)]
213 -
214 - lora_logits = (
215 - lora_logits.reshape(
216 - lora_logits.shape[0] * lora_logits.shape[1],
217 - lora_logits.shape[2],
218 - )
219 - .index_select(0, indices_padded)
220 - .nan_to_num_(nan=neg_inf, posinf=pos_inf, neginf=neg_inf)
221 - )
222 -
223 - logits[
224 - :,
225 - self.base_layer.org_vocab_size : self.base_layer.org_vocab_size
226 - + lora_logits.shape[1],
227 - ] = lora_logits
228 -
178 lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_logits(
179 logits, hidden_states, self.lora_a_stacked, self.lora_b_stacked, 1.0
180 )
vllm/lora/layers/vocal_parallel_embedding.py
+2 -31
@@ -46,19 +46,10 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
46 self.embeddings_slice = None
47 self.embeddings_weights = None
48
49 - self.embeddings_tensors = torch.zeros(
50 - (
51 - max_loras,
52 - lora_config.lora_extra_vocab_size,
53 - self.base_layer.embedding_dim,
54 - ),
55 - dtype=self.base_layer.weight.dtype,
56 - device=self.base_layer.weight.device,
57 - )
49 self.lora_a_stacked = torch.zeros(
50 (
51 max_loras,
61 - self.base_layer.org_vocab_size + lora_config.lora_extra_vocab_size,
52 + self.base_layer.org_vocab_size,
53 lora_config.max_lora_rank,
54 ),
55 dtype=lora_config.lora_dtype,
@@ -82,14 +73,12 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
73 def reset_lora(self, index: int):
74 self.lora_a_stacked[index] = 0
75 self.lora_b_stacked[index] = 0
85 - self.embeddings_tensors[index] = 0
76
77 def set_lora(
78 self,
79 index: int,
80 lora_a: torch.Tensor,
81 lora_b: torch.Tensor,
92 - embeddings_tensor: torch.Tensor | None,
82 ):
83 self.reset_lora(index)
84 # NOTE self.lora_a_stacked is row-major, and lora_a is col-major,
@@ -100,36 +89,18 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
89 self.lora_b_stacked[index, 0, : lora_b.shape[0], : lora_b.shape[1]].copy_(
90 lora_b, non_blocking=True
91 )
103 - if embeddings_tensor is not None:
104 - self.embeddings_tensors[
105 - index,
106 - : embeddings_tensor.shape[0],
107 - : embeddings_tensor.shape[1],
108 - ].copy_(embeddings_tensor, non_blocking=True)
109 - if self.embeddings_slice is not None:
110 - # TODO(yard1): Optimize this copy, we don't need to copy
111 - # everything, just the modified part
112 - embeddings = self.embeddings_tensors.view(
113 - self.embeddings_tensors.shape[0] * self.embeddings_tensors.shape[1],
114 - self.embeddings_tensors.shape[2],
115 - )[self.embeddings_slice[0] : self.embeddings_slice[1]]
116 - assert self.embeddings_weights is not None
117 - self.embeddings_weights[: embeddings.shape[0]].copy_(embeddings)
92
93 def forward(self, x: torch.Tensor) -> torch.Tensor:
120 - added_tokens_mask = torch.where(x > self.base_layer.org_vocab_size - 1, 1, 0)
121 -
94 # NB: Don't use torch.narrow here. torch.narrow triggers some
95 # Dynamic Shape specialization in torch.compile
96 num_tokens = x.shape[0]
97 indices_1 = self.punica_wrapper._embeddings_indices[1][:num_tokens]
126 - indices_0 = self.punica_wrapper._embeddings_indices[0][:num_tokens]
98
99 full_lora_a_embeddings = F.embedding(
100 x + indices_1,
101 self.lora_a_stacked_2d,
102 )
132 - full_output = self.base_layer.forward(x + (indices_0 * added_tokens_mask))
103 + full_output = self.base_layer.forward(x)
104
105 full_output_org = full_output
106 if full_output.ndim == 3:
vllm/lora/lora_weights.py
-24
@@ -21,7 +21,6 @@ class LoRALayerWeights:
21 lora_alpha: int,
22 lora_a: torch.Tensor,
23 lora_b: torch.Tensor,
24 - embeddings_tensor: torch.Tensor | None = None,
24 scaling: float | None = None,
25 ) -> None:
26 self.module_name = module_name
@@ -29,7 +28,6 @@ class LoRALayerWeights:
28 self.lora_alpha = lora_alpha
29 self.lora_a = lora_a
30 self.lora_b = lora_b
32 - self.embeddings_tensor = embeddings_tensor
31
32 if scaling is None:
33 self.scaling = self.lora_alpha / self.rank
@@ -56,18 +54,11 @@ class LoRALayerWeights:
54 def is_packed(self) -> bool:
55 return False
56
59 - @property
60 - def extra_vocab_size(self) -> int:
61 - return (
62 - self.embeddings_tensor.shape[0] if self.embeddings_tensor is not None else 0
63 - )
64 -
57 @classmethod
58 def from_config(
59 cls,
60 module_name: str,
61 peft_helper: PEFTHelper,
70 - embeddings_tensor: torch.Tensor | None = None,
62 ) -> "LoRALayerWeights":
63 # lora_a and lora_b are set to None for config-based construction
64 return cls(
@@ -76,7 +67,6 @@ class LoRALayerWeights:
67 peft_helper.lora_alpha,
68 None,
69 None,
79 - embeddings_tensor,
70 peft_helper.vllm_lora_scaling_factor,
71 )
72
@@ -89,7 +79,6 @@ class LoRALayerWeights:
79 rank: int,
80 dtype: torch.dtype,
81 device: torch.types.Device,
92 - embeddings_tensor_dim: int | None = None,
82 ) -> "LoRALayerWeights":
83 pin_memory = str(device) == "cpu" and is_pin_memory_available()
84 lora_a = torch.zeros(
@@ -99,24 +88,12 @@ class LoRALayerWeights:
88 [output_dim, rank], dtype=dtype, device=device, pin_memory=pin_memory
89 )
90
102 - embeddings_tensor = (
103 - torch.rand(
104 - 10,
105 - embeddings_tensor_dim,
106 - dtype=dtype,
107 - device=device,
108 - pin_memory=pin_memory,
109 - )
110 - if embeddings_tensor_dim
111 - else None
112 - )
91 return cls(
92 module_name,
93 rank=rank,
94 lora_alpha=1,
95 lora_a=lora_a,
96 lora_b=lora_b,
119 - embeddings_tensor=embeddings_tensor,
97 )
98
99
@@ -139,7 +116,6 @@ class PackedLoRALayerWeights(LoRALayerWeights):
116 lora_a=lora_a,
117 lora_b=lora_b,
118 scaling=scaling, # type: ignore
142 - embeddings_tensor=None,
119 )
120 self.lora_alphas = lora_alphas
121 if scaling is None:
vllm/lora/models.py
+10 -44
@@ -21,6 +21,7 @@ from vllm.lora.utils import (
21 from_layer,
22 from_layer_logits_processor,
23 get_supported_lora_modules,
24 + is_base_embeddding_weights,
25 is_regex_target_modules,
26 parse_fine_tuned_lora_name,
27 process_packed_modules_mapping,
@@ -93,14 +94,6 @@ class LoRAModel:
94 loras=self.loras.copy(),
95 )
96
96 - @property
97 - def extra_vocab_size(self) -> int:
98 - return (
99 - max(lora.extra_vocab_size for lora in self.loras.values())
100 - if self.loras
101 - else 0
102 - )
103 -
97 def get_lora(self, module_name: str) -> LoRALayerWeights | None:
98 """Get LoRA for a given module by name"""
99 return self.loras.get(module_name, None)
@@ -117,7 +110,6 @@ class LoRAModel:
110 peft_helper: PEFTHelper,
111 device: str = "cuda",
112 dtype: torch.dtype | None = None,
120 - embeddings: dict[str, torch.Tensor] | None = None,
113 target_embedding_padding: int | None = None,
114 embedding_modules: dict[str, str] | None = None,
115 embedding_padding_modules: list[str] | None = None,
@@ -127,24 +119,14 @@ class LoRAModel:
119 pin_memory = str(device) == "cpu" and is_pin_memory_available()
120 loras: dict[str, LoRALayerWeights] = {}
121 for tensor_name, tensor in tensors.items():
122 + if is_base_embeddding_weights(tensor_name):
123 + continue
124 module_name, is_lora_a = parse_fine_tuned_lora_name(
125 tensor_name, weights_mapper
126 )
127 if module_name not in loras:
134 - lora_embeddings_tensor = None
135 - if embeddings:
136 - assert embedding_modules is not None
137 - embeddings_module = next(
138 - (k for k in embedding_modules if k in module_name), None
139 - )
140 - if embeddings_module:
141 - lora_embeddings_tensor = embeddings[
142 - embedding_modules[embeddings_module]
143 - ].to(device=device, dtype=dtype)
144 - if pin_memory:
145 - lora_embeddings_tensor = lora_embeddings_tensor.pin_memory()
128 loras[module_name] = LoRALayerWeights.from_config(
147 - module_name, peft_helper, lora_embeddings_tensor
129 + module_name, peft_helper
130 )
131
132 if is_lora_a:
@@ -206,15 +188,17 @@ class LoRAModel:
188 lora_tensor_path = os.path.join(lora_dir, "adapter_model.safetensors")
189 lora_bin_file_path = os.path.join(lora_dir, "adapter_model.bin")
190 lora_pt_file_path = os.path.join(lora_dir, "adapter_model.pt")
209 - new_embeddings_tensor_path = os.path.join(
210 - lora_dir, "new_embeddings.safetensors"
211 - )
212 - new_embeddings_bin_file_path = os.path.join(lora_dir, "new_embeddings.bin")
191 + # new_embeddings_tensor_path = os.path.join(
192 + # lora_dir, "new_embeddings.safetensors"
193 + # )
194 + # new_embeddings_bin_file_path = os.path.join(lora_dir, "new_embeddings.bin")
195 tensors: dict[str, torch.Tensor] = {}
196 unexpected_modules: list[list[str] | str] = []
197
198 def check_unexpected_modules(modules: dict):
199 for lora_module in modules.keys(): # noqa
200 + if is_base_embeddding_weights(lora_module):
201 + continue
202 module_name, _ = parse_fine_tuned_lora_name(lora_module, weights_mapper)
203 # Handle FSDP file format where experts.base_layer is the
204 # gate_up_proj and experts is the down_proj
@@ -300,21 +284,12 @@ class LoRAModel:
284 else:
285 raise ValueError(f"{lora_dir} doesn't contain tensors")
286
303 - embeddings = None
304 - if os.path.isfile(new_embeddings_tensor_path):
305 - embeddings = safetensors.torch.load_file(new_embeddings_tensor_path)
306 - elif os.path.isfile(new_embeddings_bin_file_path):
307 - embeddings = torch.load(
308 - new_embeddings_bin_file_path, map_location=device, weights_only=True
309 - )
310 -
287 return cls.from_lora_tensors(
288 lora_model_id=get_lora_id() if lora_model_id is None else lora_model_id,
289 tensors=tensors,
290 peft_helper=peft_helper,
291 device=device,
292 dtype=dtype,
317 - embeddings=embeddings,
293 target_embedding_padding=target_embedding_padding,
294 embedding_modules=embedding_modules,
295 embedding_padding_modules=embedding_padding_modules,
@@ -474,7 +449,6 @@ class LoRAModelManager:
449 index,
450 module_lora.lora_a,
451 module_lora.lora_b,
477 - module_lora.embeddings_tensor,
452 )
453 else:
454 module.reset_lora(index)
@@ -505,7 +479,6 @@ class LoRAModelManager:
479 self.lora_index_to_id,
480 self.lora_slots + 1,
481 self.vocab_size,
508 - self.lora_config.lora_extra_vocab_size,
482 )
483
484 def remove_all_adapters(self):
@@ -616,7 +589,6 @@ class LoRAModelManager:
589 if parts[-1] in embedding_modules:
590 input_dim = (
591 module.base_layer.org_vocab_size
619 - + self.lora_config.lora_extra_vocab_size
592 if hasattr(module.base_layer, "org_vocab_size")
593 else module.base_layer.weight.shape[1]
594 )
@@ -625,11 +597,6 @@ class LoRAModelManager:
597 if hasattr(module.base_layer, "embedding_dim")
598 else module.base_layer.weight.shape[0]
599 )
628 - embeddings_tensor_dim = (
629 - module.base_layer.embedding_dim
630 - if hasattr(module.base_layer, "embedding_dim")
631 - else module.base_layer.weight.shape[1]
632 - )
600 lora = LoRALayerWeights.create_dummy_lora_weights(
601 module_name,
602 input_dim,
@@ -637,7 +604,6 @@ class LoRAModelManager:
604 rank,
605 module.lora_a_stacked[0].dtype,
606 "cpu",
640 - embeddings_tensor_dim=embeddings_tensor_dim,
607 )
608 else:
609 lora = LoRALayerWeights.create_dummy_lora_weights(
vllm/lora/punica_wrapper/punica_base.py
+5 -6
@@ -31,7 +31,6 @@ class PunicaWrapperABC(ABC):
31 lora_index_to_id: list[int | None],
32 max_loras: int,
33 vocab_size: int,
34 - extra_vocab_size: int,
34 **kwargs,
35 ) -> None:
36 """
@@ -172,8 +171,11 @@ class PunicaWrapperBase(PunicaWrapperABC):
171 lora_index_to_id: list[int | None],
172 max_loras: int,
173 vocab_size: int,
175 - extra_vocab_size: int,
174 ):
175 + # NOTE We have remove lora extra vocab support for now. So we set
176 + # extra_vocab_size alwayzs to 0, and extra_vocab_size will be removed.
177 +
178 + extra_vocab_size = 0
179 (
180 base_indices,
181 sampler_indices,
@@ -285,12 +287,9 @@ class PunicaWrapperBase(PunicaWrapperABC):
287 lora_index_to_id: list[int | None],
288 max_loras: int,
289 vocab_size: int,
288 - extra_vocab_size: int,
290 **kwargs,
291 ):
291 - self._update_base_metadata(
292 - mapping, lora_index_to_id, max_loras, vocab_size, extra_vocab_size
293 - )
292 + self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
293
294 if mapping.is_prefill:
295 # Update metadata required for prefill-related operators.
vllm/lora/punica_wrapper/punica_gpu.py
+1 -4
@@ -65,13 +65,10 @@ class PunicaWrapperGPU(PunicaWrapperBase):
65 lora_index_to_id: list[int | None],
66 max_loras: int,
67 vocab_size: int,
68 - extra_vocab_size: int,
68 **kwargs,
69 ):
70 self.is_prefill = mapping.is_prefill
72 - self._update_base_metadata(
73 - mapping, lora_index_to_id, max_loras, vocab_size, extra_vocab_size
74 - )
71 + self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
72
73 # Prepare cuda kernel metadata tensors
74 self.token_mapping_meta.prepare_tensors(self.token_lora_indices)
vllm/lora/punica_wrapper/punica_tpu.py
+1 -2
@@ -292,7 +292,6 @@ class PunicaWrapperTPU(PunicaWrapperBase):
292 lora_index_to_id: list[int | None],
293 max_loras: int,
294 vocab_size: int,
295 - extra_vocab_size: int,
295 ):
296 # Make sure we don't accidentally collect outside operations
297 torch_xla.sync()
@@ -313,7 +312,7 @@ class PunicaWrapperTPU(PunicaWrapperBase):
312 lora_index_to_id,
313 max_loras,
314 vocab_size,
316 - extra_vocab_size,
315 + 0, # extra_vocab_size
316 "cpu",
317 )
318 self._token_lora_indices = self._pad_to_shape(
vllm/lora/punica_wrapper/punica_xpu.py
+1 -4
@@ -43,13 +43,10 @@ class PunicaWrapperXPU(PunicaWrapperBase):
43 lora_index_to_id: list[int | None],
44 max_loras: int,
45 vocab_size: int,
46 - extra_vocab_size: int,
46 **kwargs,
47 ):
48 self.is_prefill = mapping.is_prefill
50 - self._update_base_metadata(
51 - mapping, lora_index_to_id, max_loras, vocab_size, extra_vocab_size
52 - )
49 + self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
50
51 def _get_token_lora_indices(self, x: torch.Tensor) -> torch.IntTensor:
52 return torch.narrow(self._token_lora_indices, 0, 0, x.size(0))
vllm/lora/utils.py
+10
@@ -166,6 +166,16 @@ def parse_fine_tuned_lora_name(
166 raise ValueError(f"{name} is unsupported LoRA weight")
167
168
169 +def is_base_embeddding_weights(name: str) -> bool:
170 + # hardcoded subfixes for input & output embedding weights
171 + input_embedding_subfix = ".embed_tokens.base_layer.weight"
172 + output_embedding_subfix = ".lm_head.base_layer.weight"
173 +
174 + return name.endswith(input_embedding_subfix) or name.endswith(
175 + output_embedding_subfix
176 + )
177 +
178 +
179 def is_regex_target_modules(
180 load_modules: str | list[str], expected_lora_modules: list[str]
181 ) -> bool:
vllm/lora/worker_manager.py
+1 -8
@@ -121,8 +121,7 @@ class WorkerLoRAManager:
121 lora_model_id=lora_request.lora_int_id,
122 device="cpu",
123 dtype=self.lora_config.lora_dtype,
124 - target_embedding_padding=self.vocab_size
125 - + self.lora_config.lora_extra_vocab_size,
124 + target_embedding_padding=self.vocab_size,
125 embedding_modules=self.embedding_modules,
126 embedding_padding_modules=self.embedding_padding_modules,
127 tensorizer_config_dict=lora_request.tensorizer_config_dict,
@@ -143,12 +142,6 @@ class WorkerLoRAManager:
142 # For BadRequestError
143 raise e
144
146 - if lora.extra_vocab_size > self.lora_config.lora_extra_vocab_size:
147 - raise ValueError(
148 - f"LoRA added vocab size {lora.extra_vocab_size} "
149 - f"is greater than lora_extra_vocab_size "
150 - f"{self.lora_config.lora_extra_vocab_size}."
151 - )
145 return lora
146
147 def add_dummy_lora(self, lora_request: LoRARequest, rank: int) -> bool:
vllm/model_executor/models/granite.py
+5 -29
@@ -46,7 +46,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
46 from vllm.model_executor.layers.quantization import QuantizationConfig
47 from vllm.model_executor.layers.rotary_embedding import get_rope
48 from vllm.model_executor.layers.vocab_parallel_embedding import (
49 - DEFAULT_VOCAB_PADDING_SIZE,
49 ParallelLMHead,
50 VocabParallelEmbedding,
51 )
@@ -261,29 +260,16 @@ class GraniteModel(nn.Module):
260 config = vllm_config.model_config.hf_config
261 cache_config = vllm_config.cache_config
262 quant_config = vllm_config.quant_config
264 - lora_config = vllm_config.lora_config
263
264 self.config = config
265 self.quant_config = quant_config
268 - lora_vocab = (
269 - (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
270 - if lora_config
271 - else 0
272 - )
273 - self.vocab_size = config.vocab_size + lora_vocab
274 - self.org_vocab_size = config.vocab_size
266 +
267 if get_pp_group().is_first_rank or (
268 config.tie_word_embeddings and get_pp_group().is_last_rank
269 ):
270 self.embed_tokens = VocabParallelEmbedding(
279 - self.vocab_size,
271 + config.vocab_size,
272 config.hidden_size,
281 - org_num_embeddings=config.vocab_size,
282 - padding_size=DEFAULT_VOCAB_PADDING_SIZE
283 - # We need bigger padding if using lora for kernel
284 - # compatibility
285 - if not lora_config
286 - else lora_config.lora_vocab_padding_size,
273 quant_config=quant_config,
274 )
275 else:
@@ -420,28 +406,18 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
406 super().__init__()
407 config = vllm_config.model_config.hf_config
408 quant_config = vllm_config.quant_config
423 - lora_config = vllm_config.lora_config
409
410 self.config = config
426 - self.lora_config = lora_config
411 +
412 self.quant_config = quant_config
413
414 self.model = GraniteModel(
415 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
416 )
417 if get_pp_group().is_last_rank:
433 - self.unpadded_vocab_size = config.vocab_size
434 - if lora_config:
435 - self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
418 self.lm_head = ParallelLMHead(
437 - self.unpadded_vocab_size,
419 + config.vocab_size,
420 config.hidden_size,
439 - org_num_embeddings=config.vocab_size,
440 - padding_size=DEFAULT_VOCAB_PADDING_SIZE
441 - # We need bigger padding if using lora for kernel
442 - # compatibility
443 - if not lora_config
444 - else lora_config.lora_vocab_padding_size,
421 quant_config=quant_config,
422 prefix=maybe_prefix(prefix, "lm_head"),
423 )
@@ -453,7 +429,7 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
429 logit_scale /= config.logits_scaling
430
431 self.logits_processor = LogitsProcessor(
456 - self.unpadded_vocab_size, config.vocab_size, scale=logit_scale
432 + config.vocab_size, scale=logit_scale
433 )
434 else:
435 self.lm_head = PPMissingLayer()
vllm/model_executor/models/llama.py
+5 -25
@@ -47,7 +47,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
47 from vllm.model_executor.layers.quantization import QuantizationConfig
48 from vllm.model_executor.layers.rotary_embedding import get_rope
49 from vllm.model_executor.layers.vocab_parallel_embedding import (
50 - DEFAULT_VOCAB_PADDING_SIZE,
50 ParallelLMHead,
51 VocabParallelEmbedding,
52 )
@@ -368,24 +367,18 @@ class LlamaModel(nn.Module):
367
368 config = vllm_config.model_config.hf_config
369 quant_config = vllm_config.quant_config
371 - lora_config = vllm_config.lora_config
370
371 self.config = config
372 self.quant_config = quant_config
375 - lora_vocab = (
376 - (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
377 - if lora_config
378 - else 0
379 - )
380 - self.vocab_size = config.vocab_size + lora_vocab
381 - self.org_vocab_size = config.vocab_size
373 +
374 + self.vocab_size = config.vocab_size
375 +
376 if get_pp_group().is_first_rank or (
377 config.tie_word_embeddings and get_pp_group().is_last_rank
378 ):
379 self.embed_tokens = VocabParallelEmbedding(
380 self.vocab_size,
381 config.hidden_size,
388 - org_num_embeddings=config.vocab_size,
382 quant_config=quant_config,
383 )
384 else:
@@ -562,9 +555,7 @@ class LlamaForCausalLM(
555 super().__init__()
556 config = vllm_config.model_config.hf_config
557 quant_config = vllm_config.quant_config
565 - lora_config = vllm_config.lora_config
558 self.config = config
567 - self.lora_config = lora_config
559
560 self.model = self._init_model(
561 vllm_config=vllm_config,
@@ -573,20 +564,9 @@ class LlamaForCausalLM(
564 )
565
566 if get_pp_group().is_last_rank:
576 - self.unpadded_vocab_size = config.vocab_size
577 - if lora_config:
578 - self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
567 self.lm_head = ParallelLMHead(
580 - self.unpadded_vocab_size,
568 + config.vocab_size,
569 config.hidden_size,
582 - org_num_embeddings=config.vocab_size,
583 - padding_size=(
584 - DEFAULT_VOCAB_PADDING_SIZE
585 - # We need bigger padding if using lora for kernel
586 - # compatibility
587 - if not lora_config
588 - else lora_config.lora_vocab_padding_size
589 - ),
570 quant_config=quant_config,
571 prefix=maybe_prefix(prefix, "lm_head"),
572 )
@@ -595,7 +575,7 @@ class LlamaForCausalLM(
575
576 logit_scale = getattr(config, "logit_scale", 1.0)
577 self.logits_processor = LogitsProcessor(
598 - self.unpadded_vocab_size, config.vocab_size, logit_scale
578 + config.vocab_size, scale=logit_scale
579 )
580 else:
581 self.lm_head = PPMissingLayer()
vllm/model_executor/models/mixtral.py
+8 -24
@@ -51,7 +51,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
51 from vllm.model_executor.layers.quantization import QuantizationConfig
52 from vllm.model_executor.layers.rotary_embedding import get_rope
53 from vllm.model_executor.layers.vocab_parallel_embedding import (
54 - DEFAULT_VOCAB_PADDING_SIZE,
54 ParallelLMHead,
55 VocabParallelEmbedding,
56 )
@@ -301,23 +300,18 @@ class MixtralModel(nn.Module):
300 config = vllm_config.model_config.hf_config
301 cache_config = vllm_config.cache_config
302 quant_config = vllm_config.quant_config
304 - lora_config = vllm_config.lora_config
303 +
304 parallel_config = vllm_config.parallel_config
305
306 self.config = config
307 self.quant_config = quant_config
309 - lora_vocab = (
310 - (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
311 - if lora_config
312 - else 0
313 - )
314 - self.vocab_size = config.vocab_size + lora_vocab
308 +
309 + self.vocab_size = config.vocab_size
310 self.org_vocab_size = config.vocab_size
311
312 self.embed_tokens = VocabParallelEmbedding(
313 self.vocab_size,
314 config.hidden_size,
320 - org_num_embeddings=config.vocab_size,
315 )
316
317 self.enable_eplb = parallel_config.enable_eplb
@@ -508,34 +502,24 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
502 super().__init__()
503 config = vllm_config.model_config.hf_config
504 quant_config = vllm_config.quant_config
511 - lora_config = vllm_config.lora_config
505 +
506 self.config = config
513 - self.lora_config = lora_config
507 +
508 self.quant_config = quant_config
509
510 self.model = MixtralModel(
511 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
512 )
519 - self.unpadded_vocab_size = config.vocab_size
520 - if lora_config:
521 - self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
513 +
514 self.lm_head = ParallelLMHead(
523 - self.unpadded_vocab_size,
515 + config.vocab_size,
516 config.hidden_size,
525 - org_num_embeddings=config.vocab_size,
526 - padding_size=DEFAULT_VOCAB_PADDING_SIZE
527 - # We need bigger padding if using lora for kernel
528 - # compatibility
529 - if not lora_config
530 - else lora_config.lora_vocab_padding_size,
517 quant_config=quant_config,
518 prefix=maybe_prefix(prefix, "lm_head"),
519 )
520 if self.config.tie_word_embeddings:
521 self.lm_head.weight = self.model.embed_tokens.weight
536 - self.logits_processor = LogitsProcessor(
537 - self.unpadded_vocab_size, config.vocab_size
538 - )
522 + self.logits_processor = LogitsProcessor(config.vocab_size)
523 self.make_empty_intermediate_tensors = (
524 self.model.make_empty_intermediate_tensors
525 )
vllm/model_executor/models/teleflm.py
+1 -1
@@ -74,5 +74,5 @@ class TeleFLMForCausalLM(LlamaForCausalLM):
74 self.output_mult = self.config.output_mult / self.mup_scale_factor
75 logit_scale = self.output_mult
76 self.logits_processor = LogitsProcessor(
77 - self.unpadded_vocab_size, self.config.vocab_size, logit_scale
77 + self.config.vocab_size, scale=logit_scale
78 )
vllm/v1/worker/tpu_model_runner.py
-3
@@ -219,9 +219,6 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
219 self.hidden_size = model_config.get_hidden_size()
220 self.vocab_size = model_config.get_vocab_size()
221
222 - if self.lora_config is not None:
223 - self.vocab_size += self.lora_config.lora_extra_vocab_size
224 -
222 # Multi-modal data support
223 self.mm_registry = MULTIMODAL_REGISTRY
224 self.uses_mrope = model_config.uses_mrope