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 15, 19:45 UTC
[Misc] add `use_tqdm_on_load` to reduce logs (#14407)
Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
Aaron Pham committed
Mar 8, 2025 at 08:57 UTC
0b7f06b447e513dabfb87f490713516943c7c371
4 files changed
+54
-22
vllm/config.py
+3
@@ -1277,6 +1277,8 @@ class LoadConfig:
1277
ignore_patterns: The list of patterns to ignore when loading the model.
1278
Default to "original/**/*" to avoid repeated loading of llama's
1279
checkpoints.
1280
+ use_tqdm_on_load: Whether to enable tqdm for showing progress bar during
1281
+ loading. Default to True
1282
"""
1283
1284
load_format: Union[str, LoadFormat, "BaseModelLoader"] = LoadFormat.AUTO
@@ -1284,6 +1286,7 @@ class LoadConfig:
1286
model_loader_extra_config: Optional[Union[str, dict]] = field(
1287
default_factory=dict)
1288
ignore_patterns: Optional[Union[list[str], str]] = None
1289
+ use_tqdm_on_load: bool = True
1290
1291
def compute_hash(self) -> str:
1292
"""
vllm/engine/arg_utils.py
+10
@@ -217,6 +217,7 @@ class EngineArgs:
217
additional_config: Optional[Dict[str, Any]] = None
218
enable_reasoning: Optional[bool] = None
219
reasoning_parser: Optional[str] = None
220
+ use_tqdm_on_load: bool = True
221
222
def __post_init__(self):
223
if not self.tokenizer:
@@ -751,6 +752,14 @@ class EngineArgs:
752
default=1,
753
help=('Maximum number of forward steps per '
754
'scheduler call.'))
755
+ parser.add_argument(
756
+ '--use-tqdm-on-load',
757
+ dest='use_tqdm_on_load',
758
+ action=argparse.BooleanOptionalAction,
759
+ default=EngineArgs.use_tqdm_on_load,
760
+ help='Whether to enable/disable progress bar '
761
+ 'when loading model weights.',
762
+ )
763
764
parser.add_argument(
765
'--multi-step-stream-outputs',
@@ -1179,6 +1188,7 @@ class EngineArgs:
1188
download_dir=self.download_dir,
1189
model_loader_extra_config=self.model_loader_extra_config,
1190
ignore_patterns=self.ignore_patterns,
1191
+ use_tqdm_on_load=self.use_tqdm_on_load,
1192
)
1193
1194
def create_engine_config(self,
vllm/model_executor/model_loader/loader.py
+21
-5
@@ -354,11 +354,18 @@ class DefaultModelLoader(BaseModelLoader):
354
self.load_config.download_dir,
355
hf_folder,
356
hf_weights_files,
357
+ self.load_config.use_tqdm_on_load,
358
)
359
elif use_safetensors:
359
- weights_iterator = safetensors_weights_iterator(hf_weights_files)
360
+ weights_iterator = safetensors_weights_iterator(
361
+ hf_weights_files,
362
+ self.load_config.use_tqdm_on_load,
363
+ )
364
else:
361
- weights_iterator = pt_weights_iterator(hf_weights_files)
365
+ weights_iterator = pt_weights_iterator(
366
+ hf_weights_files,
367
+ self.load_config.use_tqdm_on_load,
368
+ )
369
370
if current_platform.is_tpu():
371
# In PyTorch XLA, we should call `xm.mark_step` frequently so that
@@ -806,9 +813,15 @@ class BitsAndBytesModelLoader(BaseModelLoader):
813
814
def _hf_weight_iter(self, hf_weights_files, use_safetensors: bool):
815
if use_safetensors:
809
- iterator = safetensors_weights_iterator(hf_weights_files)
816
+ iterator = safetensors_weights_iterator(
817
+ hf_weights_files,
818
+ self.load_config.use_tqdm_on_load,
819
+ )
820
else:
811
- iterator = pt_weights_iterator(hf_weights_files)
821
+ iterator = pt_weights_iterator(
822
+ hf_weights_files,
823
+ self.load_config.use_tqdm_on_load,
824
+ )
825
for org_name, param in iterator:
826
# mapping weight names from transformers to vllm while preserving
827
# original names.
@@ -1396,7 +1409,10 @@ class RunaiModelStreamerLoader(BaseModelLoader):
1409
revision: str) -> Generator[Tuple[str, torch.Tensor], None, None]:
1410
"""Get an iterator for the model weights based on the load format."""
1411
hf_weights_files = self._prepare_weights(model_or_path, revision)
1399
- return runai_safetensors_weights_iterator(hf_weights_files)
1412
+ return runai_safetensors_weights_iterator(
1413
+ hf_weights_files,
1414
+ self.load_config.use_tqdm_on_load,
1415
+ )
1416
1417
def download_model(self, model_config: ModelConfig) -> None:
1418
"""Download model if necessary"""
vllm/model_executor/model_loader/weight_utils.py
+20
-17
@@ -366,16 +366,22 @@ def filter_files_not_needed_for_inference(
366
_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501
367
368
369
+def enable_tqdm(use_tqdm_on_load: bool):
370
+ return use_tqdm_on_load and (not torch.distributed.is_initialized()
371
+ or torch.distributed.get_rank() == 0)
372
+
373
+
374
def np_cache_weights_iterator(
370
- model_name_or_path: str, cache_dir: Optional[str], hf_folder: str,
371
- hf_weights_files: List[str]
375
+ model_name_or_path: str,
376
+ cache_dir: Optional[str],
377
+ hf_folder: str,
378
+ hf_weights_files: List[str],
379
+ use_tqdm_on_load: bool,
380
) -> Generator[Tuple[str, torch.Tensor], None, None]:
381
"""Iterate over the weights in the model np files.
382
383
Will dump the model weights to numpy files if they are not already dumped.
384
"""
377
- enable_tqdm = not torch.distributed.is_initialized(
378
- ) or torch.distributed.get_rank() == 0
385
# Convert the model weights from torch tensors to numpy arrays for
386
# faster loading.
387
np_folder = os.path.join(hf_folder, "np")
@@ -389,7 +395,7 @@ def np_cache_weights_iterator(
395
for bin_file in tqdm(
396
hf_weights_files,
397
desc="Loading np_cache checkpoint shards",
392
- disable=not enable_tqdm,
398
+ disable=not enable_tqdm(use_tqdm_on_load),
399
bar_format=_BAR_FORMAT,
400
):
401
state = torch.load(bin_file,
@@ -414,15 +420,14 @@ def np_cache_weights_iterator(
420
421
422
def safetensors_weights_iterator(
417
- hf_weights_files: List[str]
423
+ hf_weights_files: List[str],
424
+ use_tqdm_on_load: bool,
425
) -> Generator[Tuple[str, torch.Tensor], None, None]:
426
"""Iterate over the weights in the model safetensor files."""
420
- enable_tqdm = not torch.distributed.is_initialized(
421
- ) or torch.distributed.get_rank() == 0
427
for st_file in tqdm(
428
hf_weights_files,
429
desc="Loading safetensors checkpoint shards",
425
- disable=not enable_tqdm,
430
+ disable=not enable_tqdm(use_tqdm_on_load),
431
bar_format=_BAR_FORMAT,
432
):
433
with safe_open(st_file, framework="pt") as f:
@@ -432,16 +437,15 @@ def safetensors_weights_iterator(
437
438
439
def runai_safetensors_weights_iterator(
435
- hf_weights_files: List[str]
440
+ hf_weights_files: List[str],
441
+ use_tqdm_on_load: bool,
442
) -> Generator[Tuple[str, torch.Tensor], None, None]:
443
"""Iterate over the weights in the model safetensor files."""
438
- enable_tqdm = not torch.distributed.is_initialized(
439
- ) or torch.distributed.get_rank() == 0
444
with SafetensorsStreamer() as streamer:
445
for st_file in tqdm(
446
hf_weights_files,
447
desc="Loading safetensors using Runai Model Streamer",
444
- disable=not enable_tqdm,
448
+ disable=not enable_tqdm(use_tqdm_on_load),
449
bar_format=_BAR_FORMAT,
450
):
451
streamer.stream_file(st_file)
@@ -449,15 +453,14 @@ def runai_safetensors_weights_iterator(
453
454
455
def pt_weights_iterator(
452
- hf_weights_files: List[str]
456
+ hf_weights_files: List[str],
457
+ use_tqdm_on_load: bool,
458
) -> Generator[Tuple[str, torch.Tensor], None, None]:
459
"""Iterate over the weights in the model bin/pt files."""
455
- enable_tqdm = not torch.distributed.is_initialized(
456
- ) or torch.distributed.get_rank() == 0
460
for bin_file in tqdm(
461
hf_weights_files,
462
desc="Loading pt checkpoint shards",
460
- disable=not enable_tqdm,
463
+ disable=not enable_tqdm(use_tqdm_on_load),
464
bar_format=_BAR_FORMAT,
465
):
466
state = torch.load(bin_file, map_location="cpu", weights_only=True)