diff --git a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py index a98d882e5d67..0cf222292e48 100644 --- a/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py +++ b/tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py @@ -2,6 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import glob +import inspect +import json +import struct import sys from contextlib import contextmanager from types import SimpleNamespace @@ -11,6 +14,9 @@ from safetensors.torch import save_file import vllm.model_executor.model_loader.weight_utils as weight_utils +from vllm.model_executor.model_loader.reload.layerwise import ( + _own_deferred_accelerator_tensors, +) from vllm.model_executor.model_loader.weight_utils import ( download_weights_from_hf, instanttensor_weights_iterator, @@ -19,6 +25,52 @@ from vllm.platforms import current_platform +def _safetensors_tensor_metadata(filename): + with open(filename, "rb") as checkpoint: + header_size = struct.unpack(" None: self.post_load_called = True +class _SourceLifetimeQuantMethod(QuantizeMethodBase): + """Observes checkpoint-source ownership when online processing starts.""" + + uses_meta_device = True + + def __init__(self, source_refs): + self.source_refs = source_refs + self.sources_alive_at_process = None + + def create_weights(self, layer, *weight_args, **extra_weight_attrs): + pass + + def apply(self, layer, *args, **kwargs): + raise NotImplementedError + + def process_weights_after_loading(self, layer): + gc.collect() + self.sources_alive_at_process = [ + source_ref() is not None for source_ref in self.source_refs + ] + + +class _DeferredOnlineQuantAttention(_ReloadableAttentionLayer): + """Attention layer whose checkpoint tensor is processed after loading.""" + + def __init__(self): + torch.nn.Module.__init__(self) + self.source_refs = [] + self.quant_method = _SourceLifetimeQuantMethod(self.source_refs) + + def tracking_weight_loader(param, loaded_weight): + self.source_refs.append(ref(loaded_weight)) + default_weight_loader(param, loaded_weight) + + weight = torch.nn.Parameter(torch.empty(2, 2, device="meta")) + weight.weight_loader = tracking_weight_loader + self.register_parameter("weight", weight) + self.post_load_called = False + initialize_online_processing(self) + + def test_move_metatensors(): tensor = torch.empty((1, 2, 3)) meta_tensor = to_meta_tensor(tensor) @@ -192,6 +234,24 @@ def test_attention_first_load_processes_weights(default_vllm_config, layer_cls): assert torch.equal(layer.weight, loaded_weight) +def test_attention_first_load_releases_sources_before_online_quantization( + default_vllm_config, +): + default_vllm_config.model_config = ModelConfig() + layer = _DeferredOnlineQuantAttention() + model = torch.nn.Sequential(layer) + + layer.weight.weight_loader(layer.weight, torch.full((2, 2), 7.0)) + + assert layer.source_refs[0]() is not None + finalize_layerwise_reload(model, default_vllm_config.model_config) + + assert layer.quant_method.sources_alive_at_process + assert not any(layer.quant_method.sources_alive_at_process) + assert layer.post_load_called + assert torch.equal(layer.weight, torch.full((2, 2), 7.0)) + + def test_reload_lifecycle(): layer = torch.nn.Linear(2, 3) info = LayerReloadingInfo( @@ -686,6 +746,65 @@ def test_online_processing_waits_for_late_registered_bias(): assert torch.equal(quant_method.bias_at_process, loaded_bias) +def test_initial_online_processing_loads_into_materialized_parameters(): + quant_method = _RecordingQuantMethod() + layer = _LateBiasLayer(quant_method) + loaded_weight = torch.full((4, 2), 2.0) + loaded_bias = torch.full((4,), 3.0) + + layer.weight.weight_loader(layer.weight, loaded_weight) + + assert not layer.weight.is_meta + assert torch.equal(layer.weight, loaded_weight) + assert not get_layerwise_info(layer).loaded_weights + assert quant_method.bias_at_process is None + + layer.bias.weight_loader(layer.bias, loaded_bias) + + assert torch.equal(quant_method.bias_at_process, loaded_bias) + assert not get_layerwise_info(layer).loaded_weights + + +def test_online_processing_finalizes_checkpoint_omitted_padding(): + class RecordingWeightQuantMethod(QuantizeMethodBase): + uses_meta_device = True + + def __init__(self): + self.weight_at_process = None + + def create_weights(self, layer, *args, **kwargs): + raise NotImplementedError + + def apply(self, layer, *args, **kwargs): + raise NotImplementedError + + def process_weights_after_loading(self, layer): + self.weight_at_process = layer.weight.detach().clone() + + def shard_loader(param, loaded_weight, shard_id): + start = shard_id * loaded_weight.shape[0] + param.data[start : start + loaded_weight.shape[0]].copy_(loaded_weight) + + quant_method = RecordingWeightQuantMethod() + layer = torch.nn.Module() + layer.quant_method = quant_method + weight = torch.nn.Parameter(torch.empty(6, 2, device="meta")) + weight.weight_loader = shard_loader + layer.register_parameter("weight", weight) + initialize_online_processing(layer) + layer._vllm_online_processing_unloaded = {"weight": 4} + + first = torch.full((2, 2), 3.0) + second = torch.full((2, 2), 7.0) + layer.weight.weight_loader(layer.weight, first, 0) + assert quant_method.weight_at_process is None + layer.weight.weight_loader(layer.weight, second, 1) + + expected = torch.cat((first, second, torch.zeros(2, 2))) + assert torch.equal(quant_method.weight_at_process, expected) + assert not get_layerwise_info(layer).can_load() + + def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch): layer = _AliasedBufferLayer() model = torch.nn.Sequential(layer) diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index e66c5836ebfd..4d7281e1c8c4 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -90,6 +90,8 @@ def __init__(self, load_config: LoadConfig): "enable_multithread_load", "num_threads", "enable_weights_track", + "instanttensor_priority_weight_name_prefixes", + "instanttensor_small_checkpoint_max_bytes", } unexpected_keys = set(extra_config.keys()) - allowed_keys @@ -113,6 +115,30 @@ def __init__(self, load_config: LoadConfig): raise ValueError( f"num_threads must be a positive integer, got {num_threads!r}" ) + priority_prefixes = extra_config.get( + "instanttensor_priority_weight_name_prefixes" + ) + if priority_prefixes is not None and not ( + isinstance(priority_prefixes, list) + and priority_prefixes + and all(isinstance(prefix, str) and prefix for prefix in priority_prefixes) + ): + raise ValueError( + "instanttensor_priority_weight_name_prefixes must be a " + "non-empty list of non-empty strings" + ) + small_checkpoint_max_bytes = extra_config.get( + "instanttensor_small_checkpoint_max_bytes" + ) + if small_checkpoint_max_bytes is not None and not ( + isinstance(small_checkpoint_max_bytes, int) + and not isinstance(small_checkpoint_max_bytes, bool) + and small_checkpoint_max_bytes > 0 + ): + raise ValueError( + "instanttensor_small_checkpoint_max_bytes must be a positive " + f"integer, got {small_checkpoint_max_bytes!r}" + ) self.enable_weights_track: bool | None = extra_config.get( "enable_weights_track", None @@ -300,6 +326,12 @@ def _get_weights_iterator( self.load_config.use_tqdm_on_load, weight_name_prefixes=source.weight_name_prefixes, indexed_tensor_files=indexed_tensor_files, + priority_weight_name_prefixes=extra_config.get( + "instanttensor_priority_weight_name_prefixes" + ), + small_checkpoint_max_bytes=extra_config.get( + "instanttensor_small_checkpoint_max_bytes" + ), ) else: if extra_config.get("enable_multithread_load"): diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index 5aa510a6242a..24021fbe0261 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -52,6 +52,59 @@ # Global set used to track loading for logging purposes only LOADING_LAYERS: WeakSet[torch.nn.Module] = WeakSet() +_ONLINE_PROCESSING_UNLOADED_ATTR = "_vllm_online_processing_unloaded" + + +def _online_processing_load_numel_total(layer: torch.nn.Module) -> int: + """Count checkpoint elements required before processing a layer. + + A layer can allocate padded tensor tails that have no corresponding + checkpoint payload. The layer declares those omitted element counts in a + ``dict[str, int]`` named ``_vllm_online_processing_unloaded``. + + Args: + layer: Layer whose registered tensors and unloaded-tail annotations + define the expected checkpoint payload. + + Returns: + Number of tensor elements that checkpoint weight loaders must provide. + + Raises: + TypeError: If the unloaded-tail annotation is not a dictionary. + ValueError: If an annotated tensor is absent or its unloaded element + count is outside the tensor's bounds. + """ + total = get_layer_size(layer) + unloaded = getattr(layer, _ONLINE_PROCESSING_UNLOADED_ATTR, {}) + if not isinstance(unloaded, dict): + raise TypeError( + f"{_ONLINE_PROCESSING_UNLOADED_ATTR} must be a dict, " + f"got {type(unloaded).__name__}" + ) + for name, numel in unloaded.items(): + tensor = get_layer_tensors(layer).get(name) + if tensor is None: + raise ValueError(f"Annotated unloaded tensor {name!r} is missing") + if not isinstance(numel, int) or not 0 <= numel <= tensor.numel(): + raise ValueError( + f"Invalid unloaded element count for {name}: {numel} " + f"(tensor numel={tensor.numel()})" + ) + total -= numel + return total + + +def _zero_online_processing_unloaded(layer: torch.nn.Module) -> None: + """Zero checkpoint-omitted parameter tails before online processing. + + Args: + layer: Materialized layer with unloaded-tail annotations. + """ + unloaded = getattr(layer, _ONLINE_PROCESSING_UNLOADED_ATTR, {}) + for name, numel in unloaded.items(): + if numel: + getattr(layer, name).data.view(-1)[-numel:].zero_() + def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo: """ @@ -132,7 +185,7 @@ def initialize_online_processing(layer: torch.nn.Module): # Track loading progress to determine when to process/copy info.load_numel = 0 - info.load_numel_total = get_layer_size(layer) + info.load_numel_total = _online_processing_load_numel_total(layer) _wrap_parameters_weight_loader(layer) @@ -146,6 +199,28 @@ def _wrap_parameters_weight_loader(layer: torch.nn.Module) -> None: tensor.weight_loader = make_online_process_loader(layer, name) +def _own_deferred_accelerator_tensors(bound_args: inspect.BoundArguments) -> None: + """Give deferred weight-loader arguments independent device storage. + + Streaming model loaders may yield views into reusable staging storage. + Layerwise processing replays bound arguments only after every checkpoint + shard for a layer arrives, so accelerator tensors must be cloned before + they enter the deferred queue. This also covers views created by name or + shard mapping because custom attributes on the source tensor are not + guaranteed to propagate to those views. + + Args: + bound_args: Normalized weight-loader arguments to update in place. + """ + for name, value in tuple(bound_args.arguments.items()): + if ( + name != "param" + and isinstance(value, torch.Tensor) + and value.device.type not in ("meta", "cpu") + ): + bound_args.arguments[name] = value.clone() + + def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Callable: """Create a wrapped weight loader that defers processing.""" info = get_layerwise_info(layer) @@ -176,16 +251,29 @@ def online_process_loader(*args, **kwargs): # Re-run on each load: layers may register parameters later (e.g., `bias`). # Wrap late parameters and refresh `load_numel_total` so processing waits # until all parameters are loaded. - info.load_numel_total = get_layer_size(layer) + info.load_numel_total = _online_processing_load_numel_total(layer) _wrap_parameters_weight_loader(layer) # Bind and normalize arguments bound_args = loader_signature.bind(*args, **kwargs) bound_args.apply_defaults() - # Buffer loaded weights, track loading progress - info.loaded_weights.append((param_name, bound_args)) - num_loaded, ret = get_numel_loaded(original_loader, bound_args) + direct_load = info.kernel_tensors is None and not is_deferred_attention_layer( + layer + ) + if direct_load: + # Initial online quantization can write each checkpoint shard into + # its materialized TP-local destination before the iterator + # advances. Reloading and deferred attention retain their queued + # arguments because their processing has different lifetime rules. + materialize_layer(layer, info) + bound_args.arguments["param"] = getattr(layer, param_name) + num_loaded, ret = get_numel_loaded(original_loader, bound_args) + else: + _own_deferred_accelerator_tensors(bound_args) + info.loaded_weights.append((param_name, bound_args)) + num_loaded, ret = get_numel_loaded(original_loader, bound_args) + info.load_numel += num_loaded logger.debug( @@ -200,7 +288,7 @@ def online_process_loader(*args, **kwargs): return ret # Log warnings allocating excessive buffers on device - if has_device_tensors(bound_args): + if not direct_load and has_device_tensors(bound_args): LOADING_LAYERS.add(layer) if len(LOADING_LAYERS) >= 2: names = sorted([layer.__class__.__name__ for layer in LOADING_LAYERS]) @@ -352,10 +440,21 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo): param.weight_loader = _get_original_loader(param) # Load all buffered weights into materialized layer (using original loaders) - for name, args in info.loaded_weights: + loaded_args = None + for name, loaded_args in info.loaded_weights: param = getattr(layer, name) - args.arguments["param"] = param - param.weight_loader(*args.args, **args.kwargs) + loaded_args.arguments["param"] = param + param.weight_loader(*loaded_args.args, **loaded_args.kwargs) + + if info.kernel_tensors is None: + # Initial online quantization no longer needs checkpoint tensors after + # their values have reached the materialized parameters. Release the + # buffered sources before quantization and kernel repacking allocate + # their transient workspaces. + info.loaded_weights.clear() + loaded_args = None + + _zero_online_processing_unloaded(layer) # Process weights (quantization, repacking, etc.) quant_method = getattr(layer, "quant_method", None) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index c63912c576dd..540213edd733 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -14,7 +14,8 @@ import time from collections import defaultdict from collections.abc import Callable, Generator, Iterable, Sequence -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass from pathlib import Path from typing import IO, Any @@ -1192,12 +1193,28 @@ def _make_loader(nogds: bool) -> "ParallelLoader": pl.close() +@dataclass(frozen=True) +class _InstantTensorCpuFallback: + position: int + name: str + filename: str + + +@dataclass(frozen=True) +class _InstantTensorSelection: + gpu_tensor_count: int + selected_tensor_count: int + cpu_fallbacks: tuple[_InstantTensorCpuFallback, ...] + + def instanttensor_weights_iterator( hf_weights_files: list[str], use_tqdm_on_load: bool, weight_name_prefixes: Sequence[str] | None = None, *, indexed_tensor_files: dict[str, str] | None = None, + priority_weight_name_prefixes: Sequence[str] | None = None, + small_checkpoint_max_bytes: int | None = None, ) -> Generator[tuple[str, torch.Tensor], None, None]: """Iterate over weights in model safetensor files with InstantTensor. @@ -1215,6 +1232,113 @@ def instanttensor_weights_iterator( if not current_platform.is_cuda(): raise ValueError("InstantTensor requires NVIDIA GPUs") + if small_checkpoint_max_bytes is not None and ( + isinstance(small_checkpoint_max_bytes, bool) or small_checkpoint_max_bytes <= 0 + ): + raise ValueError( + "InstantTensor small-checkpoint threshold must be a positive integer" + ) + + if priority_weight_name_prefixes: + if not all( + isinstance(prefix, str) and prefix + for prefix in priority_weight_name_prefixes + ): + raise ValueError( + "InstantTensor priority weight prefixes must be non-empty strings" + ) + priority_names_by_file: dict[str, set[str]] = defaultdict(set) + tensor_files: Iterable[tuple[str, str]] + if indexed_tensor_files is not None: + tensor_files = indexed_tensor_files.items() + else: + # The target and every speculative-model source share one loader + # configuration. Header discovery lets an unindexed draft ignore + # target-only priorities without disabling them for the target. + discovered_tensor_files: list[tuple[str, str]] = [] + for filename in hf_weights_files: + absolute_filename = os.path.abspath(filename) + with safe_open(filename, framework="pt", device="cpu") as reader: + discovered_tensor_files.extend( + (name, absolute_filename) for name in reader.offset_keys() + ) + tensor_files = discovered_tensor_files + + for name, filename in tensor_files: + if not name.startswith(tuple(priority_weight_name_prefixes)): + continue + if weight_name_prefixes and not _matches_weight_name_prefixes( + name, weight_name_prefixes + ): + continue + priority_names_by_file[os.path.abspath(filename)].add(name) + if not priority_names_by_file: + logger.info_once( + "InstantTensor priority prefixes %s match no tensors in this " + "checkpoint source; all selected tensors use the normal " + "loading schedule.", + tuple(priority_weight_name_prefixes), + ) + else: + logger.info_once( + "InstantTensor loads %d tensors from %d/%d checkpoint shards " + "through CPU safetensors before GPU staging because they match " + "priority prefixes %s.", + sum(len(names) for names in priority_names_by_file.values()), + len(priority_names_by_file), + len(hf_weights_files), + tuple(priority_weight_name_prefixes), + ) + priority_names: set[str] = set() + for filename in hf_weights_files: + names = priority_names_by_file.get(os.path.abspath(filename)) + if not names: + continue + with safe_open(filename, framework="pt", device="cpu") as reader: + physical_names = list(reader.offset_keys()) + missing_names = names.difference(physical_names) + if missing_names: + raise RuntimeError( + "Safetensors index maps priority tensors to a shard " + "that does not contain them: " + f"{sorted(missing_names)[:3]} in {filename}" + ) + for name in physical_names: + if name not in names: + continue + priority_names.add(name) + yield name, reader.get_tensor(name) + expected_priority_names = set().union(*priority_names_by_file.values()) + if priority_names != expected_priority_names: + missing_names = expected_priority_names.difference(priority_names) + raise RuntimeError( + "InstantTensor priority tensors reference checkpoint shards " + f"outside the selected file set: {sorted(missing_names)[:3]}" + ) + else: + priority_names = set() + + if ( + small_checkpoint_max_bytes is not None + and indexed_tensor_files is None + and not priority_names + ): + checkpoint_size = _get_checkpoints_size_bytes(hf_weights_files) + if checkpoint_size <= small_checkpoint_max_bytes: + logger.info_once( + "Loading the %.2f GiB unindexed checkpoint through CPU " + "safetensors because it does not exceed the configured %.2f GiB " + "InstantTensor small-checkpoint threshold.", + checkpoint_size / 1024**3, + small_checkpoint_max_bytes / 1024**3, + ) + yield from safetensors_weights_iterator( + hf_weights_files, + use_tqdm_on_load, + weight_name_prefixes=weight_name_prefixes, + ) + return + try: world_group = get_world_group() except AssertionError: @@ -1229,7 +1353,24 @@ def instanttensor_weights_iterator( raise ValueError(f"INSTANTTENSOR_COPY must be 0 or 1, got {copy_setting!r}") copy_tensors = copy_setting == "1" - restrict_before_io = indexed_tensor_files is not None or bool(weight_name_prefixes) + configured_buffer_size = os.getenv("INSTANTTENSOR_BUFFER_SIZE") + max_gpu_tensor_size: int | None = None + if configured_buffer_size is not None: + try: + max_gpu_tensor_size = int(configured_buffer_size) + except ValueError as e: + raise ValueError( + "INSTANTTENSOR_BUFFER_SIZE must be an integer number of bytes" + ) from e + if max_gpu_tensor_size <= 0: + raise ValueError("INSTANTTENSOR_BUFFER_SIZE must be greater than zero") + + restrict_before_io = ( + indexed_tensor_files is not None + or bool(weight_name_prefixes) + or max_gpu_tensor_size is not None + or bool(priority_names) + ) open_kwargs: dict[str, Any] = {} if restrict_before_io: open_kwargs["load_now"] = False @@ -1244,41 +1385,121 @@ def instanttensor_weights_iterator( copy=copy_tensors, **open_kwargs, ) + selection: _InstantTensorSelection | None = None if restrict_before_io: - _restrict_instanttensor_to_selected_ranges( + selection = _restrict_instanttensor_to_selected_ranges( instant_open, indexed_tensor_files=indexed_tensor_files, weight_name_prefixes=weight_name_prefixes, + max_tensor_size=max_gpu_tensor_size, + excluded_tensor_names=priority_names, ) + if selection.selected_tensor_count == 0: + return + + cpu_fallbacks = selection.cpu_fallbacks if selection is not None else () + if not cpu_fallbacks: + with instant_open as f: + # Track bytes so the bar reports load throughput (GB/s). + pbar = tqdm( + total=f.total_tensor_size, + desc="Loading safetensors using InstantTensor loader", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + position=tqdm._get_free_pos(), + unit="B", + unit_scale=True, + unit_divisor=1024, + mininterval=1.0, + ) + try: + tensor: torch.Tensor | None = None + for name, tensor in f.tensors(): + pbar.update(tensor.numel() * tensor.element_size()) + if weight_name_prefixes and not _matches_weight_name_prefixes( + name, weight_name_prefixes + ): + continue + if not copy_tensors: + # Parameter loaders consume borrowed tensors synchronously. + # A deferred loader must own the tensor before advancing. + tensor._vllm_instanttensor_borrowed = True + yield name, tensor + # A generator frame retains its final DLPack-backed view. Drop it + # before the InstantTensor context releases the staging ring. + tensor = None + finally: + pbar.close() + return + + assert max_gpu_tensor_size is not None + assert selection is not None + logger.info_once( + "Loading %d tensors larger than the %d-byte InstantTensor buffer " + "through CPU safetensors", + len(cpu_fallbacks), + max_gpu_tensor_size, + ) + fallback_by_position = {fallback.position: fallback for fallback in cpu_fallbacks} + with ExitStack() as stack: + fallback_readers: dict[str, Any] = {} + for cpu_fallback in cpu_fallbacks: + if cpu_fallback.filename not in fallback_readers: + fallback_readers[cpu_fallback.filename] = stack.enter_context( + safe_open(cpu_fallback.filename, framework="pt", device="cpu") + ) + + pbar = None + gpu_tensors: Any = iter(()) + if selection.gpu_tensor_count: + instant_reader = stack.enter_context(instant_open) + pbar = tqdm( + total=instant_reader.total_tensor_size, + desc="Loading safetensors using InstantTensor loader", + disable=not enable_tqdm(use_tqdm_on_load), + bar_format=_BAR_FORMAT, + position=tqdm._get_free_pos(), + unit="B", + unit_scale=True, + unit_divisor=1024, + mininterval=1.0, + ) + gpu_tensors = iter(instant_reader.tensors()) - with instant_open as f: - # Track bytes so the bar reports load throughput (GB/s). - pbar = tqdm( - total=f.total_tensor_size, - desc="Loading safetensors using InstantTensor loader", - disable=not enable_tqdm(use_tqdm_on_load), - bar_format=_BAR_FORMAT, - position=tqdm._get_free_pos(), - unit="B", - unit_scale=True, - unit_divisor=1024, - mininterval=1.0, - ) try: - for name, tensor in f.tensors(): - pbar.update(tensor.numel() * tensor.element_size()) - if weight_name_prefixes and not _matches_weight_name_prefixes( - name, weight_name_prefixes - ): + gpu_tensor: torch.Tensor | None = None + for position in range(selection.selected_tensor_count): + scheduled_fallback = fallback_by_position.get(position) + if scheduled_fallback is not None: + yield ( + scheduled_fallback.name, + fallback_readers[scheduled_fallback.filename].get_tensor( + scheduled_fallback.name + ), + ) continue + try: + name, gpu_tensor = next(gpu_tensors) + except StopIteration as e: + raise RuntimeError( + "InstantTensor produced fewer tensors than the selected " + "checkpoint schedule" + ) from e + if pbar is not None: + pbar.update(gpu_tensor.numel() * gpu_tensor.element_size()) if not copy_tensors: - # Parameter loaders consume borrowed tensors synchronously. - # Loaders retaining a tensor past this yield must materialize - # owned storage before InstantTensor advances its ring buffer. - tensor._vllm_instanttensor_borrowed = True - yield name, tensor + gpu_tensor._vllm_instanttensor_borrowed = True + yield name, gpu_tensor + + if next(gpu_tensors, None) is not None: + raise RuntimeError( + "InstantTensor produced more tensors than the selected " + "checkpoint schedule" + ) + gpu_tensor = None finally: - pbar.close() + if pbar is not None: + pbar.close() def _restrict_instanttensor_to_selected_ranges( @@ -1286,8 +1507,28 @@ def _restrict_instanttensor_to_selected_ranges( *, indexed_tensor_files: dict[str, str] | None, weight_name_prefixes: Sequence[str] | None, -) -> None: - """Replace an unopened InstantTensor layout with selected byte ranges.""" + max_tensor_size: int | None = None, + excluded_tensor_names: set[str] | None = None, +) -> _InstantTensorSelection: + """Restrict an InstantTensor loader to requested checkpoint ranges. + + Args: + instant_open: Unopened InstantTensor safetensors loader whose metadata + can be restricted before I/O starts. + indexed_tensor_files: Mapping from tensor names to their canonical + checkpoint files. ``None`` accepts tensors from every input file. + weight_name_prefixes: Model parameter prefixes selected for loading. + ``None`` or an empty sequence accepts every prefix. + max_tensor_size: Largest tensor payload, in bytes, retained in the GPU + staging path. Larger selected tensors are assigned to CPU loading. + excluded_tensor_names: Tensor names already emitted by an earlier + loading pass. Excluded names are removed before InstantTensor opens + its GPU staging context. + + Returns: + Selected GPU tensor count, total selected tensor count, and CPU + fallbacks with their positions in checkpoint order. + """ required_attrs = ( "filename", "ordered_tensor_metadatas", @@ -1313,6 +1554,9 @@ def _restrict_instanttensor_to_selected_ranges( selected_filenames: list[str] = [] selected_metadata: list[tuple[str, dict[str, Any]]] = [] selected_offsets: list[tuple[int, int]] = [] + cpu_fallbacks: list[_InstantTensorCpuFallback] = [] + selected_position = 0 + excluded_selected_count = 0 metadata_pos = 0 offset_pos = 0 @@ -1334,17 +1578,40 @@ def _restrict_instanttensor_to_selected_ranges( f"for {filename}" ) - keep = [ - ( - indexed_tensor_files is None - or indexed_tensor_files.get(name) == filename_abs + keep = [] + for item_index, name in enumerate(physical_names): + indexed_here = indexed_tensor_files is None or ( + indexed_tensor_files.get(name) == filename_abs ) - and ( - not weight_name_prefixes - or _matches_weight_name_prefixes(name, weight_name_prefixes) + prefix_matches = not weight_name_prefixes or _matches_weight_name_prefixes( + name, weight_name_prefixes ) - for name in physical_names - ] + otherwise_selected = indexed_here and prefix_matches + excluded_here = ( + otherwise_selected + and excluded_tensor_names is not None + and name in excluded_tensor_names + ) + excluded_selected_count += int(excluded_here) + selected_here = otherwise_selected and not excluded_here + item_data_offsets = file_metadata[item_index][1]["data_offsets"] + tensor_size = int(item_data_offsets[1]) - int(item_data_offsets[0]) + use_cpu_fallback = ( + selected_here + and max_tensor_size is not None + and tensor_size > max_tensor_size + ) + keep.append(selected_here and not use_cpu_fallback) + if selected_here: + if use_cpu_fallback: + cpu_fallbacks.append( + _InstantTensorCpuFallback( + position=selected_position, + name=name, + filename=filename, + ) + ) + selected_position += 1 run_start = 0 while run_start < tensor_count: @@ -1373,7 +1640,7 @@ def _restrict_instanttensor_to_selected_ranges( raise RuntimeError( "InstantTensor layout contains unaccounted metadata or offsets" ) - if not selected_metadata: + if not selected_metadata and not cpu_fallbacks and excluded_selected_count == 0: raise RuntimeError("InstantTensor index/prefix selection matched no tensors") selected_names = [name for name, _ in selected_metadata] @@ -1381,19 +1648,36 @@ def _restrict_instanttensor_to_selected_ranges( raise RuntimeError( "InstantTensor index-aware selection still contains duplicate tensor names" ) + fallback_names = [fallback.name for fallback in cpu_fallbacks] + if len(fallback_names) != len(set(fallback_names)): + raise RuntimeError( + "InstantTensor CPU fallback selection contains duplicate tensor names" + ) + overlap = set(selected_names).intersection(fallback_names) + if overlap: + raise RuntimeError( + f"InstantTensor GPU and CPU selections overlap: {sorted(overlap)[:3]}" + ) selected_sizes = [ int(item["data_offsets"][1]) - int(item["data_offsets"][0]) for _, item in selected_metadata ] - instant_open.filename = selected_filenames - instant_open.ordered_tensor_metadatas = selected_metadata - instant_open.tensor_name_to_index = { - name: index for index, name in enumerate(selected_names) - } - instant_open.tensor_offsets = selected_offsets - instant_open.tensor_sizes = selected_sizes - instant_open.total_tensor_size = sum(selected_sizes) - instant_open._determine_buffer_size(None) + if selected_metadata: + instant_open.filename = selected_filenames + instant_open.ordered_tensor_metadatas = selected_metadata + instant_open.tensor_name_to_index = { + name: index for index, name in enumerate(selected_names) + } + instant_open.tensor_offsets = selected_offsets + instant_open.tensor_sizes = selected_sizes + instant_open.total_tensor_size = sum(selected_sizes) + instant_open._determine_buffer_size(None) + + return _InstantTensorSelection( + gpu_tensor_count=len(selected_metadata), + selected_tensor_count=selected_position, + cpu_fallbacks=tuple(cpu_fallbacks), + ) def pt_weights_iterator(