Use instanttensor == 0.1.8

Signed-off-by: mgoin <mgoin64@gmail.com>

requirements/test/cuda.in

@@ -59,7 +59,7 @@ numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
-instanttensor>=0.1.5
+instanttensor==0.1.8
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0; platform_machine == "x86_64"
 terratorch >= 1.2.2 # Required for Prithvi tests

requirements/test/cuda.txt

@@ -401,7 +401,7 @@ inflect==5.6.2
     # via datamodel-code-generator
 iniconfig==2.0.0
     # via pytest
-instanttensor==0.1.5
+instanttensor==0.1.8
     # via -r requirements/test/cuda.in
 isodate==0.7.2
     # via azure-storage-blob

requirements/test/nightly-torch.txt

@@ -44,5 +44,5 @@ numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2
-instanttensor>=0.1.5
+instanttensor==0.1.8
 pydantic>=2.12 # 2.11 leads to error on python 3.13

requirements/test/rocm.in

@@ -57,7 +57,7 @@ numba==0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
-instanttensor>=0.1.5
+instanttensor==0.1.8
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0
 

requirements/test/rocm.txt

@@ -467,7 +467,7 @@ inflect==7.5.0
     # via datamodel-code-generator
 iniconfig==2.3.0
     # via pytest
-instanttensor==0.1.6
+instanttensor==0.1.8
     # via -r requirements/test/rocm.in
 interegular==0.3.3
     # via lm-format-enforcer

setup.py

@@ -1079,7 +1079,7 @@ setup(
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.2.2"],
-        "instanttensor": ["instanttensor >= 0.1.5"],
+        "instanttensor": ["instanttensor == 0.1.8"],
         "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
         "audio": [
             "av",

vllm/model_executor/model_loader/default_loader.py

@@ -34,8 +34,10 @@ from vllm.model_executor.model_loader.weight_utils import (
     pt_weights_iterator,
     safetensors_weights_iterator,
 )
+from vllm.platforms import current_platform
 from vllm.tracing import instrument
 from vllm.transformers_utils.repo_utils import list_filtered_repo_files
+from vllm.utils.import_utils import has_instanttensor
 
 logger = init_logger(__name__)
 
@@ -228,7 +230,13 @@ class DefaultModelLoader(BaseModelLoader):
                     hf_weights_files,
                     self.load_config.use_tqdm_on_load,
                 )
-            elif self.load_config.load_format == "instanttensor":
+            elif self.load_config.load_format == "instanttensor" or (
+                self.load_config.load_format == "auto"
+                and has_instanttensor()
+                and current_platform.is_cuda()
+            ):
+                if self.load_config.load_format == "auto":
+                    logger.info("Using InstantTensor weight loader (auto-detected).")
                 weights_iterator = instanttensor_weights_iterator(
                     hf_weights_files,
                     self.load_config.use_tqdm_on_load,

vllm/model_executor/model_loader/weight_utils.py

@@ -1167,18 +1167,44 @@ def instanttensor_weights_iterator(
 
     device = current_platform.current_device()
 
+    # Cap InstantTensor's GPU buffer. Since we clone every yielded tensor out
+    # of the buffer on the next line, a large ring buffer is wasted GPU memory
+    # InstantTensor auto-enlarges this to ``max(tensor_sizes)`` if it's smaller
+    # than the largest tensor, so the cap is safe for any model.
+    buffer_size = 128 * 1024 * 1024
     with instanttensor.safe_open(
-        hf_weights_files, framework="pt", device=device, process_group=process_group
+        hf_weights_files,
+        framework="pt",
+        device=device,
+        process_group=process_group,
+        buffer_size=buffer_size,
     ) as f:
-        yield from tqdm(
-            f.tensors(),
+        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(),
-            total=len(f.keys()),
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
             mininterval=1.0,
         )
+        try:
+            for name, tensor in f.tensors():
+                # InstantTensor tensors are views into the ``safe_open`` ring
+                # buffer above and are invalidated when the context exits or
+                # the buffer is reused for a later tensor. Callers such as
+                # ``RobertaEmbeddingModel.load_weights`` fully materialize the
+                # iterator into a list before consuming it, which would leave
+                # those tensors dangling. Clone so each tensor owns its GPU
+                # storage independently of InstantTensor's buffer.
+                # Reference: https://github.com/scitix/InstantTensor/blob/45763a4a2eb4d1df7f05b988c01c76111c821b59/instanttensor/_impl.py#L535-L540
+                cloned = tensor.clone()
+                pbar.update(cloned.numel() * cloned.element_size())
+                yield name, cloned
+        finally:
+            pbar.close()
 
 
 def pt_weights_iterator(

vllm/utils/import_utils.py

@@ -471,3 +471,8 @@ def has_mori() -> bool:
 def has_fbgemm_gpu() -> bool:
     """Whether the optional `fbgemm_gpu` package is available."""
     return _has_module("fbgemm_gpu")
+
+
+def has_instanttensor() -> bool:
+    """Whether the optional `instanttensor` package is available."""
+    return _has_module("instanttensor")