Fix RunAI streamer tensor buffer reuse during weight loading (#43464)

Signed-off-by: bbartels <[email protected]>

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

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(