The previous commit (93730cfd) replaced the blanket module.to() in
_process_weights_after_loading with an unconditional per-parameter move.
That broke the FSDP/HSDP path: after apply_hsdp_to_model the params are
FSDP-sharded DTensors whose storage cannot be cross-device reassigned via
`param.data = param.data.to(cpu)`, raising at load time:
RuntimeError: Attempted to set the storage of a tensor on device
"cuda:0" to a storage on different device "cpu". This is no longer
allowed; the devices must match.
Regressed nightly test_zimage[layerwise_hsdp] (Diffusion X2I Function
Test with L4), which passed pre-fix. The FP8 online-quant variants it was
meant to fix (qwen/bagel fp8, parallel_*_fp8_*) all passed.
Meta-device params only ever arise from online quant (uses_meta_device),
so gate the meta-safe per-parameter move to that case and keep the
original FSDP/HSDP-aware whole-module module.to() for everything else.
Online quant initializes on the accelerator, so its params are normally
already on the target device (the per-param loop is a no-op move) — the
only behavioral change vs. the original is skipping meta stragglers.
Co-Authored-By: Claude <noreply@anthropic.com>
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>
