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 <noreply@anthropic.com>
This commit is contained in:
tzhouamandClaude committed 2026-07-11 14:21:16 +08:00
1 parent 49f4891e3e
commit f74e98119e
1 file changed
+47 -10
@@ -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)