A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 01:30 UTC
fix: Use iterator as not to store all the file loads in memory at once (#36149)
Signed-off-by: Shaun Kotek - Nvidia <skotek@nvidia.com>
Shaun Kotek committed
Mar 9, 2026 at 05:25 UTC
90512b2e8bff5bddca5fca30dc4f0136d682f7d4
1 file changed
+6
-2
vllm/model_executor/model_loader/weight_utils.py
+6
-2
@@ -773,7 +773,9 @@ def multi_thread_safetensors_weights_iterator(
773
return result
774
775
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
776
- futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files]
776
+ # Note to use generator here so we do not store all the loaded files in memory
777
+ # at the same time, which can cause OOM for large models.
778
+ futures = (executor.submit(_load_file, st_file) for st_file in hf_weights_files)
779
futures_iter = tqdm(
780
concurrent.futures.as_completed(futures),
781
total=len(hf_weights_files),
@@ -784,7 +786,9 @@ def multi_thread_safetensors_weights_iterator(
786
787
for future in futures_iter:
788
state_dict = future.result()
787
- yield from state_dict.items()
789
+ del future
790
+ for key in list(state_dict):
791
+ yield key, state_dict.pop(key)
792
793
794
def runai_safetensors_weights_iterator(