From f74e98119e1b871143984360ef611c6b94b5eeeb Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 8 Jul 2026 21:31:29 +0800 Subject: [PATCH] fix(diffusion): make FP8 online-quant weight loading meta-device safe Newer upstream vLLM (65a7b4628409, 585 commits past the v0.24.0 tag that vllm-omni main tracks) reworked online quantization: online-quant linear methods now declare `uses_meta_device=True` and create weights on the `meta` device, materializing them just-in-time via the layerwise online-process loader as each layer's weights finish loading. The diffusion loader's `_process_weights_after_loading` did a blanket `module.to(target_device)`, which raises `NotImplementedError: Cannot copy out of meta tensor; no data!` on any layer whose weights were not fully materialized during load (padded / partially-loaded "straggler" layers), crashing diffusion worker startup. This regressed the nightly diffusion FP8 tests on dev/vllm-align (all green on main nightly): Quantization Test (H100 + L4, test_single_stage_qwen_image_fp8 / test_bagel_fp8_generates_image) and the Qwen-Image Multi-GPU vae_patch_parallel_2 case (fp8). Mirror upstream's loader contract: - Materialize online-quant stragglers via finalize_layerwise_processing() before processing (imported lazily, only when online quant is in use, so older vLLM is unaffected). model_config=None is safe: it is only used for vLLM Attention/MLAAttention layers, which diffusion DiT models do not have. - Replace the blanket module.to() with a per-parameter move that mirrors upstream's device_loading_context: move only real (non-meta) params to the target device for processing and restore them afterward, leaving meta params for the quant method / finalize step. Co-Authored-By: Claude --- .../model_loader/diffusers_loader.py | 57 +++++++++++++++---- 1 file changed, 47 insertions(+), 10 deletions(-) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 7ee2b4395..c6bf97c3c 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -397,28 +397,65 @@ class DiffusersPipelineLoader: return model.eval() + @staticmethod + def _has_online_quant(model: nn.Module) -> bool: + """Whether any layer uses an online-quant method that defers weight + materialization onto the ``meta`` device (upstream vLLM + ``uses_meta_device=True``, e.g. online FP8).""" + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if getattr(quant_method, "uses_meta_device", False): + return True + return False + def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None: """Process weights after loading for quantization methods. This handles vLLM's quantization methods that need to process weights after loading (e.g., FP8 online quantization from BF16/FP16 weights). """ + # Newer upstream vLLM online-quant methods (uses_meta_device=True) create + # weights on the ``meta`` device and materialize them just-in-time as each + # layer's weights finish loading (via the layerwise online-process loader). + # Any "straggler" layers whose weights were not fully materialized during + # load (padded / partially-loaded layers) remain on ``meta``. Upstream's + # base_loader calls finalize_layerwise_processing() to materialize them; + # the diffusion loader must mirror that, otherwise the module.to() below + # raises "Cannot copy out of meta tensor; no data!". Import lazily and only + # when online quant is actually in use, so older vLLM (no meta-device quant) + # is unaffected. + if self._has_online_quant(model): + from vllm.model_executor.model_loader.reload.layerwise import ( + finalize_layerwise_processing, + ) + + # model_config is only dereferenced by finalize for vLLM Attention / + # MLAAttention layers; diffusion DiT models use their own attention and + # have none, so passing None is safe here. + finalize_layerwise_processing(model, model_config=None) + for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) if quant_method is not None and isinstance(quant_method, QuantizeMethodBase): - # Move module to target device for processing if needed - module_device = next(module.parameters(), None) - if module_device is not None: - module_device = module_device.device - needs_device_move = module_device != target_device - - if needs_device_move: - module.to(target_device) + # Move only real (non-meta) params onto the target device for + # processing and restore them afterward, mirroring upstream vLLM's + # device_loading_context. A blanket module.to(target_device) would + # raise NotImplementedError on any param still on the meta device + # (deferred online-quant materialization); those are handled by the + # quant method / finalize step above and must not be moved with .to(). + original_devices: dict[str, torch.device] = {} + for name, param in module.named_parameters(): + if param.device.type != "meta" and param.device != target_device: + original_devices[name] = param.device + param.data = param.data.to(target_device) quant_method.process_weights_after_loading(module) - if needs_device_move: - module.to(module_device) + # Restore pre-existing params to their original device; leave any + # newly created (e.g. quantized) params on the target device. + for name, param in module.named_parameters(): + if name in original_devices: + param.data = param.data.to(original_devices[name]) def load_weights(self, model: nn.Module) -> None: weights_to_load = self._get_expected_parameter_names(model)