[Bugfix] Fix cpu offload and quantization compatibility (#1473)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -256,8 +256,14 @@ class DiffusersPipelineLoader:
         load_device: str,
         load_format: str = "default",
         custom_pipeline_name: str | None = None,
+        device: torch.device | None = None,
     ) -> nn.Module:
         """Load a model with the given configurations."""
+        # CPU offload + FP8: load weights on device for FP8 quantization
+        if load_device == "cpu" and od_config.quantization and od_config.quantization.lower() != "none":
+            load_device = device.type
+            logger.info(f"Quantization enabled with CPU offload, using {load_device} for weight loading")
+
         target_device = torch.device(load_device)
         with set_default_torch_dtype(od_config.dtype):
             if od_config.parallel_config.use_hsdp:

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -127,6 +127,7 @@ class DiffusionModelRunner:
                     load_device=load_device,
                     load_format=load_format,
                     custom_pipeline_name=custom_pipeline_name,
+                    device=self.device,
                 )
         time_after_load = time.perf_counter()
 

vllm_omni/entrypoints/async_omni.py

@@ -225,6 +225,7 @@ class AsyncOmni(OmniBase):
                     "enforce_eager": kwargs.get("enforce_eager", False),
                     "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
                     "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
+                    "quantization": kwargs.get("quantization", None),
                     "worker_extension_cls": kwargs.get("worker_extension_cls", None),
                     "enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
                     "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),