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 3, 04:15 UTC
[CI] De-flake test_models for bigscience/bloom-560m (#43197)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
May 21, 2026 at 14:30 UTC
a950e9447e38727fc956afdc242bc6e3796ccb77
1 file changed
+37
-16
vllm/transformers_utils/config.py
+37
-16
@@ -1204,6 +1204,15 @@ def try_get_dense_modules(
1204
return None
1205
1206
1207
+def _read_safetensors_metadata_in_dir(local_dir: Path) -> dict[str, Any]:
1208
+ return {
1209
+ param_name: info
1210
+ for file_path in local_dir.glob("*.safetensors")
1211
+ if file_path.is_file()
1212
+ for param_name, info in parse_safetensors_file_metadata(file_path).items()
1213
+ }
1214
+
1215
+
1216
def get_safetensors_params_metadata(
1217
model: str,
1218
*,
@@ -1212,24 +1221,36 @@ def get_safetensors_params_metadata(
1221
"""
1222
Get the safetensors parameters metadata for remote/local model repository.
1223
"""
1215
- full_metadata = {}
1224
if (model_path := Path(model)).exists():
1217
- safetensors_to_check = model_path.glob("*.safetensors")
1218
- full_metadata = {
1219
- param_name: info
1220
- for file_path in safetensors_to_check
1221
- if file_path.is_file()
1222
- for param_name, info in parse_safetensors_file_metadata(file_path).items()
1225
+ return _read_safetensors_metadata_in_dir(model_path)
1226
+
1227
+ repo_mt = try_get_safetensors_metadata(model, revision=revision)
1228
+ if repo_mt and (files_mt := repo_mt.files_metadata):
1229
+ return {
1230
+ param_name: asdict(info)
1231
+ for file_mt in files_mt.values()
1232
+ for param_name, info in file_mt.tensors.items()
1233
}
1224
- else:
1225
- repo_mt = try_get_safetensors_metadata(model, revision=revision)
1226
- if repo_mt and (files_mt := repo_mt.files_metadata):
1227
- full_metadata = {
1228
- param_name: asdict(info)
1229
- for file_mt in files_mt.values()
1230
- for param_name, info in file_mt.tensors.items()
1231
- }
1232
- return full_metadata
1234
+
1235
+ # Hub fetch failed (e.g. 429, network unreachable). Fall back to the
1236
+ # local HF cache: weights may already be cached from a prior run, and
1237
+ # weight loading itself uses the same cache.
1238
+ try:
1239
+ local_dir = huggingface_hub.snapshot_download(
1240
+ repo_id=model,
1241
+ revision=revision,
1242
+ allow_patterns=["*.safetensors"],
1243
+ local_files_only=True,
1244
+ )
1245
+ except huggingface_hub.errors.LocalEntryNotFoundError as e:
1246
+ logger.warning_once(
1247
+ "Could not retrieve safetensors metadata for %s "
1248
+ "(Hub fetch failed and no local cache snapshot is available): %s.",
1249
+ model,
1250
+ str(e),
1251
+ )
1252
+ return {}
1253
+ return _read_safetensors_metadata_in_dir(Path(local_dir))
1254
1255
1256
def _download_mistral_config_file(model, revision) -> dict: