@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 109e15a335

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

Add `pt_load_map_location` to allow loading to cuda (#16869)

Signed-off-by: Jerry Zhang <jerryzh168@gmail.com>

Jerry Zhang committed May 1, 2025 at 23:23 UTC 109e15a335a20251cbefa0a81bf51cd7624eae27

6 files changed +74 -3

tests/quantization/test_torchao.py

+26

@@ -3,6 +3,7 @@ import importlib.metadata

3 import importlib.util

5 import pytest

6 +import torch

8 DTYPE = ["bfloat16"]

@@ -21,5 +22,30 @@ def test_pre_quantized_model(vllm_runner):

22 print(output)

25 +@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")

26 +@pytest.mark.parametrize(

27 + "pt_load_map_location",

28 + [

29 + "cuda:0",

30 + # {"": "cuda"},

31 + ])

32 +def test_opt_125m_int4wo_model_loading_with_params(vllm_runner,

33 + pt_load_map_location):

34 + """

35 + Test loading roberta-base model with no lm_head.

36 + """

37 + torch._dynamo.reset()

38 + model_name = "jerryzh168/opt-125m-int4wo"

39 + with vllm_runner(model_name=model_name,

40 + quantization="torchao",

41 + dtype="bfloat16",

42 + pt_load_map_location=pt_load_map_location) as llm:

43 + output = llm.generate_greedy(["The capital of France is"],

44 + max_tokens=32)

45 +

46 + assert output

47 + print(output)

48 +

49 +

50 if __name__ == "__main__":

51 pytest.main([__file__])

tests/test_config.py

+15 -1

@@ -5,7 +5,8 @@ from typing import Literal, Union

6 import pytest

8 -from vllm.config import ModelConfig, PoolerConfig, config, get_field

8 +from vllm.config import (LoadConfig, ModelConfig, PoolerConfig, VllmConfig,

9 + config, get_field)

10 from vllm.model_executor.layers.pooler import PoolingType

11 from vllm.platforms import current_platform

@@ -410,3 +411,16 @@ def test_generation_config_loading():

411 override_generation_config=override_generation_config)

412

413 assert model_config.get_diff_sampling_param() == override_generation_config

414 +

415 +

416 +@pytest.mark.parametrize("pt_load_map_location", [

417 + "cuda",

418 + {

419 + "": "cuda"

420 + },

421 +])

422 +def test_load_config_pt_load_map_location(pt_load_map_location):

423 + load_config = LoadConfig(pt_load_map_location=pt_load_map_location)

424 + config = VllmConfig(load_config=load_config)

425 +

426 + assert config.load_config.pt_load_map_location == pt_load_map_location

vllm/config.py

+10

@@ -1564,6 +1564,16 @@ class LoadConfig:

1564 use_tqdm_on_load: bool = True

1565 """Whether to enable tqdm for showing progress bar when loading model

1566 weights."""

1567 + pt_load_map_location: Union[str, dict[str, str]] = "cpu"

1568 + """

1569 + pt_load_map_location: the map location for loading pytorch checkpoint, to

1570 + support loading checkpoints can only be loaded on certain devices like

1571 + "cuda", this is equivalent to {"": "cuda"}. Another supported format is

1572 + mapping from different devices like from GPU 1 to GPU 0:

1573 + {"cuda:1": "cuda:0"}. Note that when passed from command line, the strings

1574 + in dictionary needs to be double quoted for json parsing. For more details,

1575 + see original doc for `map_location` in https://pytorch.org/docs/stable/generated/torch.load.html

1576 + """

1577

1578 def compute_hash(self) -> str:

1579 """

vllm/engine/arg_utils.py

+17 -1

@@ -64,6 +64,13 @@ def optional_type(

64 return _optional_type

67 +def union_dict_and_str(val: str) -> Optional[Union[str, dict[str, str]]]:

68 + if not re.match("^{.*}$", val):

69 + return str(val)

70 + else:

71 + return optional_type(json.loads)(val)

72 +

73 +

74 @deprecated(

75 "Passing a JSON argument as a string containing comma separated key=value "

76 "pairs is deprecated. This will be removed in v0.10.0. Please use a JSON "

@@ -187,6 +194,10 @@ def get_kwargs(cls: ConfigType) -> dict[str, Any]:

194 kwargs[name]["type"] = human_readable_int

195 elif contains_type(type_hints, float):

196 kwargs[name]["type"] = float

197 + elif contains_type(type_hints,

198 + dict) and (contains_type(type_hints, str) or any(

199 + is_not_builtin(th) for th in type_hints)):

200 + kwargs[name]["type"] = union_dict_and_str

201 elif contains_type(type_hints, dict):

202 # Dict arguments will always be optional

203 kwargs[name]["type"] = optional_type(json.loads)

@@ -371,6 +382,7 @@ class EngineArgs:

382 reasoning_parser: str = DecodingConfig.reasoning_backend

383

384 use_tqdm_on_load: bool = LoadConfig.use_tqdm_on_load

385 + pt_load_map_location: str = LoadConfig.pt_load_map_location

386

387 def __post_init__(self):

388 # support `EngineArgs(compilation_config={...})`

@@ -491,6 +503,8 @@ class EngineArgs:

503 type=str,

504 default=None,

505 help='Name or path of the QLoRA adapter.')

506 + load_group.add_argument('--pt-load-map-location',

507 + **load_kwargs["pt_load_map_location"])

508

509 # Guided decoding arguments

510 guided_decoding_kwargs = get_kwargs(DecodingConfig)

@@ -883,12 +897,14 @@ class EngineArgs:

897

898 if self.quantization == "bitsandbytes":

899 self.load_format = "bitsandbytes"

900 +

901 return LoadConfig(

902 load_format=self.load_format,

903 download_dir=self.download_dir,

904 model_loader_extra_config=self.model_loader_extra_config,

905 ignore_patterns=self.ignore_patterns,

906 use_tqdm_on_load=self.use_tqdm_on_load,

907 + pt_load_map_location=self.pt_load_map_location,

908 )

909

910 def create_speculative_config(

@@ -1513,7 +1529,7 @@ def _warn_or_fallback(feature_name: str) -> bool:

1529 def human_readable_int(value):

1530 """Parse human-readable integers like '1k', '2M', etc.

1531 Including decimal values with decimal multipliers.

1516 -

1532 +

1533 Examples:

1534 - '1k' -> 1,000

1535 - '1K' -> 1,024

vllm/model_executor/model_loader/loader.py

@@ -384,6 +384,7 @@ class DefaultModelLoader(BaseModelLoader):

384 weights_iterator = pt_weights_iterator(

385 hf_weights_files,

386 self.load_config.use_tqdm_on_load,

387 + self.load_config.pt_load_map_location,

388 )

389

390 if current_platform.is_tpu():

@@ -890,6 +891,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

891 iterator = pt_weights_iterator(

892 hf_weights_files,

893 self.load_config.use_tqdm_on_load,

894 + self.load_config.pt_load_map_location,

895 )

896 for org_name, param in iterator:

897 # mapping weight names from transformers to vllm while preserving

vllm/model_executor/model_loader/weight_utils.py

+4 -1

@@ -502,6 +502,7 @@ def fastsafetensors_weights_iterator(

502 def pt_weights_iterator(

503 hf_weights_files: List[str],

504 use_tqdm_on_load: bool,

505 + pt_load_map_location: Union[str, dict[str, str]] = "cpu",

506 ) -> Generator[Tuple[str, torch.Tensor], None, None]:

507 """Iterate over the weights in the model bin/pt files."""

508 for bin_file in tqdm(

@@ -510,7 +511,9 @@ def pt_weights_iterator(

511 disable=not enable_tqdm(use_tqdm_on_load),

512 bar_format=_BAR_FORMAT,

513 ):

513 - state = torch.load(bin_file, map_location="cpu", weights_only=True)

514 + state = torch.load(bin_file,

515 + map_location=pt_load_map_location,

516 + weights_only=True)

517 yield from state.items()

518 del state

519