@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 07ad27121f

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

Update deprecated type hinting in `model_loader` (#18130)

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

Harry Mellor committed May 15, 2025 at 12:00 UTC 07ad27121f7a24ce37f7f49d9d936936844bd058

12 files changed +80 -74

pyproject.toml

+2 -2

@@ -71,15 +71,15 @@ exclude = [

71 "vllm/third_party/**" = ["ALL"]

72 "vllm/version.py" = ["F401"]

73 "vllm/_version.py" = ["ALL"]

74 -# Python 3.8 typing. TODO: Remove these excludes after v1.0.0

74 +# Python 3.8 typing - skip V0 code

75 "vllm/attention/**/*.py" = ["UP006", "UP035"]

76 "vllm/core/**/*.py" = ["UP006", "UP035"]

77 "vllm/engine/**/*.py" = ["UP006", "UP035"]

78 "vllm/executor/**/*.py" = ["UP006", "UP035"]

79 -"vllm/model_executor/model_loader/**/*.py" = ["UP006", "UP035"]

79 "vllm/prompt_adapter/**/*.py" = ["UP006", "UP035"]

80 "vllm/spec_decode/**/*.py" = ["UP006", "UP035"]

81 "vllm/worker/**/*.py" = ["UP006", "UP035"]

82 +# Python 3.8 typing - skip utils for ROCm

83 "vllm/utils.py" = ["UP006", "UP035"]

85 [tool.ruff.lint]

vllm/model_executor/model_loader/bitsandbytes_loader.py

+13 -12

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

6 import itertools

7 import math

8 import os

9 -from typing import Any, Callable, Dict, Generator, List, Optional, Tuple

9 +from collections.abc import Generator

10 +from typing import Any, Callable, Optional

12 import numpy as np

13 import torch

@@ -49,21 +50,21 @@ class BitsAndBytesModelLoader(BaseModelLoader):

50 super().__init__(load_config)

52 # Save the module names without sharding.

52 - self.unsharded_weights_modules: List[str] = []

53 + self.unsharded_weights_modules: list[str] = []

54 # Save the module names that are sharded by column.

54 - self.column_sharded_weights_modules: List[str] = []

55 + self.column_sharded_weights_modules: list[str] = []

56 # Store all module names (from transformers) that support

57 # BNB quantization.

57 - self.target_modules: List[str] = []

58 + self.target_modules: list[str] = []

59 # mapping weight names from transformers to vllm.

60 self.weight_mapper: Callable = lambda name: name

62 def _get_weight_files(

63 self,

64 model_name_or_path: str,

64 - allowed_patterns: List[str],

65 + allowed_patterns: list[str],

66 revision: Optional[str] = None,

66 - ) -> Tuple[str, List[str], str]:

67 + ) -> tuple[str, list[str], str]:

68 """Retrieve weight files. Download the files if necessary.

70 Return the weight files and the file pattern."""

@@ -95,7 +96,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

96 f"No model weights found in: `{model_name_or_path}`")

98 def _prepare_weights(self, model_name_or_path: str,

98 - revision: Optional[str]) -> Tuple[List[str], bool]:

99 + revision: Optional[str]) -> tuple[list[str], bool]:

100 """Prepare weight files for the model."""

101

102 allowed_patterns = ["*.safetensors", "*.bin", "*.pt"]

@@ -155,7 +156,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

156 revision: Optional[str],

157 pre_quant: bool,

158 load_8bit: bool,

158 - ) -> Tuple[Generator[Tuple[str, torch.Tensor], None, None], Dict[str,

159 + ) -> tuple[Generator[tuple[str, torch.Tensor], None, None], dict[str,

160 Any]]:

161 """Get an iterator to the model weights with bitsandbytes quantization,

162 as well as the quantization state dictionary."""

@@ -175,7 +176,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

176 hf_weights_files, use_safetensors = self._prepare_weights(

177 model_name_or_path, revision)

178

178 - quant_state_dict: Dict[str, Any] = {}

179 + quant_state_dict: dict[str, Any] = {}

180

181 if pre_quant:

182 if load_8bit:

@@ -257,7 +258,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

258

259 # Closure to parse quant_state for each prequant weight

260 def _parse_quant_state(param_name: str,

260 - temp_state_dict: Dict) -> QuantState:

261 + temp_state_dict: dict) -> QuantState:

262 quant_state = {}

263 for k in temp_state_dict:

264 if param_name + "." in k:

@@ -415,7 +416,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

416

417 # Modules whose weights might have fused on disk

418 # we need their output_sizes to make shard in flight correctly with TP

418 - self.maybe_fused_weights_modules: Dict[str, List[int]] = {}

419 + self.maybe_fused_weights_modules: dict[str, list[int]] = {}

420 self._get_bnb_target_modules(model)

421 for name, module in model.named_modules():

422 # Some modules like `ReplicatedLinear` should not have their weights

@@ -480,7 +481,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

481 torch.cuda.empty_cache()

482

483 param_dict = dict(model.named_parameters())

483 - stacked_quant_state_dict: Dict[str, Dict[int, Any]] = {}

484 + stacked_quant_state_dict: dict[str, dict[int, Any]] = {}

485 # TODO: Change this lazy import to normal import

486 # after the checks are updated to run on a new version

487 from vllm.model_executor.models.utils import is_pp_missing_parameter

vllm/model_executor/model_loader/default_loader.py

+6 -5

@@ -3,7 +3,8 @@ import dataclasses

3 import glob

4 import os

5 import time

6 -from typing import Generator, Iterable, List, Optional, Tuple, cast

6 +from collections.abc import Generator, Iterable

7 +from typing import Optional, cast

9 import huggingface_hub

10 import torch

@@ -92,7 +93,7 @@ class DefaultModelLoader(BaseModelLoader):

93 revision: Optional[str],

94 fall_back_to_pt: bool,

95 allow_patterns_overrides: Optional[list[str]],

95 - ) -> Tuple[str, List[str], bool]:

96 + ) -> tuple[str, list[str], bool]:

97 """Prepare weights for the model.

99 If the model is not local, it will be downloaded."""

@@ -138,7 +139,7 @@ class DefaultModelLoader(BaseModelLoader):

139 else:

140 hf_folder = model_name_or_path

141

141 - hf_weights_files: List[str] = []

142 + hf_weights_files: list[str] = []

143 for pattern in allow_patterns:

144 hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))

145 if len(hf_weights_files) > 0:

@@ -173,7 +174,7 @@ class DefaultModelLoader(BaseModelLoader):

174

175 def _get_weights_iterator(

176 self, source: "Source"

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

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

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

179 hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(

180 source.model_or_path, source.revision, source.fall_back_to_pt,

@@ -238,7 +239,7 @@ class DefaultModelLoader(BaseModelLoader):

239 self,

240 model_config: ModelConfig,

241 model: nn.Module,

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

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

243 primary_weights = DefaultModelLoader.Source(

244 model_config.model,

245 model_config.revision,

vllm/model_executor/model_loader/gguf_loader.py

+3 -3

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 import os

3 -from typing import Dict, Generator, Tuple

3 +from collections.abc import Generator

5 import gguf

6 import torch

@@ -84,8 +84,8 @@ class GGUFModelLoader(BaseModelLoader):

84 return gguf_to_hf_name_map

86 def _get_weights_iterator(

87 - self, model_name_or_path: str, gguf_to_hf_name_map: Dict[str, str]

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

87 + self, model_name_or_path: str, gguf_to_hf_name_map: dict[str, str]

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

89 return gguf_quant_weights_iterator(model_name_or_path,

90 gguf_to_hf_name_map)

vllm/model_executor/model_loader/neuron.py

+5 -5

@@ -5,7 +5,7 @@ import ast

5 import copy

6 import importlib

7 import os

8 -from typing import Dict, List, Optional, Tuple

8 +from typing import Optional

10 import torch

11 import torch.nn as nn

@@ -33,7 +33,7 @@ TORCH_DTYPE_TO_NEURON_AMP = {

33 }

35 # Models supported by Neuron.

36 -_NEURON_SUPPORTED_MODELS: Dict[str, Tuple[str, str, str]] = {

36 +_NEURON_SUPPORTED_MODELS: dict[str, tuple[str, str, str]] = {

37 "LlamaForCausalLM": ("transformers_neuronx.llama.model",

38 "LlamaForSampling", "LlamaForCausalLM"),

39 "MistralForCausalLM": ("transformers_neuronx.mistral.model",

@@ -146,7 +146,7 @@ class NeuronSpeculationCausalLM(nn.Module):

146 self,

147 logits: torch.Tensor,

148 sampling_metadata: SamplingMetadata,

149 - ) -> Optional[List[SamplerOutput]]:

149 + ) -> Optional[list[SamplerOutput]]:

150 batch_size, num_steps = logits.shape

151 seq_ids = [

152 seq_id for sg in sampling_metadata.seq_groups

@@ -188,7 +188,7 @@ def _get_model_architecture(config: PretrainedConfig) -> str:

188 f"{list(_NEURON_SUPPORTED_MODELS.keys())}")

189

190

191 -def _get_buckets(env: str, default_value: List[int]) -> List[int]:

191 +def _get_buckets(env: str, default_value: list[int]) -> list[int]:

192 env_value = os.getenv(env)

193 if env_value is None:

194 return default_value

@@ -464,7 +464,7 @@ def get_neuron_eagle_speculation_model(model_config: ModelConfig,

464

465 draft_model.eval()

466

467 - token_tree: Dict[int, List[int]] = ast.literal_eval(

467 + token_tree: dict[int, list[int]] = ast.literal_eval(

468 speculation_config.speculative_token_tree)

469

470 speculation_model = EagleSpeculativeDecoder(draft_model.model,

vllm/model_executor/model_loader/neuronx_distributed.py

+3 -3

@@ -9,7 +9,7 @@ import importlib

9 import multiprocessing

10 import os

11 import shutil

12 -from typing import Dict, List, Optional, Tuple

12 +from typing import Optional

14 import torch

15 import torch.nn as nn

@@ -46,7 +46,7 @@ TORCH_DTYPE_TO_NEURON_AMP = {

46 }

48 # Models supported by Neuronx distributed for inference.

49 -_NEURON_SUPPORTED_MODELS: Dict[str, Tuple[str, str]] = {

49 +_NEURON_SUPPORTED_MODELS: dict[str, tuple[str, str]] = {

50 "LlamaForCausalLM":

51 ("neuronx_distributed_inference.models.llama.modeling_llama",

52 "NeuronLlamaForCausalLM"),

@@ -365,7 +365,7 @@ class NeuronSpeculationCausalLM(nn.Module):

365 self,

366 logits: torch.Tensor,

367 sampling_metadata: SamplingMetadata,

368 - ) -> Optional[List[SamplerOutput]]:

368 + ) -> Optional[list[SamplerOutput]]:

369 batch_size, num_steps = logits.shape

370 seq_ids = [

371 seq_id for sg in sampling_metadata.seq_groups

vllm/model_executor/model_loader/runai_streamer_loader.py

+4 -3

@@ -2,7 +2,8 @@

2 # ruff: noqa: SIM117

3 import glob

4 import os

5 -from typing import Generator, List, Optional, Tuple

5 +from collections.abc import Generator

6 +from typing import Optional

8 import torch

9 from torch import nn

@@ -48,7 +49,7 @@ class RunaiModelStreamerLoader(BaseModelLoader):

49 os.environ["RUNAI_STREAMER_S3_ENDPOINT"] = aws_endpoint_url

51 def _prepare_weights(self, model_name_or_path: str,

51 - revision: Optional[str]) -> List[str]:

52 + revision: Optional[str]) -> list[str]:

53 """Prepare weights for the model.

55 If the model is not local, it will be downloaded."""

@@ -87,7 +88,7 @@ class RunaiModelStreamerLoader(BaseModelLoader):

89 def _get_weights_iterator(

90 self, model_or_path: str,

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

91 + revision: str) -> Generator[tuple[str, torch.Tensor], None, None]:

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

93 hf_weights_files = self._prepare_weights(model_or_path, revision)

94 return runai_safetensors_weights_iterator(

vllm/model_executor/model_loader/sharded_state_loader.py

+7 -6

@@ -3,7 +3,8 @@

3 import collections

4 import glob

5 import os

6 -from typing import Any, Dict, Generator, List, Optional, Tuple

6 +from collections.abc import Generator

7 +from typing import Any, Optional

9 import torch

10 from torch import nn

@@ -48,12 +49,12 @@ class ShardedStateLoader(BaseModelLoader):

50 @staticmethod

51 def _filter_subtensors(

51 - tensors: Dict[str, torch.Tensor], ) -> Dict[str, torch.Tensor]:

52 + tensors: dict[str, torch.Tensor], ) -> dict[str, torch.Tensor]:

53 """

54 Filter out all tensors that share the same memory or a subset of the

55 memory of another tensor.

56 """

56 - same_storage_groups: Dict[Any, List[Tuple[str, torch.Tensor]]] = (

57 + same_storage_groups: dict[Any, list[tuple[str, torch.Tensor]]] = (

58 collections.defaultdict(list))

59 for key, tensor in tensors.items():

60 if tensor.numel():

@@ -63,7 +64,7 @@ class ShardedStateLoader(BaseModelLoader):

64 def get_end_ptr(tensor: torch.Tensor) -> int:

65 return tensor.view(-1)[-1].data_ptr() + tensor.element_size()

66 - result: Dict[str, torch.Tensor] = {}

67 + result: dict[str, torch.Tensor] = {}

68 for group in same_storage_groups.values():

69 for k, t in group:

70 a, b = t.data_ptr(), get_end_ptr(t)

@@ -160,7 +161,7 @@ class ShardedStateLoader(BaseModelLoader):

161 return model.eval()

162

163 def iterate_over_files(

163 - self, paths) -> Generator[Tuple[str, torch.Tensor], None, None]:

164 + self, paths) -> Generator[tuple[str, torch.Tensor], None, None]:

165 if self.runai_model_streamer:

166 yield from runai_safetensors_weights_iterator(paths, True)

167 else:

@@ -188,7 +189,7 @@ class ShardedStateLoader(BaseModelLoader):

189 part_idx = 0

190 total_size = 0

191 state_dict = ShardedStateLoader._filter_subtensors(model.state_dict())

191 - state_dict_part: Dict[str, torch.Tensor] = {}

192 + state_dict_part: dict[str, torch.Tensor] = {}

193 for key, tensor in state_dict.items():

194 param_size = tensor.nelement() * tensor.element_size()

195 if max_size is not None and total_size + param_size > max_size:

vllm/model_executor/model_loader/tensorizer.py

+4 -3

@@ -6,9 +6,10 @@ import io

6 import os

7 import re

8 import time

9 +from collections.abc import Generator

10 from dataclasses import dataclass

11 from functools import partial

11 -from typing import BinaryIO, Generator, Optional, Tuple, Type, Union

12 +from typing import BinaryIO, Optional, Union

14 import torch

15 from torch import nn

@@ -67,7 +68,7 @@ class TensorizerConfig:

68 s3_access_key_id: Optional[str] = None

69 s3_secret_access_key: Optional[str] = None

70 s3_endpoint: Optional[str] = None

70 - model_class: Optional[Type[torch.nn.Module]] = None

71 + model_class: Optional[type[torch.nn.Module]] = None

72 hf_config: Optional[PretrainedConfig] = None

73 dtype: Optional[Union[str, torch.dtype]] = None

74 _is_sharded: bool = False

@@ -365,7 +366,7 @@ class TensorizerAgent:

366

367 def tensorizer_weights_iterator(

368 tensorizer_args: "TensorizerArgs"

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

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

370 logger.warning("Deserializing HuggingFace models is not optimized for "

371 "loading on vLLM, as tensorizer is forced to load to CPU. "

372 "Consider deserializing a vLLM model instead for faster "

vllm/model_executor/model_loader/tensorizer_loader.py

+2 -2

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # ruff: noqa: SIM117

3 import copy

4 -from typing import Generator, Tuple

4 +from collections.abc import Generator

6 import torch

7 from torch import nn

@@ -36,7 +36,7 @@ class TensorizerLoader(BaseModelLoader):

36 self.tensorizer_config.verify_with_parallel_config(parallel_config)

38 def _get_weights_iterator(

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

39 + self, ) -> Generator[tuple[str, torch.Tensor], None, None]:

40 tensorizer_args = self.tensorizer_config._construct_tensorizer_args()

41 return tensorizer_weights_iterator(tensorizer_args)

vllm/model_executor/model_loader/utils.py

+7 -7

@@ -5,7 +5,7 @@ import inspect

5 import warnings

6 from contextlib import contextmanager

7 from dataclasses import dataclass, field

8 -from typing import Dict, List, Optional, Tuple, Type

8 +from typing import Optional

10 import torch

11 import transformers

@@ -124,7 +124,7 @@ def device_loading_context(module: torch.nn.Module,

124 yield module

125 return

126

127 - original_device_states: Dict[str, torch.device] = {}

127 + original_device_states: dict[str, torch.device] = {}

128

129 # Store original device states and move parameters to GPU if they're on CPU

130 for name, p in module.named_parameters():

@@ -214,7 +214,7 @@ def resolve_transformers_arch(model_config: ModelConfig,

214

215

216 def get_model_architecture(

217 - model_config: ModelConfig) -> Tuple[Type[nn.Module], str]:

217 + model_config: ModelConfig) -> tuple[type[nn.Module], str]:

218 architectures = getattr(model_config.hf_config, "architectures", [])

219

220 # Special handling for quantized Mixtral.

@@ -257,8 +257,8 @@ class ParamMapping:

257 It creates a bidirectional mapping between packed parameters and their

258 constituent parts.

259 """

260 - packed_mapping: Dict[str, List[str]]

261 - inverse_packed_mapping: Dict[str, Tuple[str,

260 + packed_mapping: dict[str, list[str]]

261 + inverse_packed_mapping: dict[str, tuple[str,

262 int]] = field(default_factory=dict)

263

264 def __post_init__(self):

@@ -273,7 +273,7 @@ class ParamMapping:

273 )

274

275 def get_sub_modules(self,

276 - module_name: str) -> Optional[Tuple[str, List[str]]]:

276 + module_name: str) -> Optional[tuple[str, list[str]]]:

277 for key, value in self.packed_mapping.items():

278 if module_name.endswith(key):

279 return key, value

@@ -281,7 +281,7 @@ class ParamMapping:

281

282

283 def configure_quant_config(quant_config: QuantizationConfig,

284 - model_class: Type[nn.Module]):

284 + model_class: type[nn.Module]):

285 """

286 Pass packed_modules_mapping by reference to quant_config so that

287 quant_config can properly match fused modules

vllm/model_executor/model_loader/weight_utils.py

+24 -23

@@ -8,8 +8,9 @@ import os

8 import tempfile

9 import time

10 from collections import defaultdict

11 +from collections.abc import Generator

12 from pathlib import Path

12 -from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union

13 +from typing import Any, Callable, Optional, Union

15 import filelock

16 import gguf

@@ -221,7 +222,7 @@ def get_sparse_attention_config(

222 model_config: ModelConfig,

223 load_config: LoadConfig,

224 sparse_attention_config_filename: str = "sparse_attention_config.json",

224 -) -> Dict[str, Any]:

225 +) -> dict[str, Any]:

226 model_name_or_path = model_config.model

227 is_local = os.path.isdir(model_name_or_path)

228 if not is_local:

@@ -253,9 +254,9 @@ def get_sparse_attention_config(

254 def download_weights_from_hf(

255 model_name_or_path: str,

256 cache_dir: Optional[str],

256 - allow_patterns: List[str],

257 + allow_patterns: list[str],

258 revision: Optional[str] = None,

258 - ignore_patterns: Optional[Union[str, List[str]]] = None,

259 + ignore_patterns: Optional[Union[str, list[str]]] = None,

260 ) -> str:

261 """Download model weights from Hugging Face Hub.

262

@@ -263,11 +264,11 @@ def download_weights_from_hf(

264 model_name_or_path (str): The model name or path.

265 cache_dir (Optional[str]): The cache directory to store the model

266 weights. If None, will use HF defaults.

266 - allow_patterns (List[str]): The allowed patterns for the

267 + allow_patterns (list[str]): The allowed patterns for the

268 weight files. Files matched by any of the patterns will be

269 downloaded.

270 revision (Optional[str]): The revision of the model.

270 - ignore_patterns (Optional[Union[str, List[str]]]): The patterns to

271 + ignore_patterns (Optional[Union[str, list[str]]]): The patterns to

272 filter out the weight files. Files matched by any of the patterns

273 will be ignored.

274

@@ -347,9 +348,9 @@ def download_safetensors_index_file_from_hf(

348 # Passing both of these to the weight loader functionality breaks.

349 # So, we use the index_file to

350 # look up which safetensors files should be used.

350 -def filter_duplicate_safetensors_files(hf_weights_files: List[str],

351 +def filter_duplicate_safetensors_files(hf_weights_files: list[str],

352 hf_folder: str,

352 - index_file: str) -> List[str]:

353 + index_file: str) -> list[str]:

354 # model.safetensors.index.json is a mapping from keys in the

355 # torch state_dict to safetensors file holding that weight.

356 index_file_name = os.path.join(hf_folder, index_file)

@@ -372,7 +373,7 @@ def filter_duplicate_safetensors_files(hf_weights_files: List[str],

373

374

375 def filter_files_not_needed_for_inference(

375 - hf_weights_files: List[str]) -> List[str]:

376 + hf_weights_files: list[str]) -> list[str]:

377 """

378 Exclude files that are not needed for inference.

379

@@ -408,9 +409,9 @@ def np_cache_weights_iterator(

409 model_name_or_path: str,

410 cache_dir: Optional[str],

411 hf_folder: str,

411 - hf_weights_files: List[str],

412 + hf_weights_files: list[str],

413 use_tqdm_on_load: bool,

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

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

415 """Iterate over the weights in the model np files.

416

417 Will dump the model weights to numpy files if they are not already dumped.

@@ -424,7 +425,7 @@ def np_cache_weights_iterator(

425 # dumping the same model weights to numpy at the same time.

426 with get_lock(model_name_or_path, cache_dir):

427 if not os.path.exists(weight_names_file):

427 - weight_names: List[str] = []

428 + weight_names: list[str] = []

429 for bin_file in tqdm(

430 hf_weights_files,

431 desc="Loading np_cache checkpoint shards",

@@ -453,9 +454,9 @@ def np_cache_weights_iterator(

454

455

456 def safetensors_weights_iterator(

456 - hf_weights_files: List[str],

457 + hf_weights_files: list[str],

458 use_tqdm_on_load: bool,

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

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

460 """Iterate over the weights in the model safetensor files."""

461 for st_file in tqdm(

462 hf_weights_files,

@@ -470,9 +471,9 @@ def safetensors_weights_iterator(

471

472

473 def runai_safetensors_weights_iterator(

473 - hf_weights_files: List[str],

474 + hf_weights_files: list[str],

475 use_tqdm_on_load: bool,

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

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

477 """Iterate over the weights in the model safetensor files."""

478 with SafetensorsStreamer() as streamer:

479 for st_file in tqdm(

@@ -486,9 +487,9 @@ def runai_safetensors_weights_iterator(

487

488

489 def fastsafetensors_weights_iterator(

489 - hf_weights_files: List[str],

490 + hf_weights_files: list[str],

491 use_tqdm_on_load: bool,

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

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

493 """Iterate over the weights in the model safetensor files

494 using fastsafetensor library."""

495 if torch.distributed.is_initialized():

@@ -525,10 +526,10 @@ def fastsafetensors_weights_iterator(

526

527

528 def pt_weights_iterator(

528 - hf_weights_files: List[str],

529 + hf_weights_files: list[str],

530 use_tqdm_on_load: bool,

531 pt_load_map_location: Union[str, dict[str, str]] = "cpu",

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

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

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

534 for bin_file in tqdm(

535 hf_weights_files,

@@ -544,7 +545,7 @@ def pt_weights_iterator(

545

546

547 def get_gguf_extra_tensor_names(

547 - gguf_file: str, gguf_to_hf_name_map: Dict[str, str]) -> List[str]:

548 + gguf_file: str, gguf_to_hf_name_map: dict[str, str]) -> list[str]:

549 reader = gguf.GGUFReader(gguf_file)

550 expected_gguf_keys = set(gguf_to_hf_name_map.keys())

551 exact_gguf_keys = set([tensor.name for tensor in reader.tensors])

@@ -553,8 +554,8 @@ def get_gguf_extra_tensor_names(

554

555

556 def gguf_quant_weights_iterator(

556 - gguf_file: str, gguf_to_hf_name_map: Dict[str, str]

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

557 + gguf_file: str, gguf_to_hf_name_map: dict[str, str]

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

559 """

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

561 them to torch tensors