diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 78d9ffd03..daccef343 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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: diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 105629f34..accb173e1 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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() diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index b98679033..474834d16 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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),