fix(diffusion): gate meta-safe weight move to online quant only (fix HSDP regression)

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>

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -421,10 +421,14 @@ class DiffusersPipelineLoader:
         # 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):
+        # raises "Cannot copy out of meta tensor; no data!". This whole meta-device
+        # handling is gated on online quant actually being in use, so that the
+        # proven code path for everything else (in particular FSDP/HSDP-sharded
+        # params, whose per-parameter .data cannot be cross-device reassigned) is
+        # left untouched. Import lazily so older vLLM (no meta-device quant) is
+        # unaffected.
+        has_online_quant = self._has_online_quant(model)
+        if has_online_quant:
             from vllm.model_executor.model_loader.reload.layerwise import (
                 finalize_layerwise_processing,
             )
@@ -436,13 +440,17 @@ class DiffusersPipelineLoader:
 
         for _, module in model.named_modules():
             quant_method = getattr(module, "quant_method", None)
-            if quant_method is not None and isinstance(quant_method, QuantizeMethodBase):
+            if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
+                continue
+
+            if has_online_quant:
+                # Online quant may leave straggler params on the ``meta`` 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().
+                # device_loading_context — a blanket module.to(target_device) would
+                # raise NotImplementedError on meta params. Online quant initializes
+                # on the accelerator, so params are normally already on the target
+                # device and this loop is a no-op move; the point is to skip meta.
                 original_devices: dict[str, torch.device] = {}
                 for name, param in module.named_parameters():
                     if param.device.type != "meta" and param.device != target_device:
@@ -456,6 +464,22 @@ class DiffusersPipelineLoader:
                 for name, param in module.named_parameters():
                     if name in original_devices:
                         param.data = param.data.to(original_devices[name])
+            else:
+                # No meta params possible here. Preserve the original FSDP/HSDP-aware
+                # whole-module move (module.to()), which correctly handles sharded
+                # DTensor params that per-parameter .data reassignment cannot.
+                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)
+
+                quant_method.process_weights_after_loading(module)
+
+                if needs_device_move:
+                    module.to(module_device)
 
     def load_weights(self, model: nn.Module) -> None:
         weights_to_load = self._get_expected_parameter_names(model)