From 533dbe424b9bdde6a7ffe100d507a69a24dfee1a Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sun, 13 Sep 2026 11:16:39 +0800 Subject: [PATCH] [Misc] Clean up fast loader daemon quant method verification (#55656) Signed-off-by: Isotr0py --- .../model_loader/base_loader.py | 18 ++- .../model_loader/weight_cache/__init__.py | 2 - .../model_loader/weight_cache/daemon.py | 38 +++++-- .../model_loader/weight_cache/ipc_loader.py | 28 ++--- .../model_loader/weight_cache/protocol.py | 106 ++++-------------- 5 files changed, 75 insertions(+), 117 deletions(-) diff --git a/vllm/model_executor/model_loader/base_loader.py b/vllm/model_executor/model_loader/base_loader.py index 618f874fd0..6f35646458 100644 --- a/vllm/model_executor/model_loader/base_loader.py +++ b/vllm/model_executor/model_loader/base_loader.py @@ -39,6 +39,19 @@ class BaseModelLoader(ABC): inplace weights loading for an already-initialized model""" raise NotImplementedError + def create_model( + self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" + ) -> nn.Module: + """Create a model with the given configurations.""" + model = initialize_model( + vllm_config=vllm_config, + model_config=model_config, + prefix=prefix, + ) + log_online_quantization(vllm_config) + log_model_inspection(model) + return model + @instrument(span_name="Load model") def load_model( self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" @@ -52,15 +65,12 @@ class BaseModelLoader(ABC): target_device = torch.device(load_device) with set_default_torch_dtype(model_config.dtype): with target_device: - model = initialize_model( + model = self.create_model( vllm_config=vllm_config, model_config=model_config, prefix=prefix, ) - log_online_quantization(vllm_config) - log_model_inspection(model) - logger.debug("Loading weights on %s ...", load_device) self.load_weights(model, model_config) diff --git a/vllm/model_executor/model_loader/weight_cache/__init__.py b/vllm/model_executor/model_loader/weight_cache/__init__.py index 0cea5cf72a..f9dbe9bb32 100644 --- a/vllm/model_executor/model_loader/weight_cache/__init__.py +++ b/vllm/model_executor/model_loader/weight_cache/__init__.py @@ -11,7 +11,6 @@ from vllm.model_executor.model_loader.weight_cache.protocol import ( WeightCacheUnavailableError, check_ipc_platform_support, check_ipc_quant_support, - is_ipc_quant_supported, ) __all__ = [ @@ -24,5 +23,4 @@ __all__ = [ "WeightCacheUnavailableError", "check_ipc_platform_support", "check_ipc_quant_support", - "is_ipc_quant_supported", ] diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 083910abb4..31f59a6ac5 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -39,6 +39,8 @@ from vllm.distributed import ( init_distributed_environment, ) from vllm.logger import init_logger +from vllm.model_executor.model_loader import get_model_loader +from vllm.model_executor.model_loader.utils import process_weights_after_loading from vllm.model_executor.model_loader.weight_cache.protocol import ( TensorEntry, WeightCacheKey, @@ -53,6 +55,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import ( verify_peer_is_owner, ) from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_default_torch_dtype logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon") @@ -98,6 +101,30 @@ def export_entries( return entries, aliases +def get_daemon_model(vllm_config: VllmConfig) -> torch.nn.Module: + """Load the daemon's model, composed from the configured loader. + + Runs the quantization check after model creation but before the slow + weight load, so an unsupported method fails fast. Online quantization + always fails the check, so load_model's finalize step for it is + unnecessary here. + """ + model_config = vllm_config.model_config + load_config = vllm_config.load_config + loader = get_model_loader(load_config) + device_config = vllm_config.device_config + target_device = torch.device( + device_config.device if load_config.device is None else load_config.device + ) + with set_default_torch_dtype(model_config.dtype): + with target_device: + model = loader.create_model(vllm_config, model_config) + check_ipc_quant_support(model) + loader.load_weights(model, model_config) + process_weights_after_loading(model, model_config, target_device) + return model.eval() + + class WeightCacheDaemon: """Per-GPU process that loads one TP shard and serves CUDA IPC handles.""" @@ -124,8 +151,6 @@ class WeightCacheDaemon: ) def load_model(self) -> None: - from vllm.model_executor.model_loader import get_model - tp_size = self.cache_config.tp_size torch.accelerator.set_device_index(self.tp_rank) init_distributed_environment( @@ -137,7 +162,7 @@ class WeightCacheDaemon: ) with set_current_vllm_config(self.vllm_config): ensure_model_parallel_initialized(tp_size, 1) - self.model = get_model(vllm_config=self.vllm_config) + self.model = get_daemon_model(self.vllm_config) self._export_entries() logger.info( "Weight cache daemon rank %d cached %d tensors", @@ -331,10 +356,9 @@ def main() -> None: "The weight cache daemon itself must load from disk; use the " "default --load-format" ) - # Checked before loading anything: an unsupported quantization method would - # otherwise only surface in the engine, after a full load. - check_ipc_platform_support(where="daemon") - check_ipc_quant_support(vllm_config.model_config, where="daemon") + # Config-only so it can fail before any model loading; the quant method + # check needs the created model and runs in get_daemon_model. + check_ipc_platform_support() parallel_config = vllm_config.parallel_config _reject_unsupported_parallelism(parallel_config) tp_size = parallel_config.tensor_parallel_size diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index 2132a55fbe..be2e5ebae4 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -119,9 +119,9 @@ class IpcModelLoader(BaseModelLoader): def load_model( self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = "" ) -> nn.Module: - # Unsupported quantization is a permanent misconfiguration rather than a - # transient daemon outage, so it is raised even when fallback is on. - self._check_supported(vllm_config, model_config) + # An unsupported platform is a permanent misconfiguration rather than + # a transient daemon outage, so it is raised even when fallback is on. + check_ipc_platform_support() state_fetched = False try: entries, aliases = self._fetch_entries(model_config) @@ -135,6 +135,11 @@ class IpcModelLoader(BaseModelLoader): logger.warning( "Weight cache unusable (%s); falling back to disk loading", e ) + except UnsupportedQuantForIPCError: + # Unsupported quantization is a permanent misconfiguration rather + # than a transient daemon outage, so it is raised even when + # fallback is on. + raise except Exception: if not self.fallback: raise @@ -176,6 +181,7 @@ class IpcModelLoader(BaseModelLoader): model_config=model_config, prefix=prefix, ) + check_ipc_quant_support(model) self._apply_entries(model, entries, aliases, device_index) # The daemon exports tensors that already went through # process_weights_after_loading; re-run it in pre-processed mode @@ -197,22 +203,6 @@ class IpcModelLoader(BaseModelLoader): ) return model.eval() - @staticmethod - def _check_supported(vllm_config: VllmConfig, model_config: ModelConfig) -> None: - check_ipc_platform_support(where="engine") - check_ipc_quant_support(model_config, where="engine") - cache_dtype = vllm_config.cache_config.cache_dtype - if cache_dtype != "auto" and not str(cache_dtype).startswith("fp8"): - # BaseKVCacheMethod.process_weights_after_loading turns the loaded - # k/v scale parameters into plain float attributes. For fp8 cache - # dtypes those are rebuilt from the exported scale buffers when - # process_weights_after_loading runs in pre-processed mode; other - # quantized cache dtypes are not verified. - raise UnsupportedQuantForIPCError( - f"[weight_cache:engine] kv cache dtype {cache_dtype!r} is not " - "supported by the weight cache; use --kv-cache-dtype auto." - ) - def _apply_entries( self, model: nn.Module, diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 08e5e725bd..fff1fcca2a 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -26,6 +26,7 @@ import torch import vllm.version from vllm.config import ModelConfig +from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.platforms import current_platform from vllm.utils.hashing import safe_hash @@ -59,15 +60,12 @@ class UnsupportedPlatformForIPCError(Exception): """Raised when the current platform cannot share CUDA IPC handles.""" -def check_ipc_platform_support(*, where: str) -> None: +def check_ipc_platform_support() -> None: """Hard-error unless the current platform can share CUDA IPC handles. Only CUDA/ROCm tensors get a real IPC handle from ``TensorEntry``; other platforms (e.g. XPU) would silently ship every tensor by value instead. - Args: - where: Short tag ("daemon"/"engine") used in the error message. - Raises: UnsupportedPlatformForIPCError: If the current platform is not CUDA/ROCm. @@ -75,95 +73,33 @@ def check_ipc_platform_support(*, where: str) -> None: if current_platform.is_cuda_alike(): return raise UnsupportedPlatformForIPCError( - f"[weight_cache:{where}] platform {current_platform.device_name!r} " - "does not support CUDA IPC weight sharing; only CUDA and ROCm are " - "supported. Use the default --load-format for this platform." + f"platform {current_platform.device_name!r} does not support CUDA IPC " + "weight sharing; only CUDA and ROCm are supported. Use the default " + "--load-format for this platform." ) -# The daemon exports tensor data only, so sharing is correct just for methods -# whose post-load effect is either fully captured by that data or rebuilt when -# process_weights_after_loading runs under weights_already_processed. Anything -# absent from this registry hard-errors: a method that silently repacks into -# shapes the client cannot reproduce or stamps unrebuildable Python-side state -# would serve wrong numerics. Extend it only after an end-to-end check against -# a disk-loaded baseline. -# -# This check is config-only so the daemon fails fast before a slow full load, -# rather than surfacing the incompatibility only after the engine has built the -# model (where process_weights_after_loading's supports_pre_processed_weights -# guard would eventually catch it). - - -def _quant_config_field(quant_config: Any, key: str) -> Any: - if quant_config is None: - return None - if isinstance(quant_config, dict): - return quant_config.get(key) - return getattr(quant_config, key, None) - - -def _fp8_round_trips_via_ipc(quant_config: Any) -> bool: - """Only block-wise FP8 is verified. - - Block-wise FP8 preserves the weight shape, while per-tensor FP8 transposes - ``layer.weight`` during post-processing -- a shape the client's - meta-initialized model cannot reproduce. - """ - return _quant_config_field(quant_config, "weight_block_size") is not None - - -# quantization name -> predicate(quant_config) -> True when verified safe. -IPC_QUANT_ALLOWLIST: dict[str | None, Any] = { - None: lambda _quant_config: True, # unquantized - "fp8": _fp8_round_trips_via_ipc, - "modelopt_fp4": lambda _quant_config: True, - # Kimi-K3 routed experts (mxfp4-pack): the MegaMoE experts keep their packed - # weights and scales as plain parameters and defer the DeepGEMM transform to - # forward(), so the exported tensors are complete. Any other mxfp4 method - # that repacks or drops parameters still fails closed at the engine's - # supports_pre_processed_weights guard. - "mxfp4": lambda _quant_config: True, -} - - -def is_ipc_quant_supported(quantization: str | None, quant_config: Any) -> bool: - predicate = IPC_QUANT_ALLOWLIST.get(quantization) - return False if predicate is None else bool(predicate(quant_config)) - - -def check_ipc_quant_support(model_config: ModelConfig, *, where: str) -> None: - """Hard-error unless the model's quantization is verified for IPC sharing. +# The daemon transfers post processed weights directly +def check_ipc_quant_support(model: torch.nn.Module) -> None: + """Hard-error unless every quant method supports pre-processed weights. Args: - model_config: Model configuration to inspect. - where: Short tag ("daemon"/"engine") used in the error message. + model: The model to inspect (weights need not be loaded). Raises: - UnsupportedQuantForIPCError: If the quantization method is not on the - verified allowlist. + UnsupportedQuantForIPCError: If any quant method does not declare + ``supports_pre_processed_weights``. """ - quantization = model_config.quantization - # Prefer the canonical, nested-aware config (multimodal models keep it under - # text_config); fall back to the raw hf_config for older code paths. - quant_config = getattr( - getattr(model_config, "model_arch_config", None), "quantization_config", None - ) - if quant_config is None: - quant_config = getattr(model_config.hf_config, "quantization_config", None) - if is_ipc_quant_supported(quantization, quant_config): - return - verified = ", ".join( - "unquantized" if name is None else repr(name) for name in IPC_QUANT_ALLOWLIST - ) - raise UnsupportedQuantForIPCError( - f"[weight_cache:{where}] quantization {quantization!r} is not verified " - f"for CUDA IPC weight sharing: its post-load processing may repack " - f"weights into shapes the client cannot reproduce or stamp Python-side " - f"state that tensor export cannot carry, which would silently serve " - f"wrong numerics. Verified: {verified} (FP8 only with weight_block_size " - f"set, i.e. block-wise). Use the default --load-format for this model." - ) + for name, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if ( + isinstance(quant_method, QuantizeMethodBase) + and not quant_method.supports_pre_processed_weights + ): + raise UnsupportedQuantForIPCError( + f"layer {name or ''}: {type(quant_method).__name__} " + "does not support loading from pre-processed weights." + ) def get_physical_device_id(device_index: int) -> int | None: