@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0fa939e2d1

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 15, 13:30 UTC

Improve configs - `LoRAConfig` + `PromptAdapterConfig` (#16980)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Apr 24, 2025 at 18:29 UTC 0fa939e2d154d42e9722e51d0d004eabfc0b5e6f

3 files changed +130 -91

tests/lora/test_lora_manager.py

+22 -8

@@ -31,6 +31,8 @@ DEVICES = ([

31 f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)

32 ] if current_platform.is_cuda_alike() else ["cpu"])

34 +DEFAULT_DTYPE = torch.get_default_dtype()

35 +

37 @pytest.fixture(scope="function", autouse=True)

38 def use_v0_only(monkeypatch: pytest.MonkeyPatch):

@@ -125,8 +127,10 @@ def test_replace_submodules(dist_init, dummy_model):

127 model = dummy_model

128 manager = LoRAModelManager(

129 model, 1, 1, 1,

128 - LoRAConfig(max_lora_rank=8, max_cpu_loras=8, max_loras=8),

129 - torch.device(DEVICES[0]))

130 + LoRAConfig(max_lora_rank=8,

131 + max_cpu_loras=8,

132 + max_loras=8,

133 + lora_dtype=DEFAULT_DTYPE), torch.device(DEVICES[0]))

134 model = manager.model

135 assert isinstance(model.get_submodule("dense1"),

136 ColumnParallelLinearWithLoRA)

@@ -155,7 +159,8 @@ def test_lora_model_manager(dist_init, dummy_model, device):

159 2,

160 LoRAConfig(max_lora_rank=8,

161 max_cpu_loras=3,

158 - max_loras=2),

162 + max_loras=2,

163 + lora_dtype=DEFAULT_DTYPE),

164 device=device)

165 assert all(x is None for x in manager.lora_index_to_id)

166 assert manager.add_adapter(model_lora1)

@@ -221,7 +226,8 @@ def test_lora_lru_cache_model_manager(dist_init, dummy_model, device):

226 2,

227 LoRAConfig(max_lora_rank=8,

228 max_cpu_loras=3,

224 - max_loras=2),

229 + max_loras=2,

230 + lora_dtype=DEFAULT_DTYPE),

231 device=device)

232 assert all(x is None for x in manager.lora_index_to_id)

233 assert manager.add_adapter(model_lora1)

@@ -316,7 +322,8 @@ def test_lru_lora_model_manager(dist_init, dummy_model, device):

322 2,

323 LoRAConfig(max_lora_rank=8,

324 max_cpu_loras=2,

319 - max_loras=2),

325 + max_loras=2,

326 + lora_dtype=DEFAULT_DTYPE),

327 device=device)

328

329 assert all(x is None for x in manager.lora_index_to_id)

@@ -424,7 +431,10 @@ def test_lru_lora_model_manager(dist_init, dummy_model, device):

431 @pytest.mark.parametrize("device", DEVICES)

432 def test_lru_cache_worker_adapter_manager(llama_2_7b_model_extra_embeddings,

433 sql_lora_files, device):

427 - lora_config = LoRAConfig(max_lora_rank=8, max_cpu_loras=4, max_loras=4)

434 + lora_config = LoRAConfig(max_lora_rank=8,

435 + max_cpu_loras=4,

436 + max_loras=4,

437 + lora_dtype=DEFAULT_DTYPE)

438 worker_adapter_manager = LRUCacheWorkerLoRAManager(

439 4, 2, llama_2_7b_model_extra_embeddings.unpadded_vocab_size -

440 lora_config.lora_extra_vocab_size, lora_config, device,

@@ -504,7 +514,10 @@ def test_lru_cache_worker_adapter_manager(llama_2_7b_model_extra_embeddings,

514 def test_worker_adapter_manager(llama_2_7b_model_extra_embeddings,

515 sql_lora_files, device):

516 # Should remove every LoRA not specified in the request.

507 - lora_config = LoRAConfig(max_lora_rank=8, max_cpu_loras=4, max_loras=4)

517 + lora_config = LoRAConfig(max_lora_rank=8,

518 + max_cpu_loras=4,

519 + max_loras=4,

520 + lora_dtype=DEFAULT_DTYPE)

521 worker_adapter_manager = WorkerLoRAManager(

522 4, 2, llama_2_7b_model_extra_embeddings.unpadded_vocab_size -

523 lora_config.lora_extra_vocab_size, lora_config, device,

@@ -600,7 +613,8 @@ def test_packed_loras(dist_init, dummy_model_gate_up, device):

613 2,

614 LoRAConfig(max_lora_rank=8,

615 max_cpu_loras=2,

603 - max_loras=2),

616 + max_loras=2,

617 + lora_dtype=DEFAULT_DTYPE),

618 device=device)

619 model = manager.model

620

vllm/config.py

+38 -8

@@ -2565,18 +2565,41 @@ class SpeculativeConfig:

2565 return f"SpeculativeConfig({method=}, {model=}, {num_spec_tokens=})"

2566

2567

2568 +LoRADType = Literal["auto", "float16", "bfloat16"]

2569 +

2570 +

2571 +@config

2572 @dataclass

2573 class LoRAConfig:

2570 - max_lora_rank: int

2571 - max_loras: int

2574 + """Configuration for LoRA."""

2575 +

2576 + max_lora_rank: int = 16

2577 + """Max LoRA rank."""

2578 + max_loras: int = 1

2579 + """Max number of LoRAs in a single batch."""

2580 fully_sharded_loras: bool = False

2581 + """By default, only half of the LoRA computation is sharded with tensor

2582 + parallelism. Enabling this will use the fully sharded layers. At high

2583 + sequence length, max rank or tensor parallel size, this is likely faster.

2584 + """

2585 max_cpu_loras: Optional[int] = None

2574 - lora_dtype: Optional[Union[torch.dtype, str]] = None

2586 + """Maximum number of LoRAs to store in CPU memory. Must be >= than

2587 + `max_loras`."""

2588 + lora_dtype: Union[torch.dtype, LoRADType] = "auto"

2589 + """Data type for LoRA. If auto, will default to base model dtype."""

2590 lora_extra_vocab_size: int = 256

2591 + """Maximum size of extra vocabulary that can be present in a LoRA adapter

2592 + (added to the base model vocabulary)."""

2593 # This is a constant.

2594 lora_vocab_padding_size: ClassVar[int] = 256

2578 - long_lora_scaling_factors: Optional[tuple[float]] = None

2595 + long_lora_scaling_factors: Optional[tuple[float, ...]] = None

2596 + """Specify multiple scaling factors (which can be different from base model

2597 + scaling factor - see eg. Long LoRA) to allow for multiple LoRA adapters

2598 + trained with those scaling factors to be used at the same time. If not

2599 + specified, only adapters trained with the base model scaling factor are

2600 + allowed."""

2601 bias_enabled: bool = False

2602 + """Enable bias for LoRA adapters."""

2603

2604 def compute_hash(self) -> str:

2605 """

@@ -2641,12 +2664,19 @@ class LoRAConfig:

2664 "V1 LoRA does not support long LoRA, please use V0.")

2665

2666

2667 +@config

2668 @dataclass

2669 class PromptAdapterConfig:

2646 - max_prompt_adapters: int

2647 - max_prompt_adapter_token: int

2670 + max_prompt_adapters: int = 1

2671 + """Max number of PromptAdapters in a batch."""

2672 + max_prompt_adapter_token: int = 0

2673 + """Max number of PromptAdapters tokens."""

2674 max_cpu_prompt_adapters: Optional[int] = None

2649 - prompt_adapter_dtype: Optional[torch.dtype] = None

2675 + """Maximum number of PromptAdapters to store in CPU memory. Must be >= than

2676 + `max_prompt_adapters`."""

2677 + prompt_adapter_dtype: Union[torch.dtype, str] = "auto"

2678 + """Data type for PromptAdapter. If auto, will default to base model dtype.

2679 + """

2680

2681 def compute_hash(self) -> str:

2682 """

@@ -2678,7 +2708,7 @@ class PromptAdapterConfig:

2708 self.max_cpu_prompt_adapters = self.max_prompt_adapters

2709

2710 def verify_with_model_config(self, model_config: ModelConfig):

2681 - if self.prompt_adapter_dtype in (None, "auto"):

2711 + if self.prompt_adapter_dtype == "auto":

2712 self.prompt_adapter_dtype = model_config.dtype

2713 elif isinstance(self.prompt_adapter_dtype, str):

2714 self.prompt_adapter_dtype = getattr(torch,

vllm/engine/arg_utils.py

+70 -75

@@ -7,7 +7,7 @@ import json

7 import re

8 import threading

9 from dataclasses import MISSING, dataclass, fields

10 -from typing import (Any, Callable, Dict, List, Literal, Optional, Tuple, Type,

10 +from typing import (Any, Callable, Dict, List, Literal, Optional, Type,

11 TypeVar, Union, cast, get_args, get_origin)

13 import torch

@@ -192,18 +192,23 @@ class EngineArgs:

192 get_field(MultiModalConfig, "limit_per_prompt")

193 mm_processor_kwargs: Optional[Dict[str, Any]] = None

194 disable_mm_preprocessor_cache: bool = False

195 + # LoRA fields

196 enable_lora: bool = False

196 - enable_lora_bias: bool = False

197 - max_loras: int = 1

198 - max_lora_rank: int = 16

197 + enable_lora_bias: bool = LoRAConfig.bias_enabled

198 + max_loras: int = LoRAConfig.max_loras

199 + max_lora_rank: int = LoRAConfig.max_lora_rank

200 + fully_sharded_loras: bool = LoRAConfig.fully_sharded_loras

201 + max_cpu_loras: Optional[int] = LoRAConfig.max_cpu_loras

202 + lora_dtype: Optional[Union[str, torch.dtype]] = LoRAConfig.lora_dtype

203 + lora_extra_vocab_size: int = LoRAConfig.lora_extra_vocab_size

204 + long_lora_scaling_factors: Optional[tuple[float, ...]] = \

205 + LoRAConfig.long_lora_scaling_factors

206 + # PromptAdapter fields

207 enable_prompt_adapter: bool = False

200 - max_prompt_adapters: int = 1

201 - max_prompt_adapter_token: int = 0

202 - fully_sharded_loras: bool = False

203 - lora_extra_vocab_size: int = 256

204 - long_lora_scaling_factors: Optional[Tuple[float]] = None

205 - lora_dtype: Optional[Union[str, torch.dtype]] = 'auto'

206 - max_cpu_loras: Optional[int] = None

208 + max_prompt_adapters: int = PromptAdapterConfig.max_prompt_adapters

209 + max_prompt_adapter_token: int = \

210 + PromptAdapterConfig.max_prompt_adapter_token

211 +

212 device: Device = DeviceConfig.device

213 num_scheduler_steps: int = SchedulerConfig.num_scheduler_steps

214 multi_step_stream_outputs: bool = SchedulerConfig.multi_step_stream_outputs

@@ -338,10 +343,21 @@ class EngineArgs:

343 kwargs[name]["choices"] = choices

344 choice_type = type(choices[0])

345 assert all(type(c) is choice_type for c in choices), (

341 - f"All choices must be of the same type. "

346 + "All choices must be of the same type. "

347 f"Got {choices} with types {[type(c) for c in choices]}"

348 )

349 kwargs[name]["type"] = choice_type

350 + elif can_be_type(field_type, tuple):

351 + if is_type_in_union(field_type, tuple):

352 + field_type = get_type_from_union(field_type, tuple)

353 + dtypes = get_args(field_type)

354 + dtype = dtypes[0]

355 + assert all(

356 + d is dtype for d in dtypes if d is not Ellipsis

357 + ), ("All non-Ellipsis tuple elements must be of the same "

358 + f"type. Got {dtypes}.")

359 + kwargs[name]["type"] = dtype

360 + kwargs[name]["nargs"] = "+"

361 elif can_be_type(field_type, int):

362 kwargs[name]["type"] = optional_int if optional else int

363 elif can_be_type(field_type, float):

@@ -685,70 +701,49 @@ class EngineArgs:

701 'inputs.')

702

703 # LoRA related configs

688 - parser.add_argument('--enable-lora',

689 - action='store_true',

690 - help='If True, enable handling of LoRA adapters.')

691 - parser.add_argument('--enable-lora-bias',

692 - action='store_true',

693 - help='If True, enable bias for LoRA adapters.')

694 - parser.add_argument('--max-loras',

695 - type=int,

696 - default=EngineArgs.max_loras,

697 - help='Max number of LoRAs in a single batch.')

698 - parser.add_argument('--max-lora-rank',

699 - type=int,

700 - default=EngineArgs.max_lora_rank,

701 - help='Max LoRA rank.')

702 - parser.add_argument(

703 - '--lora-extra-vocab-size',

704 - type=int,

705 - default=EngineArgs.lora_extra_vocab_size,

706 - help=('Maximum size of extra vocabulary that can be '

707 - 'present in a LoRA adapter (added to the base '

708 - 'model vocabulary).'))

709 - parser.add_argument(

704 + lora_kwargs = get_kwargs(LoRAConfig)

705 + lora_group = parser.add_argument_group(

706 + title="LoRAConfig",

707 + description=LoRAConfig.__doc__,

708 + )

709 + lora_group.add_argument(

710 + '--enable-lora',

711 + action=argparse.BooleanOptionalAction,

712 + help='If True, enable handling of LoRA adapters.')

713 + lora_group.add_argument('--enable-lora-bias',

714 + **lora_kwargs["bias_enabled"])

715 + lora_group.add_argument('--max-loras', **lora_kwargs["max_loras"])

716 + lora_group.add_argument('--max-lora-rank',

717 + **lora_kwargs["max_lora_rank"])

718 + lora_group.add_argument('--lora-extra-vocab-size',

719 + **lora_kwargs["lora_extra_vocab_size"])

720 + lora_group.add_argument(

721 '--lora-dtype',

711 - type=str,

712 - default=EngineArgs.lora_dtype,

713 - choices=['auto', 'float16', 'bfloat16'],

714 - help=('Data type for LoRA. If auto, will default to '

715 - 'base model dtype.'))

716 - parser.add_argument(

717 - '--long-lora-scaling-factors',

718 - type=optional_str,

719 - default=EngineArgs.long_lora_scaling_factors,

720 - help=('Specify multiple scaling factors (which can '

721 - 'be different from base model scaling factor '

722 - '- see eg. Long LoRA) to allow for multiple '

723 - 'LoRA adapters trained with those scaling '

724 - 'factors to be used at the same time. If not '

725 - 'specified, only adapters trained with the '

726 - 'base model scaling factor are allowed.'))

727 - parser.add_argument(

728 - '--max-cpu-loras',

729 - type=int,

730 - default=EngineArgs.max_cpu_loras,

731 - help=('Maximum number of LoRAs to store in CPU memory. '

732 - 'Must be >= than max_loras.'))

733 - parser.add_argument(

734 - '--fully-sharded-loras',

735 - action='store_true',

736 - help=('By default, only half of the LoRA computation is '

737 - 'sharded with tensor parallelism. '

738 - 'Enabling this will use the fully sharded layers. '

739 - 'At high sequence length, max rank or '

740 - 'tensor parallel size, this is likely faster.'))

741 - parser.add_argument('--enable-prompt-adapter',

742 - action='store_true',

743 - help='If True, enable handling of PromptAdapters.')

744 - parser.add_argument('--max-prompt-adapters',

745 - type=int,

746 - default=EngineArgs.max_prompt_adapters,

747 - help='Max number of PromptAdapters in a batch.')

748 - parser.add_argument('--max-prompt-adapter-token',

749 - type=int,

750 - default=EngineArgs.max_prompt_adapter_token,

751 - help='Max number of PromptAdapters tokens')

722 + **lora_kwargs["lora_dtype"],

723 + )

724 + lora_group.add_argument('--long-lora-scaling-factors',

725 + **lora_kwargs["long_lora_scaling_factors"])

726 + lora_group.add_argument('--max-cpu-loras',

727 + **lora_kwargs["max_cpu_loras"])

728 + lora_group.add_argument('--fully-sharded-loras',

729 + **lora_kwargs["fully_sharded_loras"])

730 +

731 + # PromptAdapter related configs

732 + prompt_adapter_kwargs = get_kwargs(PromptAdapterConfig)

733 + prompt_adapter_group = parser.add_argument_group(

734 + title="PromptAdapterConfig",

735 + description=PromptAdapterConfig.__doc__,

736 + )

737 + prompt_adapter_group.add_argument(

738 + '--enable-prompt-adapter',

739 + action=argparse.BooleanOptionalAction,

740 + help='If True, enable handling of PromptAdapters.')

741 + prompt_adapter_group.add_argument(

742 + '--max-prompt-adapters',

743 + **prompt_adapter_kwargs["max_prompt_adapters"])

744 + prompt_adapter_group.add_argument(

745 + '--max-prompt-adapter-token',

746 + **prompt_adapter_kwargs["max_prompt_adapter_token"])

747

748 # Device arguments

749 device_kwargs = get_kwargs(DeviceConfig)