diff --git a/docs/user_guide/diffusion/lora.md b/docs/user_guide/diffusion/lora.md index 256698752..44705129f 100644 --- a/docs/user_guide/diffusion/lora.md +++ b/docs/user_guide/diffusion/lora.md @@ -1,14 +1,18 @@ # LoRA (Low-Rank Adaptation) Guide -LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni currently supports PEFT-style LoRA adapters, allowing you to customize model behavior without modifying the base model weights. +LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni supports two LoRA backends: **PEFT** for PEFT-style adapters, and **Distill** usually for few-steps inference. **PEFT** backend allows you to customize model behavior without modifying the base model weights. **Distill** backend fuses LoRA weights into the base model at initialization. ## Overview -LoRA adapters are lightweight, model-specific fine-tuning weights that can be dynamically loaded and applied to diffusion models. vLLM-Omni uses a unified LoRA handling mechanism similar to vLLM with LRU cache management. +LoRA adapters are lightweight, model-specific fine-tuning weights that can be applied to diffusion models in two ways: + +- **PEFT backend** (`--lora-backend peft`, default): Loads a PEFT-format adapter folder via `DiffusionLoRAManager`. Adapters are cached (LRU) and activated per request via `LoRARequest`. It uses a unified LoRA handling mechanisms similar to vLLM with LRU cache management. +- **Distill backend** (`--lora-backend distill`): Calls `pipeline.load_lora_weights` once at initialization to fuse one or more concrete checkpoint files directly into the base weights. Typically used for distilled few-step LoRAs (e.g. Lightning, LightX2V). ## LoRA Adapter Format -LoRA adapters must be in **PEFT (Parameter-Efficient Fine-Tuning)** format. A typical LoRA adapter directory structure: +### PEFT (Parameter-Efficient Fine-Tuning) format (default) +A typical PEFT-format LoRA adapter directory structure: ``` lora_adapter/ @@ -25,19 +29,20 @@ The `adapter_config.json` file contains metadata about the LoRA adapter, includi ### Offline Inference -#### Pre-loaded LoRA +#### PEFT backend: pre-loaded LoRA -Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests: +Load a PEFT-format LoRA adapter at initialization. The adapter is pre-loaded into the cache and can be activated per request: ```python from vllm_omni import Omni from vllm_omni.lora.request import LoRARequest -lora_path="/path/to/lora_adapter" +lora_path = "/path/to/lora_adapter" omni = Omni( model="stabilityai/stable-diffusion-3.5-medium", - lora_path=lora_path + lora_path=lora_path, + lora_backend="peft", # default, can be omitted ) lora_request = LoRARequest( @@ -53,6 +58,98 @@ outputs = omni.generate( ) ``` +#### Distill backend: fuse distilled LoRA at init + +For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or more concrete `.safetensors` files. The weights are fused into the base model once at init; subsequent `generate()` calls do not need a `LoRARequest`. + +##### Supported pipelines + +For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization. + +| Pipeline | Supported distilled LoRA repo | Notes | +|----------|--------------------------------|-------| +| `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. | +| `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. | +| `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. | + +Other diffusion pipelines are not currently listed as supporting distilled LoRA. Use the PEFT backend for request-time adapters, or bake converted weights into a local Diffusers directory before serving. + +Single-file example (Qwen-Image-Lightning): + +```python +from vllm_omni import Omni + +omni = Omni( + model="Qwen/Qwen-Image-2512", + lora_path="/path/to/Qwen-Image-2512-Lightning.safetensors", + lora_backend="distill", +) + +outputs = omni.generate(prompt="A piece of cheesecake") +``` + +Multi-file example (Wan2.2 MoE, high + low noise): + +```python +from vllm_omni import Omni + +omni = Omni( + model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + lora_path=[ + "/path/to/wan2.2_high_noise_lora.safetensors", # -> transformer + "/path/to/wan2.2_low_noise_lora.safetensors", # -> transformer_2 + ], + lora_backend="distill", +) +``` + +The CLI examples under `examples/offline_inference/` accept the same flags, e.g.: + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ + --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --lora-backend distill \ + --lora-path /path/to/high.safetensors /path/to/low.safetensors \ + --prompt "A cat playing with yarn" +``` + +##### Online serving + +Distilled LoRAs can also be fused when an online diffusion server starts. The +adapter remains active for the lifetime of that server, so requests do not pass +a per-request `lora` field. + +```bash +vllm serve Qwen/Qwen-Image-2512 \ + --omni \ + --port 8091 \ + --lora-backend distill \ + --lora-path /path/to/Qwen-Image-2512-Lightning-4steps.safetensors +``` + +Then send a normal generation request with the sampling settings expected by +the distilled checkpoint: + +```bash +curl -s http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + {"role": "user", "content": "A Chinese female college student, around 20 years old, with a very short haircut that conveys a gentle, artistic vibe. Her hair naturally falls to partially cover her cheeks, projecting a tomboyish yet charming demeanor. She has cool-toned fair skin and delicate features, with a slightly shy yet subtly confident expression—her mouth crooked in a playful, youthful smirk. She wears an off-shoulder top, revealing one shoulder, with a well-proportioned figure. The image is framed as a close-up selfie: she dominates the foreground, while the background clearly shows her dormitory—a neatly made bed with white linens on the top bunk, a tidy study desk with organized stationery, and wooden cabinets and drawers. The photo is captured on a smartphone under soft, even ambient lighting, with natural tones, high clarity, and a bright, lively atmosphere full of youthful, everyday energy."} + ], + "extra_body": { + "height": 1024, + "width": 1024, + "num_inference_steps": 4, + "true_cfg_scale": 1.0, + "seed": 42 + } + }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > qwen_image_distill_lora_output.png +``` + +For Wan2.2 MoE serving, pass the high-noise checkpoint first and the low-noise +checkpoint second to `--lora-path`. + !!! note "Server-side Path Requirement" The LoRA adapter path (`local_path`) must be readable on the **server** machine. If your client and server are on different machines, ensure the LoRA adapter is accessible via a shared mount or copied to the server. @@ -141,6 +238,7 @@ Notes: - This route avoids runtime LoRA loading changes in vLLM-Omni when you choose to bake converted weights into a local Diffusers directory. - Output quality and speed depend on the replacement checkpoints and sampling params you choose. +- If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path .safetensors .safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above. ## See Also diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md index f1cdb9293..f631cd962 100644 --- a/docs/user_guide/examples/offline_inference/text_to_image.md +++ b/docs/user_guide/examples/offline_inference/text_to_image.md @@ -88,8 +88,9 @@ python text_to_image.py \ | `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization | | `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel | | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models | -| `--lora-path` | str | — | Path to PEFT LoRA adapter folder | +| `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file | | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights | +| `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill | **NextStep-1.1 specific arguments:** diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index fcd3e7337..d32a1f562 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -125,6 +125,9 @@ python text_to_video.py \ - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio. - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF). - `--flow-shift`: scheduler flow_shift parameter. +- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file. +- `--lora-scale`: scale factor for LoRA weights. +- `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill. ### Wan2.2-specific diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index b2f27bb5b..bae4e390f 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -60,6 +60,8 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.lora.request import LoRARequest +from vllm_omni.lora.utils import stable_lora_int_id from vllm_omni.model_extras import ( build_image_to_video_prompt, get_extra_body_params, @@ -294,6 +296,30 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) + parser.add_argument( + "--lora-path", + type=str, + nargs="+", + default=None, + help="Path to a LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. " + "For Wan2.2 MoE models, pass the high-noise checkpoint first and the low-noise checkpoint second.", + ) + parser.add_argument( + "--lora-scale", + type=float, + default=1.0, + help="Scale factor for PEFT LoRA weights (default: 1.0).", + ) + parser.add_argument( + "--lora-backend", + type=str, + default="peft", + choices=["peft", "distill"], + help=( + "LoRA backend. 'peft' loads a PEFT-format adapter for request-time activation; " + "'distill' fuses one or more concrete LoRA checkpoint files at initialization." + ), + ) parser.add_argument( "--extra-body", type=functools.partial(parse_json_object, flag_name="--extra-body"), @@ -487,6 +513,12 @@ def main(): omni_kwargs["flow_shift"] = flow_shift if args.quantization is not None: omni_kwargs["quantization"] = args.quantization + if args.lora_path is not None: + lora_path = args.lora_path + if len(lora_path) == 1: + lora_path = lora_path[0] + omni_kwargs["lora_path"] = lora_path + omni_kwargs["lora_backend"] = args.lora_backend # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails # gate is an engine-level config (offline analog of the server's --no-guardrails). if args.extra_body and "guardrails" in args.extra_body: @@ -517,6 +549,18 @@ def main(): print(f" Video size: {width}x{height}") print(f"{'=' * 60}\n") + lora_request = None + if args.lora_path and args.lora_backend == "peft": + if len(args.lora_path) != 1: + raise ValueError("Only one LoRA path is expected for PEFT backend.") + + lora_path = args.lora_path[0] + lora_request = LoRARequest( + lora_name=Path(lora_path).stem, + lora_int_id=stable_lora_int_id(lora_path), + lora_path=lora_path, + ) + negative_prompt = args.negative_prompt if negative_prompt is None and not is_ltx2: # Preserve the historical empty-prompt behavior for non-LTX examples. @@ -530,6 +574,11 @@ def main(): width=width, num_frames=num_frames, ) + sampling_extra_args = {"sample_solver": args.sample_solver} + if lora_request is not None: + sampling_extra_args["lora_request"] = lora_request + sampling_extra_args["lora_scale"] = args.lora_scale + sampling_params = OmniDiffusionSamplingParams( height=height, width=width, @@ -540,7 +589,7 @@ def main(): num_inference_steps=num_inference_steps, num_frames=num_frames, frame_rate=frame_rate, - extra_args={"sample_solver": args.sample_solver}, + extra_args=sampling_extra_args, ) if flow_shift is not None: sampling_params.extra_args["flow_shift"] = flow_shift diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index 4485e21c3..294275cbe 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -109,8 +109,9 @@ python text_to_image.py \ | `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference | | `--ulysses-mode` | str | `"strict"` | Ulysses SP mode: `"strict"` or `"advanced_uaa"` | | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models | -| `--lora-path` | str | — | Path to PEFT LoRA adapter folder | +| `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file | | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights | +| `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill | | `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.| | `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.| | `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). | diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 03bdf7080..b2d3afe82 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -260,6 +260,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--lora-path", type=str, + nargs="+", default=None, help="Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.", ) @@ -269,6 +270,15 @@ def parse_args() -> argparse.Namespace: default=1.0, help="Scale factor for LoRA weights (default: 1.0).", ) + parser.add_argument( + "--lora-backend", + type=str, + default="peft", + choices=["peft", "distill"], + help="LoRA backend for loading LoRA adapters. Default: peft" + "'peft' loads a PEFT-format adapter folder, used e.g. for RL" + "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs", + ) parser.add_argument( "--vae-patch-parallel-size", type=int, @@ -404,8 +414,12 @@ def main(): # Prepare LoRA kwargs for Omni initialization lora_args: dict[str, Any] = {} if args.lora_path: - lora_args["lora_path"] = args.lora_path - print(f"Using LoRA from: {args.lora_path}") + lora_path = args.lora_path + if len(lora_path) == 1: + lora_path = lora_path[0] + lora_args["lora_path"] = lora_path + lora_args["lora_backend"] = args.lora_backend + print(f"Using LoRA from: {lora_path} with backend: {args.lora_backend}") # Build quantization kwargs: use quantization_config dict when # ignored_layers is specified so the list flows through OmniDiffusionConfig @@ -494,12 +508,16 @@ def main(): # Build LoRA request when --lora-path is set lora_request = None - if args.lora_path: - lora_request_id = stable_lora_int_id(args.lora_path) + if args.lora_path and args.lora_backend == "peft": + if len(args.lora_path) != 1: + raise ValueError("Only one LoRA path is expected for PEFT backend.") + + lora_path = args.lora_path[0] + lora_request_id = stable_lora_int_id(lora_path) lora_request = LoRARequest( - lora_name=Path(args.lora_path).stem, + lora_name=Path(lora_path).stem, lora_int_id=lora_request_id, - lora_path=args.lora_path, + lora_path=lora_path, ) generation_start = time.perf_counter() diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index abc522db8..50b207526 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -218,6 +218,9 @@ python text_to_video.py \ - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio. - `--quantization`: quantization method (such as `fp8` for FP8). - `--flow-shift`: scheduler flow_shift parameter. +- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file. +- `--lora-scale`: scale factor for LoRA weights. +- `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill. - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above). ### Wan2.2-specific diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index b21ec8f48..383887f5a 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -14,6 +14,8 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.lora.request import LoRARequest +from vllm_omni.lora.utils import stable_lora_int_id from vllm_omni.model_extras import get_extra_body_params, get_model_class_name from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -334,6 +336,29 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable expert parallelism for MoE layers.", ) + parser.add_argument( + "--lora-path", + type=str, + nargs="+", + default=None, + help="Path to LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. Loaded at initialization and used for generation." + "Note: for Wan2.2 MoE models, two checkpoints positionally mapped to high-noise and low-noise modules.", + ) + parser.add_argument( + "--lora-scale", + type=float, + default=1.0, + help="Scale factor for LoRA weights (default: 1.0).", + ) + parser.add_argument( + "--lora-backend", + type=str, + default="peft", + choices=["peft", "distill"], + help="LoRA backend for loading LoRA adapters. Default: peft" + "'peft' loads a PEFT-format adapter folder, used e.g. for RL" + "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs", + ) parser.add_argument( "--use-hsdp", action="store_true", @@ -440,6 +465,12 @@ def main(): omni_kwargs["cache_backend"] = args.cache_backend omni_kwargs["cache_config"] = cache_config omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary + if args.lora_path is not None: + lora_path = args.lora_path + if len(lora_path) == 1: + lora_path = lora_path[0] + omni_kwargs["lora_path"] = lora_path + omni_kwargs["lora_backend"] = args.lora_backend # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails # gate is an engine-level config (offline analog of the server's --no-guardrails). @@ -469,6 +500,19 @@ def main(): print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") + lora_request = None + if args.lora_path and args.lora_backend == "peft": + if len(args.lora_path) != 1: + raise ValueError("Only one LoRA path is expected for PEFT backend.") + + lora_path = args.lora_path[0] + lora_request_id = stable_lora_int_id(lora_path) + lora_request = LoRARequest( + lora_name=Path(lora_path).stem, + lora_int_id=lora_request_id, + lora_path=lora_path, + ) + prompt_dict = {"prompt": args.prompt} if args.negative_prompt is not None: prompt_dict["negative_prompt"] = args.negative_prompt @@ -476,6 +520,11 @@ def main(): # Preserve the historical empty-prompt behavior for non-LTX examples. prompt_dict["negative_prompt"] = "" + extra_args = {} + if lora_request: + extra_args["lora_request"] = lora_request + extra_args["lora_scale"] = args.lora_scale + sampling_kwargs = dict( height=args.height, width=args.width, @@ -483,6 +532,7 @@ def main(): guidance_scale=args.guidance_scale, num_inference_steps=args.num_inference_steps, num_frames=args.num_frames, + extra_args=extra_args, ) if args.guidance_scale_high is not None: sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high diff --git a/examples/online_serving/text_to_image/README.md b/examples/online_serving/text_to_image/README.md index 17d377ea3..567aba05e 100644 --- a/examples/online_serving/text_to_image/README.md +++ b/examples/online_serving/text_to_image/README.md @@ -115,7 +115,23 @@ python gradio_demo.py ## LoRA -This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory). +This example supports both startup-time distilled LoRAs and request-time +PEFT-compatible adapters. The LoRA path must be readable on the **server** +machine (usually a local path or a mounted directory). + +### Start Server with a Distilled LoRA + +```bash +vllm serve Qwen/Qwen-Image-2512 \ + --omni \ + --port 8091 \ + --lora-backend distill \ + --lora-path /path/to/Qwen-Image-2512-Lightning-4steps.safetensors +``` + +The distilled LoRA is fused once during server initialization and applies to +every request. Use the checkpoint's few-step sampling settings, for example +`num_inference_steps=4` and `true_cfg_scale=1.0`. ### Using Python Client with LoRA diff --git a/tests/diffusion/lora/test_loader.py b/tests/diffusion/lora/test_loader.py new file mode 100644 index 000000000..d2214f940 --- /dev/null +++ b/tests/diffusion/lora/test_loader.py @@ -0,0 +1,438 @@ +from collections import defaultdict + +import pytest +import torch +import torch.nn as nn +from diffusers.loaders.lora_conversion_utils import ( + _convert_non_diffusers_qwen_lora_to_diffusers, +) +from pytest_mock import MockerFixture +from torch.testing import assert_close + +from vllm_omni.diffusion.lora.loader import ( + LoraLoaderMixin, + QwenImageLoraLoaderMixin, + WanLoraLoaderMixin, + _prepare_lora_delta, + _remap_state_dict_keys, +) + +HEAD_DIM = 24 +RANK = 2 +NUM_LAYERS = 10 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class DummyTransformerBlock(nn.Module): + def __init__( + self, + d_model: int, + bias: bool = False, + stacked_params: bool = False, + ): + super().__init__() + self.attn = nn.Module() + if stacked_params: + self.attn.to_qkv = nn.Linear(d_model * 3, d_model, bias=bias) + else: + self.attn.to_q = nn.Linear(d_model, d_model, bias=bias) + self.attn.to_k = nn.Linear(d_model, d_model, bias=bias) + self.attn.to_v = nn.Linear(d_model, d_model, bias=bias) + self.attn.to_out = nn.Linear(d_model, d_model, bias=bias) + + +class DummyTransformer(nn.Module): + def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False): + super().__init__() + + self.blocks = nn.ModuleList( + [DummyTransformerBlock(d_model, bias=bias, stacked_params=stacked_params) for _ in range(num_layers)] + ) + if stacked_params: + self.stacked_params_mapping = [ + (".attn.to_qkv", ".attn.to_q", "q"), + (".attn.to_qkv", ".attn.to_k", "k"), + (".attn.to_qkv", ".attn.to_v", "v"), + ] + + +class DummyPipeline(nn.Module, LoraLoaderMixin): + def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False): + super().__init__() + self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=stacked_params) + + +# ====================================================== +# helper functions for test +# ====================================================== +def _filter_parameters(named_params, mode: str = "all"): + if mode == "all": + return named_params + + filtered = [] + for name, param in named_params: + if name.endswith(mode): + filtered.append((name, param)) + + return filtered + + +def assert_params_equal(module, original_params, mode: str = "all"): + for name, param in _filter_parameters(module.named_parameters(), mode): + assert_close(param, original_params[name]) + + +def assert_params_not_equal(module, original_params, mode: str = "all"): + for name, param in _filter_parameters(module.named_parameters(), mode): + assert not torch.allclose(param, original_params[name]) + + +def make_dummy_lora_state_dict( + base_key, + weight_a: torch.Tensor | None = None, + weight_b: torch.Tensor | None = None, + weight_bias: torch.Tensor | None = None, + bias: bool = False, + lora_a_suffix: str = "lora_A.weight", + lora_b_suffix: str = "lora_B.weight", + lora_bias_suffix: str = "bias", +): + """ + Create a dummy LoRA state dict. + + Returns: + tuple: (state_dict, expected_delta) + """ + key_a = f"{base_key}.{lora_a_suffix}" + key_b = f"{base_key}.{lora_b_suffix}" + key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None + + if weight_a is None: + weight_a = torch.randn(RANK, HEAD_DIM) + if weight_b is None: + weight_b = torch.randn(HEAD_DIM, RANK) + + state_dict = {} + expected_delta = None + + if bias: + state_dict[key_a] = weight_a + state_dict[key_b] = weight_b + weight_bias = torch.randn(HEAD_DIM) + state_dict[key_bias] = weight_bias + expected_delta = weight_bias + else: + state_dict[key_a] = weight_a + state_dict[key_b] = weight_b + expected_delta = weight_b @ weight_a + + return state_dict, expected_delta + + +def make_lora_state_dict_for_module( + module, + lora_rank=RANK, + bias=False, + lora_a_suffix: str = "lora_A.weight", + lora_b_suffix: str = "lora_B.weight", + lora_bias_suffix: str = "bias", + prefix: str = "transformer", + remap_proj_out: bool = True, +): + param_to_weight_names = defaultdict(list) + if hasattr(module, "stacked_params_mapping"): + for param_name, weight_name, _ in module.stacked_params_mapping: + param_to_weight_names[param_name].append(weight_name) + + state_dict = {} + for name, param in module.named_parameters(prefix=prefix): + if name.endswith(".weight"): + base_key = name[: -len(".weight")] + if remap_proj_out and base_key.endswith(".to_out"): + base_key = base_key.replace(".to_out", ".to_out.0") + + d_in, d_out = param.shape + is_stacked_param = False + for param_name, weight_names in param_to_weight_names.items(): + if param_name not in name: + continue + is_stacked_param = True + target_in = d_in // len(weight_names) + for weight_name in weight_names: + lora_key = base_key.replace(param_name, weight_name) + weight_a = torch.randn(lora_rank, d_out) + weight_b = torch.randn(target_in, lora_rank) + state_dict[f"{lora_key}.{lora_a_suffix}"] = weight_a + state_dict[f"{lora_key}.{lora_b_suffix}"] = weight_b + + if is_stacked_param: + continue + + weight_a = torch.randn(lora_rank, d_out) + weight_b = torch.randn(d_in, lora_rank) + key_a = f"{base_key}.{lora_a_suffix}" + key_b = f"{base_key}.{lora_b_suffix}" + state_dict[key_a] = weight_a + state_dict[key_b] = weight_b + elif name.endswith(".bias") and bias: + base_key = name[: -len(".bias")] + if remap_proj_out and base_key.endswith(".to_out"): + base_key = base_key.replace(".to_out", ".to_out.0") + + d_bias = param.shape[0] + is_stacked_param = False + for param_name, weight_names in param_to_weight_names.items(): + if param_name not in name: + continue + is_stacked_param = True + target_bias = d_bias // len(weight_names) + for weight_name in weight_names: + lora_key = base_key.replace(param_name, weight_name) + weight_bias = torch.randn(target_bias) + state_dict[f"{lora_key}.{lora_bias_suffix}"] = weight_bias + + if is_stacked_param: + continue + + key_bias = f"{base_key}.{lora_bias_suffix}" + state_dict[key_bias] = torch.randn_like(param) + + return state_dict + + +# ====================================================== +# Test loader helper functions +# ====================================================== +def test_prepare_lora_delta_with_stacked_params(): + base_key = "transformer.attn.to_qkv" + + sd0, delta0 = make_dummy_lora_state_dict("transformer.attn.to_q") + sd1, delta1 = make_dummy_lora_state_dict("transformer.attn.to_k") + sd2, delta2 = make_dummy_lora_state_dict("transformer.attn.to_v") + + lora_state_dict = {} + lora_state_dict.update(sd0) + lora_state_dict.update(sd1) + lora_state_dict.update(sd2) + expected_delta = torch.concat([delta0, delta1, delta2]) + + stacked_mapping = {".attn.to_qkv": [".attn.to_q", ".attn.to_k", ".attn.to_v"]} + + actual_delta, used_keys = _prepare_lora_delta( + lora_state_dict, + base_key, + stacked_mapping, + ) + + # validate the delta + assert actual_delta is not None + assert_close(actual_delta, expected_delta) + + # validate the used keys + key_q_a = "transformer.attn.to_q.lora_A.weight" + key_q_b = "transformer.attn.to_q.lora_B.weight" + key_k_a = "transformer.attn.to_k.lora_A.weight" + key_k_b = "transformer.attn.to_k.lora_B.weight" + key_v_a = "transformer.attn.to_v.lora_A.weight" + key_v_b = "transformer.attn.to_v.lora_B.weight" + + for key in [key_q_a, key_q_b, key_k_a, key_k_b, key_v_a, key_v_b]: + assert key in used_keys + + +def test_remap_state_dict_keys_matched_two(): + lora_state_dict, _ = make_dummy_lora_state_dict( + "transformer.attn.to_out.0", bias=True, lora_bias_suffix="lora_B.bias" + ) + + lora_state_dict = _remap_state_dict_keys( + lora_state_dict, + [(".to_out.0.", ".to_out."), (".lora_B.bias", ".bias")], + ) + + assert len(lora_state_dict) == 3 + assert "transformer.attn.to_out.bias" in lora_state_dict + + +# ====================================================== +# Test LoraLoaderMixin +# ====================================================== +class TestLoraLoaderMixin: + @pytest.mark.parametrize( + "bias, stacked_params", + [ + (False, False), + (True, False), + (False, True), + (True, True), + ], + ids=[ + "basic", + "bias", + "stacked_params", + "bias_stacked_params", + ], + ) + def test_load_lora_into_module(self, bias: bool, stacked_params: bool): + pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM, bias=bias, stacked_params=stacked_params) + original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()} + + lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer, bias=bias) + assert len(lora_state_dict) > 0 + + # remap the keys to vllm-omni format + lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")]) + + used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer) + # validate the parameters are updated + assert_params_not_equal(pipeline.transformer, original_params) + + # validate all keys are used + missing_keys = set(lora_state_dict.keys()) - used_keys + assert len(missing_keys) == 0 + + def test_unload_module_lora(self): + pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM) + original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()} + + lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer) + # validate the state dict is valid + assert len(lora_state_dict) > 0 + + # remap the keys to vllm-omni format + lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")]) + + used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer) + + # validate the weights are updated + assert_params_not_equal(pipeline.transformer, original_params) + + # validate all keys are used + missing_keys = set(lora_state_dict.keys()) - used_keys + assert len(missing_keys) == 0 + + unload_used_keys = pipeline.unload_module_lora(lora_state_dict, pipeline.transformer) + missing_keys = set(lora_state_dict.keys()) - unload_used_keys + assert len(missing_keys) == 0 + assert len(unload_used_keys - used_keys) == 0 + + # validate the weights are restored + assert_params_equal(pipeline.transformer, original_params) + + +# ====================================================== +# Test QwenImageLoraLoaderMixin +# ====================================================== +class DummyQwenImagePipeline(nn.Module, QwenImageLoraLoaderMixin): + def __init__(self, num_layers: int, d_model: int, bias: bool = False): + super().__init__() + self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=True) + + +class TestQwenImageLoraLoaderMixin: + def test_load_lora_weights(self, mocker: MockerFixture): + pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM) + original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()} + + lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer) + assert len(lora_state_dict) > 0 + + mocker.patch( + "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline", + return_value=_convert_non_diffusers_qwen_lora_to_diffusers, + ) + pipeline.load_lora_weights(lora_state_dict, "adapter0") + + # validate the weights are updated after lora loaded + assert_params_not_equal(pipeline.transformer, original_params) + + # validate lora_loaded map is updated after lora loaded + assert "adapter0" in pipeline.lora_loaded + + def test_unload_lora_weights(self, mocker: MockerFixture): + pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM) + original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()} + + lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer) + assert len(lora_state_dict) > 0 + + mocker.patch( + "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline", + return_value=_convert_non_diffusers_qwen_lora_to_diffusers, + ) + pipeline.load_lora_weights(lora_state_dict, "adapter0") + # validate the weights are updated after lora loaded + assert_params_not_equal(pipeline.transformer, original_params) + + # validate lora_loaded map is updated after lora loaded + assert "adapter0" in pipeline.lora_loaded + + pipeline.unload_lora_weights("adapter0") + # validate the weights are restored after lora unloaded + assert_params_equal(pipeline.transformer, original_params) + + # validate lora_loaded map is updated + assert "adapter0" not in pipeline.lora_loaded + + def test_load_diffusers_lora_applies_non_unit_alpha(self, mocker: MockerFixture): + pipeline = DummyQwenImagePipeline(1, HEAD_DIM) + to_out = pipeline.transformer.blocks[0].attn.to_out + original_weight = to_out.weight.detach().clone() + base_key = "blocks.0.attn.to_out.0" + lora_a = torch.randn(RANK, HEAD_DIM) + lora_b = torch.randn(HEAD_DIM, RANK) + alpha = torch.tensor(6.0) + state_dict = { + f"{base_key}.lora_A.weight": lora_a, + f"{base_key}.lora_B.weight": lora_b, + f"{base_key}.alpha": alpha, + } + mocker.patch( + "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline", + return_value=_convert_non_diffusers_qwen_lora_to_diffusers, + ) + + pipeline.load_lora_weights(state_dict, "adapter0") + + expected_weight = original_weight + (lora_b @ lora_a) * (alpha.item() / RANK) + assert_close(to_out.weight, expected_weight) + + pipeline.unload_lora_weights("adapter0") + assert_close(to_out.weight, original_weight) + + +# ====================================================== +# Test WanLoraLoaderMixin +# ====================================================== +class DummyWanPipeline(nn.Module, WanLoraLoaderMixin): + def __init__(self, num_layers: int, d_model: int): + super().__init__() + self.has_transformer_2 = True + self.transformer = DummyTransformer(num_layers, d_model, stacked_params=True) + self.transformer_2 = DummyTransformer(num_layers, d_model, stacked_params=True) + + +def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture): + pipeline = DummyWanPipeline(NUM_LAYERS, HEAD_DIM) + original_high = {name: param.clone() for name, param in pipeline.transformer.named_parameters()} + original_low = {name: param.clone() for name, param in pipeline.transformer_2.named_parameters()} + high_state_dict = make_lora_state_dict_for_module(pipeline.transformer) + low_state_dict = make_lora_state_dict_for_module(pipeline.transformer_2) + + mocker.patch( + "vllm_omni.diffusion.lora.loader._load_lora_state_dict", + side_effect=[high_state_dict, low_state_dict], + ) + + pipeline.load_lora_weights(["high.safetensors", "low.safetensors"], "adapter0") + + assert_params_not_equal(pipeline.transformer, original_high) + assert_params_not_equal(pipeline.transformer_2, original_low) + assert set(pipeline.lora_loaded["adapter0"]) == {"transformer", "transformer_2"} + + pipeline.unload_lora_weights("adapter0") + + assert_params_equal(pipeline.transformer, original_high) + assert_params_equal(pipeline.transformer_2, original_low) + assert "adapter0" not in pipeline.lora_loaded diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 64b586bf2..2fa8c92bd 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -284,6 +284,37 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True +def test_serve_cli_forwards_distilled_lora_to_diffusion_stage(): + """Ensure startup distilled LoRA options reach the online diffusion stage.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + "--omni", + "--lora-backend", + "distill", + "--lora-path", + "/models/high.safetensors", + "/models/low.safetensors", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + engine_args = stage_cfg["engine_args"] + + assert explicit_kwargs["lora_backend"] == "distill" + assert engine_args["lora_backend"] == "distill" + assert engine_args["lora_path"] == [ + "/models/high.safetensors", + "/models/low.safetensors", + ] + + def test_serve_cli_forwards_distributed_offload_residency(): """Ensure the two-GPU DLO placement controls reach the diffusion stage.""" parser = TrackingArgumentParser() diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 23892f640..5ae83d1b9 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -563,8 +563,9 @@ class _DiffusionConfigProjection: diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict) diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict) diffusers_pipeline_cls: Any = None - lora_path: str | None = None + lora_path: str | list[str] | None = None lora_scale: float = 1.0 + lora_backend: str = "peft" max_cpu_loras: int | None = None output_type: str = "pil" enable_cpu_offload: bool = False diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5b6b9a01a..55c4a46cc 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -385,6 +385,9 @@ class StageDeployConfig: # Diffusion model loading and adapter construction. model_class_name: str | None = None diffusion_load_format: str | None = None + lora_path: str | list[str] | None = None + lora_backend: str | None = None + lora_scale: float | None = None diffusers_load_kwargs: dict[str, Any] | None = None diffusers_call_kwargs: dict[str, Any] | None = None diffusion_quantization_config: str | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 73051692d..137fb9860 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -27,6 +27,7 @@ from vllm_omni.diffusion.diffusion_kv.config import ( DiffusionKVCacheMode, parse_diffusion_kv_cache_mode, ) +from vllm_omni.diffusion.lora.manager import LoRABackend from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.network_utils import is_port_available from vllm_omni.errors import client_error_metadata @@ -682,8 +683,9 @@ class OmniDiffusionConfig: # pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False) # LoRA parameters - lora_path: str | None = None - lora_scale: float = 1.0 + lora_path: str | list[str] | None = None + lora_scale: float | list[float] = 1.0 + lora_backend: LoRABackend = LoRABackend.PEFT # available choices: ["peft", "distill"] max_cpu_loras: int | None = None output_type: str = "pil" diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py new file mode 100644 index 000000000..9a9fbfae5 --- /dev/null +++ b/vllm_omni/diffusion/lora/loader.py @@ -0,0 +1,461 @@ +import os +from collections import OrderedDict, defaultdict +from collections.abc import Callable + +import torch +from diffusers.loaders.lora_conversion_utils import ( + _convert_non_diffusers_qwen_lora_to_diffusers, + _convert_non_diffusers_wan_lora_to_diffusers, +) +from huggingface_hub import hf_hub_download +from safetensors.torch import load_file +from vllm.logger import init_logger + +from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline + +logger = init_logger(__name__) + +lora_convert_mapping: dict[str, Callable] = { + "QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers, + "QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers, + "QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers, + "Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers, + "Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers, +} + + +def get_converter_by_pipeline(pipeline): + for pipeline_cls in pipeline.__class__.__mro__: + converter = lora_convert_mapping.get(pipeline_cls.__name__) + if converter is not None: + return converter + return None + + +def _prepare_lora_delta( + lora_state_dict, + base_key: str, + param_to_weight_names: dict[str, list[str]] | None = None, + is_bias: bool = False, + lora_a_suffix: str = "lora_A.weight", + lora_b_suffix: str = "lora_B.weight", + lora_bias_suffix: str = "bias", +): + used_keys = set() + # stacked_params_mapping param_to_weight_names + # [(".to_qkv", ".to_q.", "q") + # (".to_qkv", ".to_k.", "k") ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]} + # (".to_qkv", ".to_v.", "v")] + + # stacked_sd is to store packed parameter deltas (format: [str, list[torch.Tensor]]) + # example: {".to_qkv": [delta_q, delta_k, delta_v]} + stacked_deltas = [] + is_stacked_param = False + if param_to_weight_names is None: + param_to_weight_names = defaultdict(list) + for param_name, weight_names in param_to_weight_names.items(): + if param_name not in base_key: + continue + is_stacked_param = True + # handle lora_a_key and lora_b_key together + if is_bias: + for weight_name in weight_names: + lora_bias_key = f"{base_key.replace(param_name, weight_name)}.{lora_bias_suffix}" + if lora_bias_key not in lora_state_dict: + return None, used_keys + delta = lora_state_dict[lora_bias_key] + stacked_deltas.append(delta) + used_keys.add(lora_bias_key) + else: + for weight_name in weight_names: + lora_a_key = f"{base_key.replace(param_name, weight_name)}.{lora_a_suffix}" + lora_b_key = lora_a_key.replace(lora_a_suffix, lora_b_suffix) + if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict: + return None, used_keys + a = lora_state_dict[lora_a_key] + b = lora_state_dict[lora_b_key] + delta = torch.matmul(b, a) + stacked_deltas.append(delta) + used_keys.add(lora_a_key) + used_keys.add(lora_b_key) + continue + + if is_stacked_param: + return torch.concat(stacked_deltas), used_keys + + if is_bias: + lora_bias_key = f"{base_key}.{lora_bias_suffix}" + if lora_bias_key not in lora_state_dict: + return None, used_keys + used_keys.add(lora_bias_key) + return lora_state_dict[lora_bias_key], used_keys + + lora_a_key = f"{base_key}.{lora_a_suffix}" + lora_b_key = f"{base_key}.{lora_b_suffix}" + if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict: + return None, used_keys + a = lora_state_dict[lora_a_key] + b = lora_state_dict[lora_b_key] + used_keys.add(lora_a_key) + used_keys.add(lora_b_key) + return torch.matmul(b, a), used_keys + + +def _load_lora_state_dict( + pretrained_model_name_or_path: str, + weights_name: str | None = None, + use_safetensors: bool = True, + subfolder: str | None = None, + cache_dir: str | None = None, + force_download: bool = False, + local_files_only: bool | None = None, +): + # first we try to load it from local disk + if use_safetensors and pretrained_model_name_or_path.endswith(".safetensors"): + return load_file(pretrained_model_name_or_path) + + if os.path.isdir(pretrained_model_name_or_path): + model_file = pretrained_model_name_or_path + if subfolder: + model_file = os.path.join(model_file, subfolder) + if weights_name is None: + raise ValueError("weights_name is required when loading from a directory") + model_file = os.path.join(model_file, weights_name) + return load_file(model_file) + + # finally, we try to load it from the internet + try: + model_file = hf_hub_download( + pretrained_model_name_or_path, + filename=weights_name, + subfolder=subfolder, + cache_dir=cache_dir, + force_download=force_download, + local_files_only=local_files_only, + ) + return load_file(model_file) + except Exception as e: + logger.error(f"Failed to download {weights_name} from {pretrained_model_name_or_path}: {e}") + raise e + + +def _remap_state_dict_keys(sd, rules): + """ + Remap keys in a state_dict by sequentially applying a list of substring substitution rules. + + This utility function transforms parameter names by applying a sequence of + (old_substring, new_substring) pairs. For each key in the input dictionary, + **all rules are evaluated in the given order**, and every rule whose + `old_substring` is found within the current key triggers a replacement. + Substitutions are cumulative: the key is updated after each match, and + subsequent rules operate on the result of previous replacements. + + The order of rules matters; users are responsible for arranging them to + achieve the desired final key. Keys that do not match any rule are copied + unchanged. + + Args: + sd (dict[str, Any]): Original state_dict mapping parameter names to tensors. + rules (List[Tuple[str, str]]): A list of (old_substring, new_substring) pairs. + Rules are applied sequentially, and **all matching rules are executed**. + + Returns: + dict[str, Any]: A new state_dict with remapped keys. Values are unchanged. + """ + new_sd = OrderedDict() + for k, v in sd.items(): + new_key = k + matched = False + for p1, p2 in rules: + if p1 not in k: + continue + new_key = new_key.replace(p1, p2) + matched = True + if matched: + new_sd[new_key] = v + else: + new_sd[k] = v + return new_sd + + +def _apply_diffusers_lora_alpha_scaling(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Fold Diffusers-format LoRA alpha values into their B matrices.""" + scaled_state_dict = dict(state_dict) + for alpha_key in (key for key in state_dict if key.endswith(".alpha")): + base_key = alpha_key[: -len(".alpha")] + lora_a_key = f"{base_key}.lora_A.weight" + lora_b_key = f"{base_key}.lora_B.weight" + if lora_a_key not in state_dict or lora_b_key not in state_dict: + raise ValueError(f"LoRA alpha key {alpha_key!r} does not have matching lora_A and lora_B weights.") + + rank = state_dict[lora_a_key].shape[0] + alpha = state_dict[alpha_key] + + scaled_state_dict[lora_b_key] = state_dict[lora_b_key] * (alpha.item() / rank) + scaled_state_dict.pop(alpha_key) + + return scaled_state_dict + + +class LoraLoaderMixin: + transformer_name = "transformer" + + # Lazy initialization to avoid MRO issues: __init__ may not be called + # when mixin with nn.Module + @property + def lora_loaded(self): + if not hasattr(self, "_lora_loaded"): + self._lora_loaded = {} + return self._lora_loaded + + @lora_loaded.setter + def lora_loaded(self, value): + self._lora_loaded = value + + @classmethod + def load_lora_into_module( + cls, + state_dict, + module, + prefix: str = "transformer", + lora_a_suffix: str = "lora_A.weight", + lora_b_suffix: str = "lora_B.weight", + lora_bias_suffix: str = "bias", + ): + lora_loaded_keys = set() + + # prepare stacked parameters mapping for later use + # stacked_params_mapping param_to_weight_names + # [(".to_qkv", ".to_q.", "q") + # (".to_qkv", ".to_k.", "k") ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]} + # (".to_qkv", ".to_v.", "v")] + param_to_weight_names = defaultdict(list) + if hasattr(module, "stacked_params_mapping"): + for param_name, weight_name, _ in module.stacked_params_mapping: + param_to_weight_names[param_name].append(weight_name) + + for name, params in module.named_parameters(prefix): + is_bias = False + if name.endswith(".bias"): + base_key = name[: -len(".bias")] + is_bias = True + elif name.endswith(".weight"): + base_key = name[: -len(".weight")] + else: + continue + + delta, used_keys = _prepare_lora_delta( + state_dict, + base_key, + param_to_weight_names, + is_bias, + lora_a_suffix, + lora_b_suffix, + lora_bias_suffix, + ) + if delta is None: + continue + lora_loaded_keys.update(used_keys) + + delta = delta.to(device=params.device, dtype=params.dtype) + with torch.no_grad(): + params.add_(delta) + del delta + + missing_keys = set(state_dict.keys()) - lora_loaded_keys + for k in missing_keys: + logger.warning(f"Missing loading lora key: {k}") + + logger.info(f"{len(lora_loaded_keys)} lora keys loaded into {module.__class__.__name__}.") + return lora_loaded_keys + + @classmethod + def unload_module_lora( + cls, + state_dict, + module, + prefix: str = "transformer", + lora_a_suffix: str = "lora_A.weight", + lora_b_suffix: str = "lora_B.weight", + lora_bias_suffix: str = "bias", + ): + lora_unloaded_keys = set() + + param_to_weight_names = defaultdict(list) + if hasattr(module, "stacked_params_mapping"): + for param_name, weight_name, _ in module.stacked_params_mapping: + param_to_weight_names[param_name].append(weight_name) + + for name, param in module.named_parameters(prefix): + is_bias = False + if name.endswith(".bias"): + base_key = name[: -len(".bias")] + is_bias = True + elif name.endswith(".weight"): + base_key = name[: -len(".weight")] + else: + continue + + delta, used_keys = _prepare_lora_delta( + state_dict, + base_key, + param_to_weight_names, + is_bias, + lora_a_suffix, + lora_b_suffix, + lora_bias_suffix, + ) + if delta is None: + continue + lora_unloaded_keys.update(used_keys) + + delta = delta.to(device=param.device, dtype=param.dtype) + with torch.no_grad(): + param.sub_(delta) + del delta + + missing_keys = set(state_dict.keys()) - lora_unloaded_keys + for k in missing_keys: + logger.warning(f"Missing unloading lora key: {k}") + + logger.info(f"{len(lora_unloaded_keys)} lora keys unloaded from {module.__class__.__name__}.") + return lora_unloaded_keys + + +class QwenImageLoraLoaderMixin(LoraLoaderMixin): + def load_lora_weights( + self, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, + ): + if adapter_name in self.lora_loaded: + return + + if isinstance(pretrained_model_name_or_path_or_dict, dict): + state_dict = pretrained_model_name_or_path_or_dict + else: + state_dict = _load_lora_state_dict(pretrained_model_name_or_path_or_dict) + + has_alpha = any(k.endswith(".alpha") for k in state_dict) + is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict) + is_diffusers_lora_format = any(k.endswith(".lora_A.weight") for k in state_dict) + if has_alpha and is_diffusers_lora_format: + state_dict = _apply_diffusers_lora_alpha_scaling(state_dict) + + if has_alpha or is_non_diffusers_format: + converter = get_converter_by_pipeline(self) + if converter is None: + raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}") + + state_dict = converter(state_dict) + + state_dict = _remap_state_dict_keys(state_dict, [(".to_out.0.", ".to_out.")]) + + self.load_lora_into_module( + state_dict, + self.transformer, + prefix=self.transformer_name, + ) + + self.lora_loaded[adapter_name] = state_dict + + def unload_lora_weights(self, adapter_name: str): + if adapter_name not in self.lora_loaded: + return + + state_dict = self.lora_loaded[adapter_name] + + transformer = get_transformer_from_pipeline(self) + self.unload_module_lora(state_dict, transformer, prefix=self.transformer_name) + + del self.lora_loaded[adapter_name] + + +class WanLoraLoaderMixin(LoraLoaderMixin): + transformer_2_name = "transformer_2" + + def load_lora_weights( + self, + pretrained_model_name_or_path: str | list[str], + adapter_name: str | None = None, + ): + """Load LoRA weights into the pipeline's transformer(s). + + Args: + pretrained_model_name_or_path: Path(s) to concrete .safetensors LoRA files. + - str: a single .safetensors file. Only valid for Wan2.1 + (single-transformer pipelines). + - list[str]: one file per transformer, matched **by position** to + [transformer, transformer_2]. Required for Wan2.2 MoE. + adapter_name: Name to register the adapter under. + """ + if adapter_name in self.lora_loaded: + return + + # Normalize to list[(lora_path, target_module_name)]. + lora_paths = None + if isinstance(pretrained_model_name_or_path, str): + lora_paths = [pretrained_model_name_or_path] + else: + lora_paths = list(pretrained_model_name_or_path) + + target_modules = ( + [self.transformer_name, self.transformer_2_name] if self.has_transformer_2 else [self.transformer_name] + ) + if len(lora_paths) != len(target_modules): + raise ValueError( + f"{self.__class__.__name__} expects {len(target_modules)} LoRA file(s) " + f"(target_modules={target_modules}, paths={lora_paths}), got {len(lora_paths)}. Pass a list of concrete " # noqa + f".safetensors files, one per transformer." + ) + + module_to_lora_sd = dict() + for lora_path, target_module_name in zip(lora_paths, target_modules): + state_dict = _load_lora_state_dict(lora_path) + is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict) + if is_non_diffusers_format: + converter = get_converter_by_pipeline(self) + if converter is None: + raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}") + + state_dict = converter(state_dict) + + state_dict = _remap_state_dict_keys( + state_dict, + [ + (".ffn.net.0.", ".ffn.net_0."), + (".ffn.net.2.", ".ffn.net_2."), + ( + ".to_out.0.", + ".to_out.", + ), + # '.bias_B.bias' is a horrible name in diffuser's. + # So we remap it to '.bias' + (".lora_B.bias", ".bias"), + ], + ) + + module = get_transformer_from_pipeline(self, transformer_name=target_module_name) + if module is None: + logger.warning( + f"Skip LoRA {lora_path}: target module '{target_module_name}' " + f"is not loaded in pipeline {self.__class__.__name__} " + f"(e.g. disabled by boundary_ratio)." + ) + continue + + self.load_lora_into_module(state_dict, module, prefix=self.transformer_name, lora_bias_suffix="bias") + module_to_lora_sd[target_module_name] = state_dict + + self.lora_loaded[adapter_name] = module_to_lora_sd + + def unload_lora_weights(self, adapter_name: str): + if adapter_name not in self.lora_loaded: + return + + module_to_lora_sd = self.lora_loaded[adapter_name] + for module_name, lora_sd in module_to_lora_sd.items(): + module = get_transformer_from_pipeline(self, transformer_name=module_name) + self.unload_module_lora(lora_sd, module, prefix=self.transformer_name) + + del self.lora_loaded[adapter_name] diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 9098fec2a..d11b1b9a2 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -4,6 +4,7 @@ import math import time from collections import OrderedDict +from enum import Enum from typing import get_args import torch @@ -33,6 +34,11 @@ from vllm_omni.lora.utils import stable_lora_int_id logger = init_logger(__name__) +class LoRABackend(str, Enum): + PEFT = "peft" + DISTILL = "distill" + + class DiffusionLoRAManager: """Manager for LoRA adapters in diffusion models. diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 0aeb8b85c..70e0ef0f1 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -25,6 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.lora.loader import QwenImageLoraLoaderMixin from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin @@ -264,7 +265,11 @@ def apply_rotary_emb_qwen( class QwenImagePipeline( - nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery + nn.Module, + QwenImageCFGParallelMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, + QwenImageLoraLoaderMixin, ): supports_request_batch = True _dit_modules: ClassVar[list[str]] = ["transformer"] diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 0fc9aca65..f9bc9f04f 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -25,6 +25,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import DenoiseProgressMixin +from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin @@ -285,6 +286,7 @@ class Wan22Pipeline( DenoiseProgressMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery, + WanLoraLoaderMixin, ): _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 180f8bb7a..a2e444bb0 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -25,6 +25,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import DenoiseProgressMixin +from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin @@ -147,6 +148,7 @@ class Wan22I2VPipeline( DiffusionPipelineProfilerMixin, DenoiseProgressMixin, SupportsComponentDiscovery, + WanLoraLoaderMixin, ): """ Wan2.2 Image-to-Video Pipeline. diff --git a/vllm_omni/diffusion/utils/tf_utils.py b/vllm_omni/diffusion/utils/tf_utils.py index 745410ec2..b79d0feb1 100644 --- a/vllm_omni/diffusion/utils/tf_utils.py +++ b/vllm_omni/diffusion/utils/tf_utils.py @@ -70,9 +70,10 @@ def find_module_with_attr(model, attr_name="transformer"): return None -def get_transformer_from_pipeline(pipeline: Any): - pipe = find_module_with_attr(pipeline, attr_name="transformer") +def get_transformer_from_pipeline(pipeline: Any, transformer_name: str = "transformer"): + """Return a named transformer from a pipeline or pipeline wrapper.""" + pipe = find_module_with_attr(pipeline, attr_name=transformer_name) if pipe is not None: - return pipe.transformer + return getattr(pipe, transformer_name) return None diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 2b927ce37..0820d82b4 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -49,7 +49,7 @@ from vllm_omni.diffusion.distributed.parallel_state import ( ) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm -from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager +from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob @@ -388,14 +388,31 @@ class DiffusionWorker: """Initialize the LoRA manager for this worker.""" if self.model_runner.pipeline is None: return - self.lora_manager = DiffusionLoRAManager( - pipeline=self.model_runner.pipeline, - device=self.device, - dtype=self.od_config.dtype, - max_cached_adapters=self.od_config.max_cpu_loras, - lora_path=self.od_config.lora_path, - lora_scale=self.od_config.lora_scale, - ) + + lora_path = self.od_config.lora_path + if isinstance(lora_path, list) and len(lora_path) == 1: + lora_path = lora_path[0] + + lora_backend = self.od_config.lora_backend + if lora_backend == LoRABackend.PEFT: + self.lora_manager = DiffusionLoRAManager( + pipeline=self.model_runner.pipeline, + device=self.device, + dtype=self.od_config.dtype, + max_cached_adapters=self.od_config.max_cpu_loras, + lora_path=lora_path, + lora_scale=self.od_config.lora_scale, + ) + elif lora_backend == LoRABackend.DISTILL: + pipeline = self.model_runner.pipeline + if hasattr(pipeline, "load_lora_weights"): + if self.od_config.lora_scale > 1.0: + logger.warning("lora_scale > 1.0 may not take any effect when using distilled LoRA backend.") + pipeline.load_lora_weights(lora_path) + else: + logger.warning("Pipeline does not support loading distilled LoRA weights for now.") + else: + raise ValueError(f"Unknown LoRA backend: {lora_backend}. Available choices: {LoRABackend.__members__}") def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: """Start or stop profiling for this GPU worker. diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 02ec3587b..a19beed80 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -471,6 +471,9 @@ class OrchestratorArgs: num_gpus: int | None = None model_class_name: str | None = None diffusion_load_format: str | None = None + lora_path: list[str] | None = None + lora_backend: str | None = None + lora_scale: float | None = None diffusers_load_kwargs: str = "{}" diffusers_call_kwargs: str = "{}" ulysses_degree: int | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 058510eda..e7ad79696 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1048,6 +1048,9 @@ class AsyncOmniEngine: "boundary_ratio": kwargs.get("boundary_ratio", None), "flow_shift": kwargs.get("flow_shift", None), "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), + "lora_path": kwargs.get("lora_path", None), + "lora_scale": kwargs.get("lora_scale", 1.0), + "lora_backend": kwargs.get("lora_backend", "peft"), "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), "worker_extension_cls": kwargs.get("worker_extension_cls", None), "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), @@ -1251,6 +1254,9 @@ class AsyncOmniEngine: if lora_scale is not None: if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None: cfg.engine_args.lora_scale = lora_scale + if kwargs.get("lora_backend") is not None: + if not hasattr(cfg.engine_args, "lora_backend") or cfg.engine_args.lora_backend is None: + cfg.engine_args.lora_backend = kwargs["lora_backend"] if ( kwargs.get("diffusion_attention_config") is not None or kwargs.get("diffusion_attention_backend") is not None diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 7fb9cb8cd..815ce4d6e 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -164,6 +164,13 @@ class OmniServeCommand(CLISubcommand): "(single-runtime) or `--omni-dp-size-local` (headless / multi-runtime)." ) + lora_backend = getattr(args, "lora_backend", None) + lora_path = getattr(args, "lora_path", None) + if lora_backend == "distill" and not lora_path: + raise ValueError("--lora-backend distill requires --lora-path") + if (lora_backend or "peft") == "peft" and isinstance(lora_path, list) and len(lora_path) > 1: + raise ValueError("--lora-backend peft accepts only one startup LoRA path") + # --omni-lb-policy is validated against the LoadBalancingPolicy enum. omni_lb_policy = getattr(args, "omni_lb_policy", None) if omni_lb_policy is not None: @@ -427,6 +434,32 @@ class OmniServeCommand(CLISubcommand): "custom_pipeline, dummy, or diffusers for the HF diffusers adapter." ), ) + omni_config_group.add_argument( + "--lora-path", + type=str, + nargs="+", + default=None, + help=( + "LoRA checkpoint path(s) loaded when the diffusion server starts. " + "The distill backend accepts one file per pipeline transformer." + ), + ) + omni_config_group.add_argument( + "--lora-backend", + type=str, + choices=["peft", "distill"], + default=None, + help=( + "Diffusion LoRA loading backend. 'distill' fuses checkpoint files " + "into the base model at server startup; 'peft' uses the adapter manager." + ), + ) + omni_config_group.add_argument( + "--lora-scale", + type=float, + default=None, + help="Scale for a startup PEFT LoRA. Distilled LoRAs are fused at their checkpoint scale.", + ) omni_config_group.add_argument( "--diffusion-compile-granularity", choices=["regional", "full"],