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, 22:15 UTC
[Misc] Check that the model can be inspected upon registration (#13743)
Cyrus Leung committed
Feb 25, 2025 at 16:18 UTC
6724e791641ed3c604bf929dfa46cfc71fa8ba71
1 file changed
+15
-1
vllm/model_executor/models/registry.py
+15
-1
@@ -347,6 +347,10 @@ class _ModelRegistry:
347
when importing the model and thus the related error
348
:code:`RuntimeError: Cannot re-initialize CUDA in forked subprocess`.
349
"""
350
+ if not isinstance(model_arch, str):
351
+ msg = f"`model_arch` should be a string, not a {type(model_arch)}"
352
+ raise TypeError(msg)
353
+
354
if model_arch in self.models:
355
logger.warning(
356
"Model architecture %s is already registered, and will be "
@@ -360,8 +364,18 @@ class _ModelRegistry:
364
raise ValueError(msg)
365
366
model = _LazyRegisteredModel(*split_str)
363
- else:
367
+
368
+ try:
369
+ model.inspect_model_cls()
370
+ except Exception as exc:
371
+ msg = f"Unable to inspect model {model_cls}"
372
+ raise RuntimeError(msg) from exc
373
+ elif isinstance(model_cls, type) and issubclass(model_cls, nn.Module):
374
model = _RegisteredModel.from_model_cls(model_cls)
375
+ else:
376
+ msg = ("`model_cls` should be a string or PyTorch model class, "
377
+ f"not a {type(model_arch)}")
378
+ raise TypeError(msg)
379
380
self.models[model_arch] = model
381