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 7, 19:00 UTC
[Misc][LoRA] Increase max vocab size limit to 258048 in logits processor (#34773)
Signed-off-by: Bhoomit Vasani <vbhoomit@amazon.com>
Bhoomit committed
Feb 25, 2026 at 07:30 UTC
42489e43c2718674828ece00eefc0f11088e801d
3 files changed
+34
-9
tests/lora/conftest.py
+6
-6
@@ -103,14 +103,14 @@ def dummy_model(default_vllm_config) -> nn.Module:
103
("output", ColumnParallelLinear(50, 10)),
104
("outact", nn.Sigmoid()),
105
# Special handling for lm_head & sampler
106
- ("lm_head", ParallelLMHead(512, 10)),
107
- ("logits_processor", LogitsProcessor(512)),
106
+ ("lm_head", ParallelLMHead(32064, 10)),
107
+ ("logits_processor", LogitsProcessor(32064)),
108
]
109
)
110
)
111
model.config = MagicMock()
112
model.embedding_modules = {"lm_head": "lm_head"}
113
- model.unpadded_vocab_size = 32000
113
+ model.unpadded_vocab_size = 32064
114
return model
115
116
@@ -136,8 +136,8 @@ def dummy_model_gate_up(default_vllm_config) -> nn.Module:
136
("gate_up_proj", MergedColumnParallelLinear(50, [5, 5])),
137
("outact", nn.Sigmoid()),
138
# Special handling for lm_head & sampler
139
- ("lm_head", ParallelLMHead(512, 10)),
140
- ("logits_processor", LogitsProcessor(512)),
139
+ ("lm_head", ParallelLMHead(32064, 10)),
140
+ ("logits_processor", LogitsProcessor(32064)),
141
]
142
)
143
)
@@ -149,7 +149,7 @@ def dummy_model_gate_up(default_vllm_config) -> nn.Module:
149
],
150
}
151
model.embedding_modules = {"lm_head": "lm_head"}
152
- model.unpadded_vocab_size = 32000
152
+ model.unpadded_vocab_size = 32064
153
154
return model
155
tests/lora/test_layers.py
+26
-1
@@ -353,7 +353,7 @@ def test_embeddings(
353
@torch.inference_mode()
354
@pytest.mark.parametrize("num_loras", [1, 2, 4])
355
@pytest.mark.parametrize("device", DEVICES)
356
-@pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 256512])
356
+@pytest.mark.parametrize("vocab_size", [64000, 256512, 258048])
357
@pytest.mark.parametrize("stage", STAGES)
358
def test_lm_head_logits_processor(
359
default_vllm_config, dist_init, num_loras, device, vocab_size, stage
@@ -468,6 +468,31 @@ def test_lm_head_logits_processor(
468
torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol)
469
470
471
+@torch.inference_mode()
472
+@pytest.mark.parametrize("vocab_size", [512, 32000, 258049, 300000])
473
+@pytest.mark.parametrize("device", DEVICES)
474
+def test_lm_head_logits_processor_invalid_vocab_size(
475
+ default_vllm_config, dist_init, vocab_size, device
476
+) -> None:
477
+ """Test that LogitsProcessorWithLoRA raises ValueError for invalid vocab sizes."""
478
+ if current_platform.is_cuda_alike():
479
+ torch.cuda.set_device(device)
480
+
481
+ torch.set_default_device(device)
482
+ max_loras = 8
483
+ lora_config = LoRAConfig(
484
+ max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16
485
+ )
486
+
487
+ logits_processor = LogitsProcessor(vocab_size)
488
+ lora_logits_processor = LogitsProcessorWithLoRA(
489
+ logits_processor, 1024, torch.float16, device, None
490
+ )
491
+
492
+ with pytest.raises(ValueError, match="vocab size must be > 32000 and <= 258048"):
493
+ lora_logits_processor.create_lora_weights(max_loras, lora_config)
494
+
495
+
496
@torch.inference_mode()
497
@pytest.mark.parametrize("num_loras", [1, 2, 4])
498
@pytest.mark.parametrize("device", DEVICES)
vllm/lora/layers/logits_processor.py
+2
-2
@@ -88,9 +88,9 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
88
model_config: PretrainedConfig | None = None,
89
) -> None:
90
# TODO: Verify if this condition can be further relaxed
91
- if 32000 < self.base_layer.vocab_size > 257024:
91
+ if self.base_layer.vocab_size <= 32000 or self.base_layer.vocab_size > 258048:
92
raise ValueError(
93
- "When using LoRA, vocab size must be 32000 >= vocab_size <= 257024"
93
+ "When using LoRA, vocab size must be > 32000 and <= 258048"
94
)
95
self.lora_a_stacked = torch.zeros(
96
(