diff --git a/tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py b/tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py index 1975eb61b25..da974131f65 100644 --- a/tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py +++ b/tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py @@ -20,7 +20,9 @@ from vllm.platforms import current_platform not current_platform.is_cuda_alike(), reason="fastsafetensors requires NVIDIA/AMD GPUs", ) -def test_fastsafetensors_model_loader(): +@pytest.mark.parametrize("queue_size", [0, 1]) +def test_fastsafetensors_model_loader(monkeypatch, queue_size): + monkeypatch.setenv("VLLM_FASTSAFETENSORS_QUEUE_SIZE", str(queue_size)) with tempfile.TemporaryDirectory() as tmpdir: huggingface_hub.constants.HF_HUB_OFFLINE = False download_weights_from_hf( @@ -45,7 +47,3 @@ def test_fastsafetensors_model_loader(): assert fastsafetensors_tensor.dtype == hf_safetensors_tensors[name].dtype assert fastsafetensors_tensor.shape == hf_safetensors_tensors[name].shape assert torch.all(fastsafetensors_tensor.eq(hf_safetensors_tensors[name])) - - -if __name__ == "__main__": - test_fastsafetensors_model_loader() diff --git a/vllm/envs.py b/vllm/envs.py index a44ca348746..8ea10c3ffae 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -107,6 +107,7 @@ if TYPE_CHECKING: VLLM_FORCE_AOT_LOAD: bool = False VLLM_USE_MEGA_AOT_ARTIFACT: bool = False VLLM_USE_TRITON_AWQ: bool = False + VLLM_FASTSAFETENSORS_QUEUE_SIZE: int = 0 VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False VLLM_SKIP_P2P_CHECK: bool = False VLLM_DISABLED_KERNELS: list[str] = [] @@ -1014,6 +1015,21 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_TEST_FORCE_LOAD_FORMAT": lambda: os.getenv( "VLLM_TEST_FORCE_LOAD_FORMAT", "dummy" ), + # Queue size for fastsafetensors ParallelLoader pipelined weight + # loading. Peak load-time VRAM is roughly + # model_weights + (1 + queue_size) * shard_size. + # Default 0 preserves the non-pipelined memory footprint so this + # change does not shrink the loadable-model envelope. Set to 1 + # (or higher) to overlap producing the next shard's device buffer + # with the consumer copying the current shard into model params, + # at the cost of `queue_size` extra shard-sized buffers resident + # at peak during loading. + "VLLM_FASTSAFETENSORS_QUEUE_SIZE": lambda: int( + os.getenv("VLLM_FASTSAFETENSORS_QUEUE_SIZE", "0") + ), + # Time in ms for the zmq client to wait for a response from the backend + # server for simple data operations + "VLLM_RPC_TIMEOUT": lambda: int(os.getenv("VLLM_RPC_TIMEOUT", "10000")), # Timeout in seconds for keeping HTTP connections alive in API server "VLLM_HTTP_TIMEOUT_KEEP_ALIVE": lambda: int( os.environ.get("VLLM_HTTP_TIMEOUT_KEEP_ALIVE", "5") diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 821c0e99de7..47c6c02be6a 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -55,10 +55,9 @@ except ImportError: SafetensorsStreamer = runai_model_streamer.placeholder_attr("SafetensorsStreamer") try: - from fastsafetensors import SafeTensorsFileLoader, SingleGroup + from fastsafetensors import SingleGroup except ImportError: fastsafetensors = PlaceholderModule("fastsafetensors") - SafeTensorsFileLoader = fastsafetensors.placeholder_attr("SafeTensorsFileLoader") SingleGroup = fastsafetensors.placeholder_attr("SingleGroup") from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least @@ -1022,25 +1021,19 @@ def runai_safetensors_weights_iterator( yield name, tensor.clone() -def _init_fastsafetensors_loader( - pg: "torch.distributed.ProcessGroup", - device: torch.device, - f_list: list[str], - *, - nogds: bool = False, -): - loader = SafeTensorsFileLoader(pg, device, nogds=nogds) - rank_file_map = {i: [f] for i, f in enumerate(f_list)} - loader.add_filenames(rank_file_map) - return loader - - def fastsafetensors_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over the weights in the model safetensor files - using fastsafetensor library.""" + using fastsafetensor library. + + Uses ParallelLoader for pipelined loading: the producer thread + prepares metadata for the next shard while the consumer yields + tensors from the current shard. + """ + from fastsafetensors.parallel_loader import ParallelLoader + if torch.distributed.is_initialized(): pg = torch.distributed.group.WORLD else: @@ -1048,48 +1041,53 @@ def fastsafetensors_weights_iterator( device = torch.device(f"cuda:{current_platform.current_device()}") hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key) - weight_files_sub_lists = [ - hf_weights_files[i : i + pg.size()] - for i in range(0, len(hf_weights_files), pg.size()) - ] # Use nogds=True for TP > 1 to avoid cuFileDriverOpen() which # initializes the GDS DMA subsystem for all visible GPUs, creating # unwanted CUDA contexts on every device. nogds = pg.size() > 1 - for f_list in tqdm( - weight_files_sub_lists, - desc="Loading safetensors using Fastsafetensor loader", - disable=not enable_tqdm(use_tqdm_on_load), - bar_format=_BAR_FORMAT, - ): - loader = _init_fastsafetensors_loader(pg, device, f_list, nogds=nogds) + queue_size = envs.VLLM_FASTSAFETENSORS_QUEUE_SIZE + tqdm_enabled = enable_tqdm(use_tqdm_on_load) + + def _make_loader(nogds: bool) -> "ParallelLoader": + return ParallelLoader( + pg=pg, + hf_weights_files=hf_weights_files, + queue_size=queue_size, + use_tqdm_on_load=tqdm_enabled, + device=str(device), + nogds=nogds, + ) + + # GDS can fail either at construction or lazily inside the producer + # thread during iteration (e.g. cuFileHandleRegister returning + # CU_FILE_HANDLE_NOT_REGISTERED on a filesystem without GDS support). + # Catch both and fall back to nogds, but only before yielding any + # tensor -- restarting mid-stream would reload earlier shards. + pl = None + yielded = False + try: try: - try: - fb = loader.copy_files_to_device() - except RuntimeError as e: - if "gds" not in str(e): - raise - - loader.close() - nogds = True - logger.warning_once( - "GDS not enabled, setting `nogds=True`.\n" - "For more information, see: https://github.com/foundation-model-stack/fastsafetensors?tab=readme-ov-file#basic-api-usages" - ) - loader = _init_fastsafetensors_loader(pg, device, f_list, nogds=nogds) - fb = loader.copy_files_to_device() - - try: - keys = list(fb.key_to_rank_lidx.keys()) - for k in keys: - t = fb.get_tensor(k) - yield k, t - finally: - fb.close() - finally: - loader.close() + pl = _make_loader(nogds) + for name, tensor in pl.iterate_weights(): + yielded = True + yield name, tensor + except RuntimeError as e: + if nogds or yielded or "gds" not in str(e): + raise + logger.warning_once( + "GDS not enabled, setting `nogds=True`.\n" + "For more information, see: https://github.com/foundation-model-stack/" + "fastsafetensors?tab=readme-ov-file#basic-api-usages" + ) + if pl is not None: + pl.close() + pl = _make_loader(nogds=True) + yield from pl.iterate_weights() + finally: + if pl is not None: + pl.close() def instanttensor_weights_iterator(