diff --git a/tests/model_executor/model_loader/test_direct_io.py b/tests/model_executor/model_loader/test_direct_io.py index f0bdd85714e7..2abc2b858fc1 100644 --- a/tests/model_executor/model_loader/test_direct_io.py +++ b/tests/model_executor/model_loader/test_direct_io.py @@ -108,6 +108,7 @@ def test_large_cufile_tensor_bounds_compat_memory_and_reads_exact_tail( payload_offset = 8 + int.from_bytes(payload[:8], "little") reads, allocated, closed = [], [], [] limit = 17 + monkeypatch.setattr(direct_io, "_try_device_buffer", lambda _: None) class File: def __init__(self, filename, flags): @@ -339,3 +340,55 @@ def test_destination_requires_current_accelerator(monkeypatch): ) assert direct_io._can_load_into(tensor) is expected assert not direct_io._can_load_into(torch.empty(1)) + + +@pytest.mark.parametrize("budget", ["available", "low", "allocation_race"]) +def test_large_hf_weight_uses_device_buffer_or_bounded_fallback( + monkeypatch, tmp_path, budget +): + value = torch.arange(37, dtype=torch.float32) + path = tmp_path / "hf.safetensors" + save_file({"weight": value}, path) + payload = path.read_bytes() + allocated, attempts = [], [] + size = value.numel() * value.element_size() + + def device_empty(length): + attempts.append(length) + if budget == "allocation_race" and length == size: + raise torch.OutOfMemoryError("competing allocation") + tensor = torch.empty(length, dtype=torch.uint8) + allocated.append(tensor) + return tensor + + @contextmanager + def open_reader(_): + def read(raw, offset, length): + raw.copy_( + torch.frombuffer( + bytearray(payload[offset : offset + length]), dtype=torch.uint8 + ) + ) + + yield read + + monkeypatch.setattr(direct_io, "_empty", device_empty) + free = 0 if budget == "low" else 1024 * 1024 * 1024 + monkeypatch.setattr(torch.accelerator, "get_memory_info", lambda: (free, free)) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + loaded = dict( + direct_io.direct_weights_iterator( + [str(path)], + backend="test", + mode="compat", + open_reader=open_reader, + coalesce_bytes=17, + ) + )["weight"] + torch.testing.assert_close(loaded, value) + if budget == "available": + assert attempts == [size] + assert loaded.data_ptr() == allocated[0].data_ptr() + else: + assert attempts == ([size, 17] if budget == "allocation_race" else [17]) + assert loaded.data_ptr() != allocated[0].data_ptr() diff --git a/vllm/model_executor/model_loader/direct_io.py b/vllm/model_executor/model_loader/direct_io.py index 66f40b9325e6..f25f042710c0 100644 --- a/vllm/model_executor/model_loader/direct_io.py +++ b/vllm/model_executor/model_loader/direct_io.py @@ -120,6 +120,19 @@ def _iter_contiguous_groups( yield group +def _try_device_buffer(length: int) -> torch.Tensor | None: + free, _ = torch.accelerator.get_memory_info() + # Leave room for the model loader's own temporary tensors. The free-memory + # snapshot is advisory: another allocation may win the race below. + reserve = max(256 * 1024 * 1024, length // 4) + if length + reserve > free: + return None + try: + return _empty(length) + except torch.OutOfMemoryError: + return None + + def _can_load_into(tensor: torch.Tensor) -> bool: return ( tensor.device.type == "cuda" @@ -161,8 +174,9 @@ def direct_weights_iterator( leaves its contents unspecified and must abort loading. Other destinations fall back to staging. P2P does not use destination callbacks. - Oversized compatibility-mode tensors are assembled on CPU using bounded - device reads, so TP slicing does not require a full tensor's temporary VRAM. + Oversized compatibility-mode tensors stay on device when free memory + permits. Otherwise bounded device reads assemble them on CPU, so TP slicing + can still proceed without a full tensor's temporary VRAM. """ if coalesce_bytes is not None and coalesce_bytes < 0: raise ValueError("coalesce_bytes must be nonnegative") @@ -204,6 +218,22 @@ def direct_weights_iterator( yield group[0][0], destination continue if mode == "compat" and length > compat_buffer_bytes: + device_raw = _try_device_buffer(length) + if device_raw is not None: + _read_into(read, device_raw, first_offset, compat_buffer_bytes) + for name, dtype, shape, offset, tensor_length in group: + tensor = ( + device_raw.narrow( + 0, offset - first_offset, tensor_length + ) + .view(dtype) + .reshape(shape) + ) + total += tensor_length + yield name, tensor + del tensor + del device_raw + continue logger.info( "%s compat bounded staging tensor=%s bytes=%d " "device_buffer_bytes=%d",