A high-throughput and memory-efficient inference and serving engine for LLMs
[Core] Use fastsafetensors ParallelLoader for weight loading (#40183)
Signed-off-by: Git Bisector <gitbisector@gmail.com> Signed-off-by: gitbisector <gitbisector@gmail.com> Signed-off-by: git bisector <gitbisector@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
gitbisector committed
Jun 15, 2026 at 22:32 UTC
9d808e2309733c4ae9782bd2c237d89e844a273d
3 files changed
+68
-56
tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py
+3
-5
@@ -20,7 +20,9 @@ from vllm.platforms import current_platform
20
not current_platform.is_cuda_alike(),
21
reason="fastsafetensors requires NVIDIA/AMD GPUs",
22
)
23
-def test_fastsafetensors_model_loader():
23
+@pytest.mark.parametrize("queue_size", [0, 1])
24
+def test_fastsafetensors_model_loader(monkeypatch, queue_size):
25
+ monkeypatch.setenv("VLLM_FASTSAFETENSORS_QUEUE_SIZE", str(queue_size))
26
with tempfile.TemporaryDirectory() as tmpdir:
27
huggingface_hub.constants.HF_HUB_OFFLINE = False
28
download_weights_from_hf(
@@ -45,7 +47,3 @@ def test_fastsafetensors_model_loader():
47
assert fastsafetensors_tensor.dtype == hf_safetensors_tensors[name].dtype
48
assert fastsafetensors_tensor.shape == hf_safetensors_tensors[name].shape
49
assert torch.all(fastsafetensors_tensor.eq(hf_safetensors_tensors[name]))
48
-
49
-
50
-if __name__ == "__main__":
51
- test_fastsafetensors_model_loader()
vllm/envs.py
+16
@@ -107,6 +107,7 @@ if TYPE_CHECKING:
107
VLLM_FORCE_AOT_LOAD: bool = False
108
VLLM_USE_MEGA_AOT_ARTIFACT: bool = False
109
VLLM_USE_TRITON_AWQ: bool = False
110
+ VLLM_FASTSAFETENSORS_QUEUE_SIZE: int = 0
111
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
112
VLLM_SKIP_P2P_CHECK: bool = False
113
VLLM_DISABLED_KERNELS: list[str] = []
@@ -1014,6 +1015,21 @@ environment_variables: dict[str, Callable[[], Any]] = {
1015
"VLLM_TEST_FORCE_LOAD_FORMAT": lambda: os.getenv(
1016
"VLLM_TEST_FORCE_LOAD_FORMAT", "dummy"
1017
),
1018
+ # Queue size for fastsafetensors ParallelLoader pipelined weight
1019
+ # loading. Peak load-time VRAM is roughly
1020
+ # model_weights + (1 + queue_size) * shard_size.
1021
+ # Default 0 preserves the non-pipelined memory footprint so this
1022
+ # change does not shrink the loadable-model envelope. Set to 1
1023
+ # (or higher) to overlap producing the next shard's device buffer
1024
+ # with the consumer copying the current shard into model params,
1025
+ # at the cost of `queue_size` extra shard-sized buffers resident
1026
+ # at peak during loading.
1027
+ "VLLM_FASTSAFETENSORS_QUEUE_SIZE": lambda: int(
1028
+ os.getenv("VLLM_FASTSAFETENSORS_QUEUE_SIZE", "0")
1029
+ ),
1030
+ # Time in ms for the zmq client to wait for a response from the backend
1031
+ # server for simple data operations
1032
+ "VLLM_RPC_TIMEOUT": lambda: int(os.getenv("VLLM_RPC_TIMEOUT", "10000")),
1033
# Timeout in seconds for keeping HTTP connections alive in API server
1034
"VLLM_HTTP_TIMEOUT_KEEP_ALIVE": lambda: int(
1035
os.environ.get("VLLM_HTTP_TIMEOUT_KEEP_ALIVE", "5")
vllm/model_executor/model_loader/weight_utils.py
+49
-51
@@ -55,10 +55,9 @@ except ImportError:
55
SafetensorsStreamer = runai_model_streamer.placeholder_attr("SafetensorsStreamer")
56
57
try:
58
- from fastsafetensors import SafeTensorsFileLoader, SingleGroup
58
+ from fastsafetensors import SingleGroup
59
except ImportError:
60
fastsafetensors = PlaceholderModule("fastsafetensors")
61
- SafeTensorsFileLoader = fastsafetensors.placeholder_attr("SafeTensorsFileLoader")
61
SingleGroup = fastsafetensors.placeholder_attr("SingleGroup")
62
63
from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least
@@ -1022,25 +1021,19 @@ def runai_safetensors_weights_iterator(
1021
yield name, tensor.clone()
1022
1023
1025
-def _init_fastsafetensors_loader(
1026
- pg: "torch.distributed.ProcessGroup",
1027
- device: torch.device,
1028
- f_list: list[str],
1029
- *,
1030
- nogds: bool = False,
1031
-):
1032
- loader = SafeTensorsFileLoader(pg, device, nogds=nogds)
1033
- rank_file_map = {i: [f] for i, f in enumerate(f_list)}
1034
- loader.add_filenames(rank_file_map)
1035
- return loader
1036
-
1037
-
1024
def fastsafetensors_weights_iterator(
1025
hf_weights_files: list[str],
1026
use_tqdm_on_load: bool,
1027
) -> Generator[tuple[str, torch.Tensor], None, None]:
1028
"""Iterate over the weights in the model safetensor files
1043
- using fastsafetensor library."""
1029
+ using fastsafetensor library.
1030
+
1031
+ Uses ParallelLoader for pipelined loading: the producer thread
1032
+ prepares metadata for the next shard while the consumer yields
1033
+ tensors from the current shard.
1034
+ """
1035
+ from fastsafetensors.parallel_loader import ParallelLoader
1036
+
1037
if torch.distributed.is_initialized():
1038
pg = torch.distributed.group.WORLD
1039
else:
@@ -1048,48 +1041,53 @@ def fastsafetensors_weights_iterator(
1041
1042
device = torch.device(f"cuda:{current_platform.current_device()}")
1043
hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key)
1051
- weight_files_sub_lists = [
1052
- hf_weights_files[i : i + pg.size()]
1053
- for i in range(0, len(hf_weights_files), pg.size())
1054
- ]
1044
1045
# Use nogds=True for TP > 1 to avoid cuFileDriverOpen() which
1046
# initializes the GDS DMA subsystem for all visible GPUs, creating
1047
# unwanted CUDA contexts on every device.
1048
nogds = pg.size() > 1
1049
1061
- for f_list in tqdm(
1062
- weight_files_sub_lists,
1063
- desc="Loading safetensors using Fastsafetensor loader",
1064
- disable=not enable_tqdm(use_tqdm_on_load),
1065
- bar_format=_BAR_FORMAT,
1066
- ):
1067
- loader = _init_fastsafetensors_loader(pg, device, f_list, nogds=nogds)
1068
- try:
1069
- try:
1070
- fb = loader.copy_files_to_device()
1071
- except RuntimeError as e:
1072
- if "gds" not in str(e):
1073
- raise
1074
-
1075
- loader.close()
1076
- nogds = True
1077
- logger.warning_once(
1078
- "GDS not enabled, setting `nogds=True`.\n"
1079
- "For more information, see: https://github.com/foundation-model-stack/fastsafetensors?tab=readme-ov-file#basic-api-usages"
1080
- )
1081
- loader = _init_fastsafetensors_loader(pg, device, f_list, nogds=nogds)
1082
- fb = loader.copy_files_to_device()
1050
+ queue_size = envs.VLLM_FASTSAFETENSORS_QUEUE_SIZE
1051
+ tqdm_enabled = enable_tqdm(use_tqdm_on_load)
1052
+
1053
+ def _make_loader(nogds: bool) -> "ParallelLoader":
1054
+ return ParallelLoader(
1055
+ pg=pg,
1056
+ hf_weights_files=hf_weights_files,
1057
+ queue_size=queue_size,
1058
+ use_tqdm_on_load=tqdm_enabled,
1059
+ device=str(device),
1060
+ nogds=nogds,
1061
+ )
1062
1084
- try:
1085
- keys = list(fb.key_to_rank_lidx.keys())
1086
- for k in keys:
1087
- t = fb.get_tensor(k)
1088
- yield k, t
1089
- finally:
1090
- fb.close()
1091
- finally:
1092
- loader.close()
1063
+ # GDS can fail either at construction or lazily inside the producer
1064
+ # thread during iteration (e.g. cuFileHandleRegister returning
1065
+ # CU_FILE_HANDLE_NOT_REGISTERED on a filesystem without GDS support).
1066
+ # Catch both and fall back to nogds, but only before yielding any
1067
+ # tensor -- restarting mid-stream would reload earlier shards.
1068
+ pl = None
1069
+ yielded = False
1070
+ try:
1071
+ try:
1072
+ pl = _make_loader(nogds)
1073
+ for name, tensor in pl.iterate_weights():
1074
+ yielded = True
1075
+ yield name, tensor
1076
+ except RuntimeError as e:
1077
+ if nogds or yielded or "gds" not in str(e):
1078
+ raise
1079
+ logger.warning_once(
1080
+ "GDS not enabled, setting `nogds=True`.\n"
1081
+ "For more information, see: https://github.com/foundation-model-stack/"
1082
+ "fastsafetensors?tab=readme-ov-file#basic-api-usages"
1083
+ )
1084
+ if pl is not None:
1085
+ pl.close()
1086
+ pl = _make_loader(nogds=True)
1087
+ yield from pl.iterate_weights()
1088
+ finally:
1089
+ if pl is not None:
1090
+ pl.close()
1091
1092
1093
def instanttensor_weights_iterator(