[Loader] Improve InstantTensor loading (#46868)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Michael Goin committed
Jul 17, 2026 at 16:30 UTC
cc25f028b76844024fddd0b4ca18d4ca169b32be
9 files changed
+27
-14
requirements/test/cpu.txt
+1
-1
@@ -379,7 +379,7 @@ inflect==5.6.2
379
# via datamodel-code-generator
380
iniconfig==2.0.0
381
# via pytest
382
-instanttensor==0.1.5
382
+instanttensor==0.1.9
383
# via -r requirements/test/cuda.in
384
interegular==0.3.3
385
# via lm-format-enforcer
requirements/test/cuda.in
+1
-1
@@ -58,7 +58,7 @@ arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix de
58
numba == 0.65.0 # Required for N-gram speculative decoding
59
runai-model-streamer[s3,gcs,azure]==0.15.7
60
fastsafetensors>=0.3.2
61
-instanttensor>=0.1.5; platform_machine == "x86_64"
61
+instanttensor>=0.1.9; platform_machine == "x86_64"
62
decord==0.6.0; platform_machine == "x86_64"
63
# terratorch is temporarily disabled while PyPI has the `lightning` package
64
# in `quarantined` status (every published terratorch version transitively
requirements/test/cuda.txt
+1
-1
@@ -398,7 +398,7 @@ inflect==5.6.2
398
# via datamodel-code-generator
399
iniconfig==2.0.0
400
# via pytest
401
-instanttensor==0.1.5
401
+instanttensor==0.1.9
402
# via -r requirements/test/cuda.in
403
interegular==0.3.3
404
# via lm-format-enforcer
requirements/test/nightly-torch.txt
+1
-1
@@ -44,5 +44,5 @@ numba == 0.65.0 # Required for N-gram speculative decoding
44
numpy
45
runai-model-streamer[s3,gcs,azure]==0.15.7
46
fastsafetensors>=0.3.2
47
-instanttensor>=0.1.5
47
+instanttensor>=0.1.9
48
pydantic>=2.12 # 2.11 leads to error on python 3.13
requirements/test/rocm.in
+1
-1
@@ -54,7 +54,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
54
numba==0.65.0 # Required for N-gram speculative decoding
55
runai-model-streamer[s3,gcs,azure]==0.15.7
56
fastsafetensors>=0.3.2
57
-instanttensor>=0.1.5
57
+instanttensor>=0.1.9
58
decord==0.6.0
59
60
# Prithvi tests
requirements/test/rocm.txt
+1
-1
@@ -391,7 +391,7 @@ inflect==7.5.0
391
# via datamodel-code-generator
392
iniconfig==2.3.0
393
# via pytest
394
-instanttensor==0.1.6
394
+instanttensor==0.1.9
395
# via -r requirements/test/rocm.in
396
interegular==0.3.3
397
# via lm-format-enforcer
setup.py
+1
-1
@@ -1268,7 +1268,7 @@ setup(
1268
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
1269
"tensorizer": ["tensorizer==2.10.1"],
1270
"fastsafetensors": ["fastsafetensors >= 0.3.2"],
1271
- "instanttensor": ["instanttensor >= 0.1.5"],
1271
+ "instanttensor": ["instanttensor >= 0.1.9"],
1272
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
1273
"audio": [
1274
"av",
tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py
-2
@@ -33,8 +33,6 @@ def test_instanttensor_model_loader():
33
hf_safetensors_tensors = {}
34
35
for name, tensor in instanttensor_weights_iterator(safetensors, True):
36
- # Copy the tensor immediately as it is a reference to the internal
37
- # buffer of instanttensor.
36
instanttensor_tensors[name] = tensor.to("cpu")
37
38
for name, tensor in safetensors_weights_iterator(safetensors, True):
vllm/model_executor/model_loader/weight_utils.py
+20
-5
@@ -1120,7 +1120,7 @@ def instanttensor_weights_iterator(
1120
import instanttensor
1121
except ImportError as e:
1122
raise ImportError(
1123
- "Please install instanttensor via `pip install instanttensor`"
1123
+ "Please install instanttensor via `pip install vllm[instanttensor]`"
1124
) from e
1125
1126
if not current_platform.is_cuda():
@@ -1136,18 +1136,33 @@ def instanttensor_weights_iterator(
1136
1137
device = current_platform.current_device()
1138
1139
+ # copy=True yields tensors that own their memory, staying valid after the
1140
+ # context exits or InstantTensor reuses its buffer.
1141
with instanttensor.safe_open(
1140
- hf_weights_files, framework="pt", device=device, process_group=process_group
1142
+ hf_weights_files,
1143
+ framework="pt",
1144
+ device=device,
1145
+ process_group=process_group,
1146
+ copy=True,
1147
) as f:
1142
- yield from tqdm(
1143
- f.tensors(),
1148
+ # Track bytes so the bar reports load throughput (GB/s).
1149
+ pbar = tqdm(
1150
+ total=f.total_tensor_size,
1151
desc="Loading safetensors using InstantTensor loader",
1152
disable=not enable_tqdm(use_tqdm_on_load),
1153
bar_format=_BAR_FORMAT,
1154
position=tqdm._get_free_pos(),
1148
- total=len(f.keys()),
1155
+ unit="B",
1156
+ unit_scale=True,
1157
+ unit_divisor=1024,
1158
mininterval=1.0,
1159
)
1160
+ try:
1161
+ for name, tensor in f.tensors():
1162
+ pbar.update(tensor.numel() * tensor.element_size())
1163
+ yield name, tensor
1164
+ finally:
1165
+ pbar.close()
1166
1167
1168
def pt_weights_iterator(