b3601da6e7 · tduncan/Nvidia-Optimization-Stack-VLLM

tools/pre_commit/mypy.py

+1 -1

@@ -29,7 +29,7 @@ SEPARATE_GROUPS = [

29 "tests",

30 # v0 related

31 "vllm/lora",

32 - "vllm/model_executor",

32 + "vllm/model_executor/layers",

33 ]

34

35 # TODO(woosuk): Include the code from Megatron and HuggingFace.

vllm/model_executor/layers/sparse_attn_indexer.py

+3

@@ -96,6 +96,7 @@ def sparse_attn_indexer(

96 topk_indices_buffer[: hidden_states.shape[0]] = -1

97 if has_prefill:

98 prefill_metadata = attn_metadata.prefill

99 + assert prefill_metadata is not None

100

101 # Get the full shared workspace buffers once (will allocate on first use)

102 workspace_manager = current_workspace_manager()

@@ -170,6 +171,8 @@ def sparse_attn_indexer(

171

172 if has_decode:

173 decode_metadata = attn_metadata.decode

174 + assert decode_metadata is not None

175 + # kv_cache shape [

176 # kv_cache size requirement [num_block, block_size, n_head, head_dim],

177 # we only have [num_block, block_size, head_dim],

178 kv_cache = kv_cache.unsqueeze(-2)

vllm/model_executor/model_loader/gguf_loader.py

+7 -4

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import os

4 from collections.abc import Generator

5 +from typing import TYPE_CHECKING, cast

6

7 import gguf

8 import regex as re

@@ -27,6 +28,9 @@ from vllm.model_executor.model_loader.weight_utils import (

28 from vllm.transformers_utils.gguf_utils import detect_gguf_multimodal

29 from vllm.utils.torch_utils import set_default_torch_dtype

30

31 +if TYPE_CHECKING:

32 + from vllm.model_executor.layers.quantization.gguf import GGUFConfig

33 +

34 logger = init_logger(__name__)

35

36

@@ -350,10 +354,9 @@ class GGUFModelLoader(BaseModelLoader):

354 for name, weight_type in weight_type_map.items()

355 if weight_type in ("F32", "F16", "BF16") and name.endswith(".weight")

356 ]

353 - logger.debug(

354 - "GGUF unquantized modules: %s",

355 - unquant_names,

356 - )

357 + logger.debug("GGUF unquantized modules: %s", unquant_names)

358 + if TYPE_CHECKING:

359 + vllm_config.quant_config = cast(GGUFConfig, vllm_config.quant_config)

360 vllm_config.quant_config.unquantized_modules.extend(unquant_names)

361

362 target_device = torch.device(device_config.device)

vllm/model_executor/model_loader/runai_streamer_loader.py

+8 -20

@@ -27,28 +27,16 @@ class RunaiModelStreamerLoader(BaseModelLoader):

27 def __init__(self, load_config: LoadConfig):

28 super().__init__(load_config)

29

30 - self._is_distributed = False

30 + self._is_distributed: bool = False

31 if load_config.model_loader_extra_config:

32 extra_config = load_config.model_loader_extra_config

33

34 - if "distributed" in extra_config and isinstance(

35 - extra_config.get("distributed"), bool

36 - ):

37 - self._is_distributed = extra_config.get("distributed")

38 -

39 - if "concurrency" in extra_config and isinstance(

40 - extra_config.get("concurrency"), int

41 - ):

42 - os.environ["RUNAI_STREAMER_CONCURRENCY"] = str(

43 - extra_config.get("concurrency")

44 - )

45 -

46 - if "memory_limit" in extra_config and isinstance(

47 - extra_config.get("memory_limit"), int

48 - ):

49 - os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] = str(

50 - extra_config.get("memory_limit")

51 - )

34 + if isinstance(distributed := extra_config.get("distributed"), bool):

35 + self._is_distributed = distributed

36 + if isinstance(concurrency := extra_config.get("concurrency"), int):

37 + os.environ["RUNAI_STREAMER_CONCURRENCY"] = str(concurrency)

38 + if isinstance(memory_limit := extra_config.get("memory_limit"), int):

39 + os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] = str(memory_limit)

40

41 runai_streamer_s3_endpoint = os.getenv("RUNAI_STREAMER_S3_ENDPOINT")

42 aws_endpoint_url = os.getenv("AWS_ENDPOINT_URL")

@@ -93,7 +81,7 @@ class RunaiModelStreamerLoader(BaseModelLoader):

81 return hf_weights_files

82

83 def _get_weights_iterator(

96 - self, model_or_path: str, revision: str

84 + self, model_or_path: str, revision: str | None

85 ) -> Generator[tuple[str, torch.Tensor], None, None]:

86 """Get an iterator for the model weights based on the load format."""

87 hf_weights_files = self._prepare_weights(model_or_path, revision)

vllm/model_executor/model_loader/sharded_state_loader.py

+2 -1

@@ -6,6 +6,7 @@ import glob

6 import os

7 import time

8 from collections.abc import Generator

9 +from copy import copy

10 from typing import Any

11

12 import torch

@@ -42,7 +43,7 @@ class ShardedStateLoader(BaseModelLoader):

43 extra_config = (

44 {}

45 if load_config.model_loader_extra_config is None

45 - else load_config.model_loader_extra_config.copy()

46 + else copy(load_config.model_loader_extra_config)

47 )

48 self.pattern = extra_config.pop("pattern", self.DEFAULT_PATTERN)

49 if extra_config:

vllm/model_executor/model_loader/tensorizer.py

+2 -1

@@ -674,7 +674,8 @@ def serialize_vllm_model(

674 key = f.read()

675 encryption_params = EncryptionParams(key=key)

676

677 - output_file = tensorizer_args.tensorizer_uri

677 + if (output_file := tensorizer_args.tensorizer_uri) is None:

678 + raise ValueError("tensorizer_uri must be specified for serialization.")

679 if tensorizer_config._is_sharded:

680 from vllm.distributed import get_tensor_model_parallel_rank

681

vllm/model_executor/model_loader/tensorizer_loader.py

+1

@@ -121,6 +121,7 @@ class TensorizerLoader(BaseModelLoader):

121 if parallel_config.tensor_parallel_size > 1:

122 from vllm.distributed import get_tensor_model_parallel_rank

123

124 + assert self.tensorizer_config.tensorizer_uri is not None

125 self.tensorizer_config.tensorizer_uri = (

126 self.tensorizer_config.tensorizer_uri % get_tensor_model_parallel_rank()

127 )

vllm/model_executor/model_loader/utils.py

+2 -1

@@ -6,6 +6,7 @@ import inspect

6 import warnings

7 from contextlib import contextmanager

8 from dataclasses import dataclass, field

9 +from typing import Any

10

11 import torch

12 from torch import nn

@@ -71,7 +72,7 @@ def initialize_model(

72 model_class,

73 )

74 # try to be compatible with old-style model class

74 - kwargs = {}

75 + kwargs: dict[str, Any] = {}

76 if "prefix" in all_params:

77 kwargs["prefix"] = prefix

78 if "config" in all_params:

vllm/model_executor/model_loader/weight_utils.py

+10 -3

@@ -257,6 +257,8 @@ def convert_bin_to_safetensor_file(

257 def get_quant_config(

258 model_config: ModelConfig, load_config: LoadConfig

259 ) -> QuantizationConfig:

260 + if model_config.quantization is None:

261 + raise ValueError("Model quantization method is not specified in the config.")

262 quant_cls = get_quantization_config(model_config.quantization)

263

264 # GGUF doesn't have config file

@@ -307,6 +309,11 @@ def get_quant_config(

309 # if hf_quant_config is None, we will try to get config from

310 # hf_overrides

311 hf_overrides = model_config.hf_overrides

312 + if not isinstance(hf_overrides, dict):

313 + raise ValueError(

314 + "hf_overrides must be a dict for get_quant_config "

315 + "to get the quantization config from it."

316 + )

317 quantization_config_file = hf_overrides.get("quantization_config_file", None)

318 if quantization_config_file is not None:

319 if hasattr(quant_cls, "from_config_file"):

@@ -1087,7 +1094,7 @@ def multi_thread_pt_weights_iterator(

1094

1095

1096 def get_gguf_extra_tensor_names(

1090 - gguf_file: str, gguf_to_hf_name_map: dict[str, str]

1097 + gguf_file: str | Path, gguf_to_hf_name_map: dict[str, str]

1098 ) -> list[str]:

1099 reader = gguf.GGUFReader(gguf_file)

1100 expected_gguf_keys = set(gguf_to_hf_name_map.keys())

@@ -1097,7 +1104,7 @@ def get_gguf_extra_tensor_names(

1104

1105

1106 def get_gguf_weight_type_map(

1100 - gguf_file: str, gguf_to_hf_name_map: dict[str, str]

1107 + gguf_file: str | Path, gguf_to_hf_name_map: dict[str, str]

1108 ) -> dict[str, str]:

1109 """

1110 Return GGUF mapped weight's name and its quant type

@@ -1111,7 +1118,7 @@ def get_gguf_weight_type_map(

1118

1119

1120 def gguf_quant_weights_iterator(

1114 - gguf_file: str, gguf_to_hf_name_map: dict[str, str]

1121 + gguf_file: str | Path, gguf_to_hf_name_map: dict[str, str]

1122 ) -> Generator[tuple[str, torch.Tensor], None, None]:

1123 """

1124 Iterate over the quant weights in the model gguf files and convert

vllm/model_executor/parameter.py

+8 -8

@@ -154,8 +154,8 @@ class _ColumnvLLMParameter(BasevLLMParameter):

154 self.data.copy_(loaded_weight)

155

156 def load_merged_column_weight(self, loaded_weight: torch.Tensor, **kwargs):

157 - shard_offset = kwargs.get("shard_offset")

158 - shard_size = kwargs.get("shard_size")

157 + shard_offset: int = kwargs["shard_offset"]

158 + shard_size: int = kwargs["shard_size"]

159

160 # TODO: move these to PackedColumnParameter and PackedvLLMParameter

161 if (

@@ -176,10 +176,10 @@ class _ColumnvLLMParameter(BasevLLMParameter):

176 param_data.copy_(loaded_weight)

177

178 def load_qkv_weight(self, loaded_weight: torch.Tensor, **kwargs):

179 - shard_offset = kwargs.get("shard_offset")

180 - shard_size = kwargs.get("shard_size")

181 - shard_id = kwargs.get("shard_id")

182 - num_heads = kwargs.get("num_heads")

179 + shard_offset: int = kwargs["shard_offset"]

180 + shard_size: int = kwargs["shard_size"]

181 + shard_id: str = kwargs["shard_id"]

182 + num_heads: int = kwargs["num_heads"]

183

184 # TODO: move these to PackedColumnParameter and PackedvLLMParameter

185 if (

@@ -191,10 +191,10 @@ class _ColumnvLLMParameter(BasevLLMParameter):

191 )

192

193 param_data = self.data

194 - shard_id = self.tp_rank if shard_id == "q" else self.tp_rank // num_heads

194 + shard_id_int = self.tp_rank if shard_id == "q" else self.tp_rank // num_heads

195 param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)

196 loaded_weight = loaded_weight.narrow(

197 - self.output_dim, shard_id * shard_size, shard_size

197 + self.output_dim, shard_id_int * shard_size, shard_size

198 )

199

200 assert param_data.shape == loaded_weight.shape

[Mypy] Fix mypy for `vllm/model_executor` (except `vllm/model_executor/layers`) (#37904)