From 05eec7120ea76383f099b8ad0f731e7da2abc75d Mon Sep 17 00:00:00 2001 From: Benjamin Bartels Date: Thu, 28 May 2026 03:16:52 +0100 Subject: [PATCH] Fix RunAI streamer tensor buffer reuse during weight loading (#43464) Signed-off-by: bbartels --- .../test_weight_utils.py | 22 +++++++++++++++++++ .../model_loader/weight_utils.py | 3 ++- 2 files changed, 24 insertions(+), 1 deletion(-) 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 03691b4a472..0d252228484 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 @@ -6,6 +6,7 @@ import tempfile import huggingface_hub.constants import torch +from safetensors.torch import save_file from vllm.model_executor.model_loader.weight_utils import ( download_weights_from_hf, @@ -14,6 +15,27 @@ from vllm.model_executor.model_loader.weight_utils import ( ) +def test_runai_safetensors_weights_iterator_clones_reused_buffers( + tmp_path, monkeypatch +): + 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) + ) + + 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/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index de15dc3d53c..972271e8c30 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -1090,7 +1090,8 @@ def runai_safetensors_weights_iterator( mininterval=2, ) - yield from tensor_iter + for name, tensor in tensor_iter: + yield name, tensor.clone() def _init_fastsafetensors_loader(