| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING, Any, Literal, TypeAlias |
| 5 | |
| 6 | from pydantic import Field, field_validator |
| 7 | |
| 8 | from vllm.config.utils import config |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.utils.hashing import safe_hash |
| 11 | |
| 12 | DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS = 8 |
| 13 | DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE = 16 * 1024 * 1024 |
| 14 | SafetensorsLoadStrategy: TypeAlias = Literal["lazy", "eager", "prefetch", "torchao"] |
| 15 | |
| 16 | if TYPE_CHECKING: |
| 17 | from vllm.model_executor.model_loader import LoadFormats |
| 18 | from vllm.model_executor.model_loader.tensorizer import TensorizerConfig |
| 19 | else: |
| 20 | LoadFormats = str |
| 21 | TensorizerConfig = Any |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | @config |
| 27 | class LoadConfig: |
| 28 | """Configuration for loading the model weights.""" |
| 29 | |
| 30 | load_format: str | LoadFormats = "auto" |
| 31 | """ |
| 32 | The format of the model weights to load. |
| 33 | |
| 34 | - "auto" will try to load the weights in the safetensors format and fall |
| 35 | back to the pytorch bin format if safetensors format is not available. |
| 36 | - "pt" will load the weights in the pytorch bin format. |
| 37 | - "safetensors" will load the weights in the safetensors format. |
| 38 | - "instanttensor" will load the Safetensors weights on CUDA devices using |
| 39 | InstantTensor, which enables distributed loading with pipelined prefetching |
| 40 | and fast direct I/O. |
| 41 | - "ipc_cache" will map post-quantized weights from a local weight cache |
| 42 | daemon via CUDA IPC for fast engine restarts. Launch the daemon with |
| 43 | `vllm preload` first. |
| 44 | - "npcache" will load the weights in pytorch format and store a numpy cache |
| 45 | to speed up the loading. |
| 46 | - "dummy" will initialize the weights with random values, which is mainly |
| 47 | for profiling. |
| 48 | - "tensorizer" will use CoreWeave's tensorizer library for fast weight |
| 49 | loading. See the Tensorize vLLM Model script in the Examples section for |
| 50 | more information. |
| 51 | - "runai_streamer" will load the Safetensors weights using Run:ai Model |
| 52 | Streamer. |
| 53 | - "runai_streamer_sharded" will load weights from pre-sharded checkpoint |
| 54 | files using Run:ai Model Streamer. |
| 55 | - "sharded_state" will load weights from pre-sharded checkpoint files, |
| 56 | supporting efficient loading of tensor-parallel models. |
| 57 | - "mistral" will load weights from consolidated safetensors files used by |
| 58 | Mistral models. |
| 59 | - "modelexpress" will load weights using ModelExpress. |
| 60 | - Other custom values can be supported via plugins. |
| 61 | """ |
| 62 | download_dir: str | None = None |
| 63 | """Directory to download and load the weights, default to the default |
| 64 | cache directory of Hugging Face.""" |
| 65 | safetensors_load_strategy: SafetensorsLoadStrategy | None = None |
| 66 | """ |
| 67 | Specifies the loading strategy for safetensors weights. |
| 68 | |
| 69 | - None (default): Uses memory-mapped (lazy) loading. When an NFS |
| 70 | filesystem is detected and the total checkpoint size fits within 90%% |
| 71 | of available RAM, prefetching is enabled automatically. |
| 72 | - "lazy": Weights are memory-mapped from the file. This enables |
| 73 | on-demand loading and is highly efficient for models on local storage. |
| 74 | Unlike the default (None), auto-prefetch on NFS is not performed. |
| 75 | - "eager": The entire file is read into CPU memory upfront before loading. |
| 76 | This is recommended for models on network filesystems (e.g., Lustre, NFS) |
| 77 | as it avoids inefficient random reads, significantly speeding up model |
| 78 | initialization. However, it uses more CPU RAM. |
| 79 | - "prefetch": Checkpoint files are read into the OS page cache before |
| 80 | workers load them, speeding up the model loading phase. Useful on |
| 81 | network or high-latency storage. |
| 82 | - "torchao": Weights are loaded in upfront and then reconstructed |
| 83 | into torchao tensor subclasses. This is used when the checkpoint |
| 84 | was quantized using torchao and saved using safetensors. |
| 85 | Needs `torchao >= 0.14.0`. |
| 86 | """ |
| 87 | safetensors_prefetch_num_threads: int = Field( |
| 88 | default=DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, ge=1 |
| 89 | ) |
| 90 | """Number of worker threads used to prefetch safetensors checkpoint files |
| 91 | into the OS page cache when safetensors prefetching is enabled.""" |
| 92 | safetensors_prefetch_block_size: int = Field( |
| 93 | default=DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE, ge=1 |
| 94 | ) |
| 95 | """Read size in bytes for each safetensors checkpoint file prefetch.""" |
| 96 | model_loader_extra_config: dict | TensorizerConfig = Field(default_factory=dict) |
| 97 | """Extra config for model loader. This will be passed to the model loader |
| 98 | corresponding to the chosen load_format.""" |
| 99 | device: str | None = None |
| 100 | """Device to which model weights will be loaded, default to |
| 101 | device_config.device""" |
| 102 | ignore_patterns: list[str] | str = Field(default_factory=lambda: ["original/**/*"]) |
| 103 | """The list of patterns to ignore when loading the model. Default to |
| 104 | "original/**/*" to avoid repeated loading of llama's checkpoints.""" |
| 105 | use_tqdm_on_load: bool = True |
| 106 | """Whether to enable tqdm for showing progress bar when loading model |
| 107 | weights.""" |
| 108 | pt_load_map_location: str | dict[str, str] = "cpu" |
| 109 | """ |
| 110 | The map location for loading pytorch checkpoint, to support loading |
| 111 | checkpoints can only be loaded on certain devices like "cuda", this |
| 112 | is equivalent to `{"": "cuda"}`. Another supported format is mapping |
| 113 | from different devices like from GPU 1 to GPU 0: `{"cuda:1": "cuda:0"}`. |
| 114 | Note that when passed from command line, the strings in dictionary |
| 115 | need to be double quoted for json parsing. For more details, see |
| 116 | the original doc for `map_location` parameter in [`torch.load`][] parameter. |
| 117 | """ |
| 118 | |
| 119 | def compute_hash(self) -> str: |
| 120 | """WARNING: Whenever a new field is added to this config, |
| 121 | ensure that it is included in the factors list if |
| 122 | it affects the computation graph. |
| 123 | |
| 124 | Provide a hash that uniquely identifies all the configs |
| 125 | that affect the structure of the computation |
| 126 | graph from input ids/embeddings to the final hidden states, |
| 127 | excluding anything before input ids/embeddings and after |
| 128 | the final hidden states. |
| 129 | """ |
| 130 | # no factors to consider. |
| 131 | # this config will not affect the computation graph. |
| 132 | factors: list[Any] = [] |
| 133 | hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() |
| 134 | return hash_str |
| 135 | |
| 136 | @field_validator("load_format", mode="after") |
| 137 | def _lowercase_load_format(cls, load_format: str) -> str: |
| 138 | return load_format.lower() |
| 139 | |
| 140 | @field_validator("ignore_patterns", mode="after") |
| 141 | def _validate_ignore_patterns( |
| 142 | cls, ignore_patterns: list[str] | str |
| 143 | ) -> list[str] | str: |
| 144 | if ignore_patterns != ["original/**/*"] and len(ignore_patterns) > 0: |
| 145 | logger.info( |
| 146 | "Ignoring the following patterns when downloading weights: %s", |
| 147 | ignore_patterns, |
| 148 | ) |
| 149 | |
| 150 | return ignore_patterns |