From 7410ffcd3fe24c716407175025c90a4c6c00b2ba Mon Sep 17 00:00:00 2001 From: Tom Date: Tue, 13 Oct 2026 01:01:17 +0000 Subject: [PATCH] [Loader] Pinned host staging for the Run:ai streamer path The Run:ai streamer reads every tensor into one reusable pageable host buffer, and the weight loaders then copy out of it with param.copy_(), i.e. a driver-staged pageable memcpy per tensor. On an 8xB300 loading a 433 GiB NVFP4 checkpoint (232k tensors) as four TP=2 replicas that copy path ran at about 2.4 GB/s per worker and was the longest span of a tuned cold start. Add a `pinned_staging` key to the loader's extra config. When set, the streamer's staging buffer is page-locked right after stream_files() (pre-faulted with MADV_POPULATE_WRITE from several threads, then cudaHostRegister), the streamer copies each tensor to the current device out of that pinned buffer, and the loaders only ever see device tensors so every TP narrow is a device-to-device copy. The buffer is bounded to 16 GB unless `memory_limit` is set, is unregistered when the iterator is exhausted, and any failure to pin falls back to the stock path with a warning. Nothing changes when the key is unset. Measured on vLLM 0.20.2 (GLM-5.2-NVFP4, runai-model-streamer 0.16.1, 4x TP=2 on 8xB300): "Model loading took" 39-42 s vs 75-79 s (-48 %), four replicas serving at 133 s instead of 170 s from script start, every GPU parameter bit-identical to the stock loader, throughput unchanged. At TP=8 the gain is small (every rank reads the whole checkpoint through the page cache, so the ceiling is a memcpy). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Tom --- .../models/extensions/runai_model_streamer.md | 8 ++ .../test_runai_model_streamer_loader.py | 24 ++++ .../test_weight_utils.py | 25 ++++ .../model_loader/runai_streamer_loader.py | 31 ++++- .../model_loader/weight_utils.py | 131 +++++++++++++++++- 5 files changed, 211 insertions(+), 8 deletions(-) diff --git a/docs/models/extensions/runai_model_streamer.md b/docs/models/extensions/runai_model_streamer.md index 965b2932ffaa..70047e36d082 100644 --- a/docs/models/extensions/runai_model_streamer.md +++ b/docs/models/extensions/runai_model_streamer.md @@ -82,6 +82,14 @@ vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \ --model-loader-extra-config '{"memory_limit":5368709120}' ``` +You can set `pinned_staging` to page-lock the streamer's CPU buffer and have the streamer copy every tensor to the GPU out of it, so the host-to-device copies run as DMA copies from pinned memory instead of pageable copies. This helps on hosts where the weight load is bound by the copy path rather than by storage (for example a large checkpoint on a fast local disk or a parallel filesystem). The buffer is registered once per load, at roughly 0.2 s per GB, so unless `memory_limit` is set it is bounded to 16 GB. If the buffer cannot be pinned, loading falls back to the regular path. + +```bash +vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \ + --load-format runai_streamer \ + --model-loader-extra-config '{"pinned_staging":true}' +``` + !!! note For further instructions about tunable parameters and additional parameters configurable through environment variables, read the [Environment Variables Documentation](https://github.com/run-ai/runai-model-streamer/blob/master/docs/src/env-vars.md). diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py index e0a083b3a8bc..f3057ad83c60 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py @@ -96,6 +96,7 @@ def _runai_loader(extra): ({"distributed": "yes"}, "distributed must be a bool"), ({"concurrency": "16"}, "concurrency must be a positive integer"), ({"concurrency": -1}, "concurrency must be a positive integer"), + ({"pinned_staging": "yes"}, "pinned_staging must be a bool"), ], ) def test_runai_rejects_invalid_extra_config(extra, match): @@ -116,6 +117,29 @@ def test_runai_accepts_valid_extra_config(): assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == "1024" +def test_runai_pinned_staging_bounds_default_memory_limit(): + # Page-locking is paid per byte of staging buffer, so the streamer's 40 GB + # default is bounded unless the user chose a limit themselves. + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("RUNAI_STREAMER_MEMORY_LIMIT", None) + loader = _runai_loader({"pinned_staging": True}) + assert loader._pinned_staging is True + assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == str( + rsl.RUNAI_PINNED_STAGING_DEFAULT_BYTES + ) + + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("RUNAI_STREAMER_MEMORY_LIMIT", None) + _runai_loader({"pinned_staging": True, "memory_limit": 1024}) + assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == "1024" + + with patch.dict(os.environ, {"RUNAI_STREAMER_MEMORY_LIMIT": "2048"}, clear=False): + _runai_loader({"pinned_staging": True}) + assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == "2048" + + assert _runai_loader({"distributed": True})._pinned_staging is False + + def test_runai_invalid_extra_config_leaves_environ_untouched(): # A later invalid key must not leave an earlier valid key applied to # os.environ (all values are validated before any global mutation). diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py b/tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py index 0d2522284845..d4d687e72f33 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py @@ -36,6 +36,31 @@ def test_runai_safetensors_weights_iterator_clones_reused_buffers( assert torch.equal(actual_tensors[name], expected_tensor) +def test_runai_safetensors_weights_iterator_pinned_staging_fallback( + tmp_path, monkeypatch +): + # Without a CUDA device the staging buffer cannot be registered; the + # iterator must then yield the same (copied-out) tensors as the plain path. + monkeypatch.setenv("RUNAI_STREAMER_MEMORY_LIMIT", "0") + weights_file = tmp_path / "model.safetensors" + expected_tensors = { + "first": torch.tensor([1.0, 2.0]), + "second": torch.tensor([3.0, 4.0]), + } + save_file(expected_tensors, weights_file) + + actual_tensors = dict( + runai_safetensors_weights_iterator( + [str(weights_file)], False, pinned_staging=True + ) + ) + + assert actual_tensors.keys() == expected_tensors.keys() + assert actual_tensors["first"].data_ptr() != actual_tensors["second"].data_ptr() + for name, expected_tensor in expected_tensors.items(): + assert torch.equal(actual_tensors[name], expected_tensor) + + def test_runai_model_loader(): with tempfile.TemporaryDirectory() as tmpdir: huggingface_hub.constants.HF_HUB_OFFLINE = False diff --git a/vllm/model_executor/model_loader/runai_streamer_loader.py b/vllm/model_executor/model_loader/runai_streamer_loader.py index e00116006212..4ec6ccb4bf2a 100644 --- a/vllm/model_executor/model_loader/runai_streamer_loader.py +++ b/vllm/model_executor/model_loader/runai_streamer_loader.py @@ -11,6 +11,7 @@ from vllm.config.load import LoadConfig from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.weight_utils import ( + RUNAI_PINNED_STAGING_DEFAULT_BYTES, download_safetensors_index_file_from_hf, download_weights_from_hf, runai_safetensors_weights_iterator, @@ -28,10 +29,16 @@ def __init__(self, load_config: LoadConfig): super().__init__(load_config) self._is_distributed: bool = False + self._pinned_staging: bool = False if load_config.model_loader_extra_config: extra_config = load_config.model_loader_extra_config - allowed_keys = {"distributed", "concurrency", "memory_limit"} + allowed_keys = { + "distributed", + "concurrency", + "memory_limit", + "pinned_staging", + } if unexpected_keys := set(extra_config) - allowed_keys: raise ValueError( "Unexpected extra config keys for runai_streamer: " @@ -44,6 +51,14 @@ def __init__(self, load_config: LoadConfig): raise ValueError(f"distributed must be a bool, got {distributed!r}") self._is_distributed = distributed + if "pinned_staging" in extra_config: + pinned_staging = extra_config["pinned_staging"] + if not isinstance(pinned_staging, bool): + raise ValueError( + f"pinned_staging must be a bool, got {pinned_staging!r}" + ) + self._pinned_staging = pinned_staging + # Validate every value before mutating os.environ, so a later # invalid key cannot leave an earlier one partially applied. env_updates: dict[str, str] = {} @@ -70,6 +85,15 @@ def __init__(self, load_config: LoadConfig): f"memory_limit must be an integer >= -1, got {memory_limit!r}" ) env_updates["RUNAI_STREAMER_MEMORY_LIMIT"] = str(memory_limit) + elif ( + self._pinned_staging + and os.getenv("RUNAI_STREAMER_MEMORY_LIMIT") is None + ): + # Page-locking is paid once per byte of staging buffer, so + # bound the streamer's default (40 GB) unless the user chose. + env_updates["RUNAI_STREAMER_MEMORY_LIMIT"] = str( + RUNAI_PINNED_STAGING_DEFAULT_BYTES + ) os.environ.update(env_updates) runai_streamer_s3_endpoint = os.getenv("RUNAI_STREAMER_S3_ENDPOINT") @@ -127,7 +151,10 @@ def _get_weights_iterator( """Get an iterator for the model weights based on the load format.""" hf_weights_files = self._prepare_weights(model_or_path, revision) return runai_safetensors_weights_iterator( - hf_weights_files, self.load_config.use_tqdm_on_load, self._is_distributed + hf_weights_files, + self.load_config.use_tqdm_on_load, + self._is_distributed, + pinned_staging=self._pinned_staging, ) def download_model(self, model_config: ModelConfig) -> None: diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index ac1da8d6987b..b33430077396 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -1243,18 +1243,121 @@ def _load_file(st_file: str): yield key, state_dict.pop(key) +# Default size of the Run:ai streamer's CPU staging buffer when it is +# page-locked (``pinned_staging``). Registering the buffer costs roughly +# 0.2 s/GB once per load; 16 GB keeps that to a few seconds while leaving the +# streamer enough room to keep every reader busy. +RUNAI_PINNED_STAGING_DEFAULT_BYTES = 16_000_000_000 +# Threads used to pre-fault the staging buffer before it is page-locked. +_RUNAI_PINNED_STAGING_TOUCH_THREADS = 16 + + +def _pin_runai_streamer_buffer(streamer) -> int | None: + """Page-lock the Run:ai streamer's CPU staging buffer. + + The streamer reads every tensor into one reusable host buffer and the + host-to-device copies out of it run as pageable copies, i.e. a driver + staged memcpy per tensor. Registering the buffer with ``cudaHostRegister`` + turns those into DMA copies out of pinned memory. The pages are faulted in + with ``MADV_POPULATE_WRITE`` from several threads first, so the + single-threaded registration does not pay for the page faults itself. + + Returns the registered address, to be passed to ``cudaHostUnregister`` + once the iterator is exhausted, or ``None`` if the buffer could not be + pinned. In that case loading continues unchanged on the pageable path. + """ + try: + buf = streamer.file_streamer.file_streamer.requests_iterator.buffer + except AttributeError as e: + logger.warning( + "Run:ai streamer staging buffer not found (%s); " + "loading from pageable memory.", + e, + ) + return None + nbytes = int(buf.nbytes) + if nbytes == 0: + return None + addr = int(buf.ctypes.data) + start_time = time.perf_counter() + try: + import ctypes + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + page = 4096 + madv_hugepage = 14 # MADV_HUGEPAGE + madv_populate_write = 23 # MADV_POPULATE_WRITE, Linux >= 5.14 + begin = addr & ~(page - 1) + end = (addr + nbytes + page - 1) & ~(page - 1) + libc.madvise( + ctypes.c_void_p(begin), ctypes.c_size_t(end - begin), madv_hugepage + ) + step = (end - begin) // _RUNAI_PINNED_STAGING_TOUCH_THREADS + step = max((step + page - 1) & ~(page - 1), page) + + def _populate(lo: int, hi: int) -> None: + libc.madvise( + ctypes.c_void_p(lo), ctypes.c_size_t(hi - lo), madv_populate_write + ) + + threads = [ + threading.Thread(target=_populate, args=(lo, min(lo + step, end))) + for lo in range(begin, end, step) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + populated_time = time.perf_counter() + + rc = torch.cuda.cudart().cudaHostRegister(addr, nbytes, 0) + if int(rc) != 0: + logger.warning( + "cudaHostRegister on the Run:ai streamer staging buffer returned " + "%s; loading from pageable memory.", + rc, + ) + return None + logger.info( + "Page-locked the %.1f GB Run:ai streamer staging buffer " + "(populate %.2f s, register %.2f s).", + nbytes / 1e9, + populated_time - start_time, + time.perf_counter() - populated_time, + ) + return addr + except Exception as e: + logger.warning( + "Could not page-lock the Run:ai streamer staging buffer (%s); " + "loading from pageable memory.", + e, + ) + return None + + def runai_safetensors_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, is_distributed: bool = False, + pinned_staging: bool = False, ) -> Generator[tuple[str, torch.Tensor], None, None]: - """Iterate over the weights in the model safetensor files.""" + """Iterate over the weights in the model safetensor files. + + Args: + hf_weights_files: safetensors files to stream. + use_tqdm_on_load: show a progress bar. + is_distributed: use the streamer's distributed mode. + pinned_staging: page-lock the streamer's CPU staging buffer and let + the streamer copy each tensor to the current device out of it, so + the weight loaders only ever see device tensors. Falls back to the + pageable path if the buffer cannot be pinned. + + """ with SafetensorsStreamer() as streamer: is_cuda_alike = current_platform.is_cuda_alike() + stream_to_device = is_cuda_alike and (is_distributed or pinned_staging) device = ( - f"cuda:{current_platform.current_device()}" - if is_distributed and is_cuda_alike - else "cpu" + f"cuda:{current_platform.current_device()}" if stream_to_device else "cpu" ) streamer.stream_files( @@ -1262,6 +1365,7 @@ def runai_safetensors_weights_iterator( device=device, is_distributed=is_distributed, ) + pinned_addr = _pin_runai_streamer_buffer(streamer) if pinned_staging else None total_tensors = sum( len(tensors_meta) for tensors_meta in streamer.files_to_tensors_metadata.values() @@ -1276,8 +1380,23 @@ def runai_safetensors_weights_iterator( mininterval=2, ) - for name, tensor in tensor_iter: - yield name, tensor.clone() + try: + for name, tensor in tensor_iter: + # The streamer reuses its staging buffers, so a tensor must be + # copied out before the next chunk lands in the same memory. + # On the pinned path (non-distributed) the streamer has already + # moved the tensor to the device, which is that copy. + if ( + pinned_addr is not None + and not is_distributed + and tensor.device.type != "cpu" + ): + yield name, tensor + else: + yield name, tensor.clone() + finally: + if pinned_addr is not None: + torch.cuda.cudart().cudaHostUnregister(pinned_addr) def fastsafetensors_weights_iterator(