From 5c72be1e51ef1716c69d0e349030bfeeb949f53d Mon Sep 17 00:00:00 2001 From: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Date: Sat, 10 Jan 2026 23:32:45 +0800 Subject: [PATCH] [diffusion] feat: support multiple LoRA adapters loading and application (#16667) Co-authored-by: niehen6174 Co-authored-by: Mick --- .../sglang/multimodal_gen/docs/openai_api.md | 49 ++- .../entrypoints/diffusion_generator.py | 29 +- .../runtime/entrypoints/openai/common_api.py | 31 +- .../runtime/entrypoints/openai/utils.py | 33 +- .../runtime/layers/lora/linear.py | 104 +++-- .../runtime/managers/gpu_worker.py | 21 +- .../runtime/pipelines_core/lora_pipeline.py | 364 ++++++++++++------ .../test/server/perf_baselines.json | 25 ++ .../test/server/test_server_common.py | 88 +++++ .../test/server/testcase_configs.py | 14 + 10 files changed, 578 insertions(+), 180 deletions(-) diff --git a/python/sglang/multimodal_gen/docs/openai_api.md b/python/sglang/multimodal_gen/docs/openai_api.md index c55fd0e34..88dabac4c 100644 --- a/python/sglang/multimodal_gen/docs/openai_api.md +++ b/python/sglang/multimodal_gen/docs/openai_api.md @@ -241,17 +241,21 @@ The server supports dynamic loading, merging, and unmerging of LoRA adapters. #### Set LoRA Adapter -Loads a LoRA adapter and merges its weights into the model. +Loads one or more LoRA adapters and merges their weights into the model. Supports both single LoRA (backward compatible) and multiple LoRA adapters. **Endpoint:** `POST /v1/set_lora` **Parameters:** -- `lora_nickname` (string, required): A unique identifier for this LoRA -- `lora_path` (string, optional): Path to the `.safetensors` file or Hugging Face repo ID. Required for the first load; optional if re-activating a cached nickname -- `target` (string, optional): Which transformer(s) to apply the LoRA to. One of "all" (default), "transformer", "transformer_2", "critic" -- `strength` (float, optional): LoRA strength for merge, default 1.0. Values < 1.0 reduce the effect, values > 1.0 amplify the effect +- `lora_nickname` (string or list of strings, required): A unique identifier for the LoRA adapter(s). Can be a single string or a list of strings for multiple LoRAs +- `lora_path` (string or list of strings/None, optional): Path to the `.safetensors` file(s) or Hugging Face repo ID(s). Required for the first load; optional if re-activating a cached nickname. If a list, must match the length of `lora_nickname` +- `target` (string or list of strings, optional): Which transformer(s) to apply the LoRA to. If a list, must match the length of `lora_nickname`. Valid values: + - `"all"` (default): Apply to all transformers + - `"transformer"`: Apply only to the primary transformer (high noise for Wan2.2) + - `"transformer_2"`: Apply only to transformer_2 (low noise for Wan2.2) + - `"critic"`: Apply only to the critic model +- `strength` (float or list of floats, optional): LoRA strength for merge, default 1.0. If a list, must match the length of `lora_nickname`. Values < 1.0 reduce the effect, values > 1.0 amplify the effect -**Curl Example:** +**Single LoRA Example:** ```bash curl -X POST http://localhost:30010/v1/set_lora \ @@ -259,10 +263,43 @@ curl -X POST http://localhost:30010/v1/set_lora \ -d '{ "lora_nickname": "lora_name", "lora_path": "/path/to/lora.safetensors", + "target": "all", "strength": 0.8 }' ``` +**Multiple LoRA Example:** + +```bash +curl -X POST http://localhost:30010/v1/set_lora \ + -H "Content-Type: application/json" \ + -d '{ + "lora_nickname": ["lora_1", "lora_2"], + "lora_path": ["/path/to/lora1.safetensors", "/path/to/lora2.safetensors"], + "target": ["transformer", "transformer_2"], + "strength": [0.8, 1.0] + }' +``` + +**Multiple LoRA with Same Target:** + +```bash +curl -X POST http://localhost:30010/v1/set_lora \ + -H "Content-Type: application/json" \ + -d '{ + "lora_nickname": ["style_lora", "character_lora"], + "lora_path": ["/path/to/style.safetensors", "/path/to/character.safetensors"], + "target": "all", + "strength": [0.7, 0.9] + }' +``` + +> [!NOTE] +> When using multiple LoRAs: +> - All list parameters (`lora_nickname`, `lora_path`, `target`, `strength`) must have the same length +> - If `target` or `strength` is a single value, it will be applied to all LoRAs +> - Multiple LoRAs applied to the same target will be merged in order + #### Merge LoRA Weights diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index d7518d8d9..d0412f439 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -11,7 +11,7 @@ diffusion models. import multiprocessing as mp import os import time -from typing import Any +from typing import Any, List, Union import numpy as np @@ -21,6 +21,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( MergeLoraWeightsReq, SetLoraReq, UnmergeLoraWeightsReq, + format_lora_message, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( post_process_sample, @@ -317,23 +318,25 @@ class DiffGenerator: def set_lora( self, - lora_nickname: str, - lora_path: str | None = None, - target: str = "all", - strength: float = 1.0, + lora_nickname: Union[str, List[str]], + lora_path: Union[str, None, List[Union[str, None]]] = None, + target: Union[str, List[str]] = "all", + strength: Union[float, List[float]] = 1.0, ) -> None: """ - Set a LoRA adapter for the specified transformer(s). + Set LoRA adapter(s) for the specified transformer(s). + Supports both single LoRA (backward compatible) and multiple LoRA adapters. Args: - lora_nickname: The nickname of the adapter. - lora_path: Path to the LoRA adapter. - target: Which transformer(s) to apply the LoRA to. One of: + lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings. + lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None. + target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings. + Valid values: - "all": Apply to all transformers (default) - "transformer": Apply only to the primary transformer (high noise for Wan2.2) - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2) - "critic": Apply only to the critic model - strength: LoRA strength for merge, default 1.0. + strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats. """ req = SetLoraReq( lora_nickname=lora_nickname, @@ -341,9 +344,13 @@ class DiffGenerator: target=target, strength=strength, ) + nickname_str, target_str, strength_str = format_lora_message( + lora_nickname, target, strength + ) + self._send_lora_request( req, - f"Successfully set LoRA adapter: {lora_nickname} (target: {target}, strength: {strength})", + f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})", "Failed to set LoRA adapter", ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py index 53e2fbc59..63f7444de 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py @@ -1,4 +1,4 @@ -from typing import Any, Optional +from typing import Any, List, Optional, Union from fastapi import APIRouter, Body, HTTPException from fastapi.responses import ORJSONResponse @@ -9,6 +9,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( MergeLoraWeightsReq, SetLoraReq, UnmergeLoraWeightsReq, + format_lora_message, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client @@ -48,23 +49,27 @@ async def _handle_lora_request(req: Any, success_msg: str, failure_msg: str): @router.post("/set_lora") async def set_lora( - lora_nickname: str = Body(..., embed=True), - lora_path: Optional[str] = Body(None, embed=True), - target: str = Body("all", embed=True), - strength: float = Body(1.0, embed=True), + lora_nickname: Union[str, List[str]] = Body(..., embed=True), + lora_path: Optional[Union[str, List[Optional[str]]]] = Body(None, embed=True), + target: Union[str, List[str]] = Body("all", embed=True), + strength: Union[float, List[float]] = Body(1.0, embed=True), ): """ - Set a LoRA adapter for the specified transformer(s). + Set LoRA adapter(s) for the specified transformer(s). + Supports both single LoRA (backward compatible) and multiple LoRA adapters. Args: - lora_nickname: The nickname of the adapter. - lora_path: Path to the LoRA adapter (local path or HF repo id). - target: Which transformer(s) to apply the LoRA to. One of: + lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings. + lora_path: Path(s) to the LoRA adapter(s) (local path or HF repo id). + Can be a string, None, or a list of strings/None. Must match the length of lora_nickname. + target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings. + If a list, must match the length of lora_nickname. Valid values: - "all": Apply to all transformers (default) - "transformer": Apply only to the primary transformer (high noise for Wan2.2) - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2) - "critic": Apply only to the critic model - strength: LoRA strength for merge, default 1.0. Values < 1.0 reduce the effect, + strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats. + If a list, must match the length of lora_nickname. Values < 1.0 reduce the effect, values > 1.0 amplify the effect. """ req = SetLoraReq( @@ -73,9 +78,13 @@ async def set_lora( target=target, strength=strength, ) + nickname_str, target_str, strength_str = format_lora_message( + lora_nickname, target, strength + ) + return await _handle_lora_request( req, - f"Successfully set LoRA adapter: {lora_nickname} (target: {target}, strength: {strength})", + f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})", "Failed to set LoRA adapter", ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index b885086a9..92445ad8b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -24,10 +24,10 @@ logger = init_logger(__name__) @dataclasses.dataclass class SetLoraReq: - lora_nickname: str - lora_path: Optional[str] = None - target: str = "all" # "all", "transformer", "transformer_2", "critic" - strength: float = 1.0 # LoRA strength for merge, default 1.0 + lora_nickname: Union[str, List[str]] + lora_path: Optional[Union[str, List[Optional[str]]]] = None + target: Union[str, List[str]] = "all" + strength: Union[float, List[float]] = 1.0 # LoRA strength for merge, default 1.0 @dataclasses.dataclass @@ -47,6 +47,31 @@ class ListLorasReq: pass +def format_lora_message( + lora_nickname: Union[str, List[str]], + target: Union[str, List[str]], + strength: Union[float, List[float]], +) -> tuple[str, str, str]: + """Format success message for single or multiple LoRAs""" + if isinstance(lora_nickname, list): + nickname_str = ", ".join(lora_nickname) + target_str = ", ".join(target) if isinstance(target, list) else target + strength_str = ( + ", ".join(f"{s:.2f}" for s in strength) + if isinstance(strength, list) + else f"{strength:.2f}" + ) + else: + nickname_str = lora_nickname + target_str = target if isinstance(target, str) else ", ".join(target) + strength_str = ( + f"{strength:.2f}" + if isinstance(strength, (int, float)) + else ", ".join(f"{s:.2f}" for s in strength) + ) + return nickname_str, target_str, strength_str + + def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]: try: parts = size.lower().replace(" ", "").split("x") diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index 2f9de9882..2ef6f4d70 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -55,6 +55,9 @@ class BaseLayerWithLoRA(nn.Module): self.disable_lora: bool = True self.lora_rank = lora_rank self.lora_alpha = lora_alpha + self.lora_weights_list: list[ + tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float] + ] = [] self.lora_path: str | None = None self.strength: float = 1.0 @@ -77,6 +80,7 @@ class BaseLayerWithLoRA(nn.Module): lora_B = self.lora_B.to_local() lora_A = self.lora_A.to_local() + # TODO: Support multiple LoRA adapters when use not merged mode if not self.merged and not self.disable_lora: lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)) lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True)) @@ -104,15 +108,70 @@ class BaseLayerWithLoRA(nn.Module): B: torch.Tensor, lora_path: str | None = None, strength: float = 1.0, + clear_existing: bool = False, ) -> None: - self.lora_A = torch.nn.Parameter( + """ + Set LoRA weights. Supports multiple LoRA adapters. + + Args: + A: LoRA A weight tensor + B: LoRA B weight tensor + lora_path: Path to the LoRA adapter (for logging) + strength: LoRA strength + clear_existing: If True, clear existing LoRA weights before adding new one. + If False, append to existing list (for multi-LoRA support). + """ + lora_A_param = torch.nn.Parameter( A ) # share storage with weights in the pipeline - self.lora_B = torch.nn.Parameter(B) - self.disable_lora = False - self.strength = strength - self.merge_lora_weights() + lora_B_param = torch.nn.Parameter(B) + + if clear_existing: + self.lora_weights_list.clear() + # Also clear backward compatibility attributes + self.lora_A = None + self.lora_B = None + self.lora_path = None + self.strength = 1.0 + + # Add to list for multi-LoRA support + self.lora_weights_list.append((lora_A_param, lora_B_param, lora_path, strength)) + + # Set backward compatibility attributes to point to the last LoRA (for single LoRA case) + # This ensures backward compatibility while supporting multiple LoRA + self.lora_A = lora_A_param + self.lora_B = lora_B_param self.lora_path = lora_path + self.strength = strength + + self.disable_lora = False + self.merge_lora_weights() + + @torch.no_grad() + def _merge_lora_into_data( + self, + data: torch.Tensor, + lora_list: list[ + tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float] + ], + ) -> None: + """ + Merge all LoRA adapters into the data tensor in-place. + + Args: + data: The base weight tensor to merge LoRA into (modified in-place) + lora_list: List of (lora_A, lora_B, lora_path, lora_strength) tuples + """ + # Merge all LoRA adapters in order + for lora_A, lora_B, _, lora_strength in lora_list: + lora_delta = self.slice_lora_b_weights( + lora_B.to(data) + ) @ self.slice_lora_a_weights(lora_A.to(data)) + # Apply lora_alpha / lora_rank scaling for consistency with forward() + if self.lora_alpha is not None and self.lora_rank is not None: + if self.lora_alpha != self.lora_rank: + lora_delta = lora_delta * (self.lora_alpha / self.lora_rank) + data += lora_strength * lora_delta @torch.no_grad() def merge_lora_weights(self, strength: float | None = None) -> None: @@ -124,9 +183,15 @@ class BaseLayerWithLoRA(nn.Module): if self.merged: self.unmerge_lora_weights() - assert ( - self.lora_A is not None and self.lora_B is not None - ), "LoRA weights not set. Please set them first." + + # Use lora_weights_list if available, otherwise fall back to single LoRA for backward compatibility + lora_list = self.lora_weights_list if self.lora_weights_list else [] + if not lora_list and self.lora_A is not None and self.lora_B is not None: + lora_list = [(self.lora_A, self.lora_B, self.lora_path, self.strength)] + + if not lora_list: + raise ValueError("LoRA weights not set. Please set them first.") + if isinstance(self.base_layer.weight, DTensor): mesh = self.base_layer.weight.data.device_mesh unsharded_base_layer = ReplicatedLinear( @@ -143,14 +208,9 @@ class BaseLayerWithLoRA(nn.Module): data = self.base_layer.weight.data.to( get_local_torch_device() ).full_tensor() - lora_delta = self.slice_lora_b_weights(self.lora_B).to( - data - ) @ self.slice_lora_a_weights(self.lora_A).to(data) - # Apply lora_alpha / lora_rank scaling for consistency with forward() - if self.lora_alpha is not None and self.lora_rank is not None: - if self.lora_alpha != self.lora_rank: - lora_delta = lora_delta * (self.lora_alpha / self.lora_rank) - data += self.strength * lora_delta + + self._merge_lora_into_data(data, lora_list) + unsharded_base_layer.weight = nn.Parameter(data.to(current_device)) if isinstance(getattr(self.base_layer, "bias", None), DTensor): unsharded_base_layer.bias = nn.Parameter( @@ -173,14 +233,9 @@ class BaseLayerWithLoRA(nn.Module): else: current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()) - lora_delta = self.slice_lora_b_weights( - self.lora_B.to(data) - ) @ self.slice_lora_a_weights(self.lora_A.to(data)) - # Apply lora_alpha / lora_rank scaling for consistency with forward() - if self.lora_alpha is not None and self.lora_rank is not None: - if self.lora_alpha != self.lora_rank: - lora_delta = lora_delta * (self.lora_alpha / self.lora_rank) - data += self.strength * lora_delta + + self._merge_lora_into_data(data, lora_list) + self.base_layer.weight.data = data.to(current_device, non_blocking=True) self.merged = True @@ -403,6 +458,7 @@ class LinearWithLoRA(BaseLayerWithLoRA): lora_B = self.lora_B.to_local() lora_A = self.lora_A.to_local() + # TODO: Support multiple LoRA adapters when use not merged mode if not self.merged and not self.disable_lora: lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)) lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True)) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 129018c0e..2f1d0996b 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -4,7 +4,7 @@ import multiprocessing as mp import os import time -from typing import List +from typing import List, Union import torch from setproctitle import setproctitle @@ -206,19 +206,20 @@ class GPUWorker: def set_lora( self, - lora_nickname: str, - lora_path: str | None = None, - target: str = "all", - strength: float = 1.0, + lora_nickname: Union[str, List[str]], + lora_path: Union[str, None, List[Union[str, None]]] = None, + target: Union[str, List[str]] = "all", + strength: Union[float, List[float]] = 1.0, ) -> OutputBatch: """ - Set the LoRA adapter for the pipeline. + Set the LoRA adapter(s) for the pipeline. + Supports both single LoRA (backward compatible) and multiple LoRA adapters. Args: - lora_nickname: The nickname of the adapter. - lora_path: Path to the LoRA adapter. - target: Which transformer(s) to apply the LoRA to. - strength: LoRA strength for merge, default 1.0. + lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings. + lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None. + target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings. + strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats. """ if not isinstance(self.pipeline, LoRAPipeline): return OutputBatch(error="Lora is not enabled") diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py index 929be44af..6079faf65 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py @@ -48,6 +48,7 @@ class LoRAPipeline(ComposedPipelineBase): cur_adapter_name: dict[str, str] cur_adapter_path: dict[str, str] cur_adapter_strength: dict[str, float] # Track current strength per module + cur_adapter_config: dict[str, tuple[list[str], list[float]]] # [dit_layer_name] = wrapped_lora_layer lora_layers: dict[str, BaseLayerWithLoRA] lora_layers_critic: dict[str, BaseLayerWithLoRA] @@ -74,6 +75,8 @@ class LoRAPipeline(ComposedPipelineBase): self.cur_adapter_name = {} self.cur_adapter_path = {} self.cur_adapter_strength = {} + # Track full LoRA config: {module_name: (nickname_list, strength_list)} + self.cur_adapter_config = {} self.lora_layers = {} self.lora_layers_critic = {} self.lora_layers_transformer_2 = {} @@ -297,50 +300,167 @@ class LoRAPipeline(ComposedPipelineBase): converted_count_critic, ) + def _normalize_lora_params( + self, + lora_nickname: str | list[str], + lora_path: str | None | list[str | None], + strength: float | list[float], + target: str | list[str], + ) -> tuple[list[str], list[str | None], list[float], list[str]]: + """ + Normalize LoRA parameters to lists for multi-LoRA support. + + Requirements: + - each nickname must have a corresponding lora_path (no implicit repeat) + - strength / target if scalar broadcast, else length must match nickname + """ + # nickname + if isinstance(lora_nickname, str): + lora_nicknames = [lora_nickname] + else: + lora_nicknames = lora_nickname + + # lora_path: require 1:1 mapping with nickname (no implicit repeat) + if isinstance(lora_path, list): + lora_paths = lora_path + else: + lora_paths = [lora_path] + if len(lora_paths) != len(lora_nicknames): + raise ValueError( + f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, " + f"but lora_path has {len(lora_paths)} items. " + "Provide one path per nickname." + ) + + # strength and target: allow scalar broadcast, else length must match + if isinstance(strength, (int, float)): + strengths = [float(strength)] * len(lora_nicknames) + else: + strengths = [float(s) for s in strength] + if len(strengths) != len(lora_nicknames): + raise ValueError( + f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, " + f"but strength has {len(strengths)} items" + ) + + if isinstance(target, str): + targets = [target] * len(lora_nicknames) + else: + targets = target + if len(targets) != len(lora_nicknames): + raise ValueError( + f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, " + f"but target has {len(targets)} items" + ) + return lora_nicknames, lora_paths, strengths, targets + + def _check_lora_config_matches( + self, + module_name: str, + target_nicknames: list[str], + target_strengths: list[float], + adapter_updated: bool, + ) -> bool: + """ + Check if current LoRA configuration matches the target configuration. + + Args: + module_name: The name of the module to check. + target_nicknames: List of LoRA nicknames to apply. + target_strengths: List of LoRA strengths to apply. + adapter_updated: Whether any adapter was updated/loaded. + + Returns: + True if the configuration matches exactly (including order and strength), False otherwise. + """ + if not self.is_lora_merged.get(module_name, False): + return False + if adapter_updated: + return False # Adapter was updated, need to reapply + + stored_config = self.cur_adapter_config.get(module_name) + if stored_config is None: + return False + + stored_nicknames, stored_strengths = stored_config + # Compare: nickname list and strength list must match exactly (including order) + return ( + stored_nicknames == target_nicknames + and stored_strengths == target_strengths + ) + def _apply_lora_to_layers( self, lora_layers: dict[str, BaseLayerWithLoRA], - lora_nickname: str, - lora_path: str | None, + lora_nicknames: list[str], + lora_paths: list[str | None], rank: int, - strength: float = 1.0, + strengths: list[float], + clear_existing: bool = False, ) -> int: """ - Apply LoRA weights to the given lora_layers. + Apply LoRA weights to the given lora_layers. Supports multiple LoRA adapters. Args: lora_layers: The dictionary of LoRA layers to apply weights to. - lora_nickname: The nickname of the LoRA adapter. - lora_path: The path to the LoRA adapter. + lora_nicknames: The list of nicknames of the LoRA adapters. + lora_paths: The list of paths to the LoRA adapters. Must match length of lora_nicknames. rank: The distributed rank (for logging). - strength: LoRA strength for merge, default 1.0. + strengths: The list of LoRA strengths for merge. Must match length of lora_nicknames. + clear_existing: If True, clear existing LoRA weights before adding new ones. Returns: The number of layers that had LoRA weights applied. """ + if len(lora_paths) != len(lora_nicknames): + raise ValueError( + f"Length mismatch: lora_nicknames has {len(lora_nicknames)} items, " + f"but lora_paths has {len(lora_paths)} items" + ) + if len(strengths) != len(lora_nicknames): + raise ValueError( + f"Length mismatch: lora_nicknames has {len(lora_nicknames)} items, " + f"but strengths has {len(strengths)} items" + ) + adapted_count = 0 for name, layer in lora_layers.items(): - lora_A_name = name + ".lora_A" - lora_B_name = name + ".lora_B" - if ( - lora_A_name in self.lora_adapters[lora_nickname] - and lora_B_name in self.lora_adapters[lora_nickname] + # Apply all LoRA adapters in order + for idx, (nickname, path, lora_strength) in enumerate( + zip(lora_nicknames, lora_paths, strengths) ): - layer.set_lora_weights( - self.lora_adapters[lora_nickname][lora_A_name], - self.lora_adapters[lora_nickname][lora_B_name], - lora_path=lora_path, - strength=strength, - ) - adapted_count += 1 - else: - if rank == 0: - logger.warning( - "LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.", - lora_path, - name, + lora_A_name = name + ".lora_A" + lora_B_name = name + ".lora_B" + if ( + lora_A_name in self.lora_adapters[nickname] + and lora_B_name in self.lora_adapters[nickname] + ): + layer.set_lora_weights( + self.lora_adapters[nickname][lora_A_name], + self.lora_adapters[nickname][lora_B_name], + lora_path=path, + strength=lora_strength, + clear_existing=( + clear_existing and idx == 0 + ), # Only clear on first LoRA ) - layer.disable_lora = True + adapted_count += 1 + else: + if rank == 0 and idx == 0: # Only warn for first missing LoRA + logger.warning( + "LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.", + path, + name, + ) + # Only disable if no LoRA was applied at all + if idx == len(lora_nicknames) - 1: + has_any_lora = any( + name + ".lora_A" in self.lora_adapters[n] + and name + ".lora_B" in self.lora_adapters[n] + for n in lora_nicknames + ) + if not has_any_lora: + layer.disable_lora = True return adapted_count def is_lora_effective(self, target: str = "all") -> bool: @@ -438,46 +558,25 @@ class LoRAPipeline(ComposedPipelineBase): def set_lora( self, - lora_nickname: str, - lora_path: str | None = None, - target: str = "all", - strength: float = 1.0, + lora_nickname: str | list[str], + lora_path: str | None | list[str | None] = None, + target: str | list[str] = "all", + strength: float | list[float] = 1.0, ): # type: ignore """ - Load a LoRA adapter into the pipeline and apply it to the specified transformer(s). - - Args: - lora_nickname: The "nick name" of the adapter when referenced in the pipeline. - lora_path: The path to the adapter, either a local path or a Hugging Face repo id. - target: Which transformer(s) to apply the LoRA to. One of: - - "all": Apply to all transformers (default, backward compatible) - - "transformer": Apply only to the primary transformer (high noise for Wan2.2) - - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2) - - "critic": Apply only to the critic model (fake_score_transformer) - strength: LoRA strength for merge, default 1.0. + Load LoRA adapter(s) into the pipeline and apply them to the specified transformer(s). + Supports both single LoRA (backward compatible) and multiple LoRA adapters. """ - if target not in self.VALID_TARGETS: - raise ValueError( - f"Invalid target: {target}. Valid targets: {self.VALID_TARGETS}" - ) + # Normalize inputs to lists for multi-LoRA support + lora_nicknames, lora_paths, strengths, targets = self._normalize_lora_params( + lora_nickname, lora_path, strength, target + ) - # Check if any target module has a different LoRA merged - target_modules, error = self._get_target_lora_layers(target) - if error: - logger.warning("set_lora: %s", error) - for module_name, _ in target_modules: - if ( - self.is_lora_merged.get(module_name, False) - and self.cur_adapter_name.get(module_name) != lora_nickname - ): - raise ValueError( - f"LoRA '{self.cur_adapter_name.get(module_name)}' is currently merged in {module_name}. " - "Please call 'unmerge_lora_weights' before setting a new LoRA." - ) - - if lora_nickname not in self.lora_adapters and lora_path is None: + # Validate targets + invalid_targets = [t for t in targets if t not in self.VALID_TARGETS] + if invalid_targets: raise ValueError( - f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it." + f"Invalid target(s): {invalid_targets}. Valid targets: {self.VALID_TARGETS}" ) # Disable layerwise offload before convert_to_lora_layers to ensure weights are accessible @@ -489,64 +588,100 @@ class LoRAPipeline(ComposedPipelineBase): ): self.convert_to_lora_layers() - # Re-fetch target_modules after convert_to_lora_layers() to get populated dicts - target_modules, error = self._get_target_lora_layers(target) - if error: - logger.warning("set_lora: %s", error) - + # Check adapter presence and load missing adapters adapter_updated = False rank = dist.get_rank() - should_load = False - if lora_path is not None: - if lora_nickname not in self.loaded_adapter_paths: - should_load = True - elif self.loaded_adapter_paths[lora_nickname] != lora_path: - should_load = True + # load required adapters + for nickname, path in zip(lora_nicknames, lora_paths): + if nickname not in self.lora_adapters and path is None: + raise ValueError( + f"Adapter {nickname} not found in the pipeline. Please provide lora_path to load it." + ) + # Check if adapter needs to be loaded + should_load = False + if path is not None: + if nickname not in self.loaded_adapter_paths: + should_load = True + elif self.loaded_adapter_paths[nickname] != path: + should_load = True + if should_load: + adapter_updated = True + self.load_lora_adapter(path, nickname, rank) - if should_load: - adapter_updated = True - self.load_lora_adapter(lora_path, lora_nickname, rank) + # Group by target to apply separately + target_to_indices = {} + for idx, tgt in enumerate(targets): + if tgt not in target_to_indices: + target_to_indices[tgt] = [] + target_to_indices[tgt].append(idx) - # Check if we can skip (same adapter already applied to all target modules with same strength) - all_already_applied = all( - not adapter_updated - and self.cur_adapter_name.get(module_name) == lora_nickname - and self.is_lora_merged.get(module_name, False) - and self.cur_adapter_strength.get(module_name) == strength - for module_name, _ in target_modules + adapted_count = 0 + for tgt, idx_list in target_to_indices.items(): + target_modules, error = self._get_target_lora_layers(tgt) + if error: + logger.warning("set_lora: %s", error) + if not target_modules: + continue + + # Disable layerwise offload if enabled: load all layers to GPU + # the LoRA weights merging process requires weights being on device + with self._temporarily_disable_offload(target_modules=target_modules): + tgt_nicknames = [lora_nicknames[i] for i in idx_list] + tgt_paths = [lora_paths[i] for i in idx_list] + tgt_strengths = [strengths[i] for i in idx_list] + + merged_name = ( + ",".join(tgt_nicknames) + if len(tgt_nicknames) > 1 + else tgt_nicknames[0] + ) + + # Skip if LoRA configuration matches exactly (including order and strength) + # Since all modules for the same target apply the same config, checking one is sufficient + first_module_name, _ = target_modules[0] + if self._check_lora_config_matches( + first_module_name, tgt_nicknames, tgt_strengths, adapter_updated + ): + logger.info("LoRA configuration matches exactly, skipping") + continue + + # Apply LoRA to modules for this target + for module_name, lora_layers_dict in target_modules: + count = self._apply_lora_to_layers( + lora_layers_dict, + tgt_nicknames, + tgt_paths, + rank, + tgt_strengths, + clear_existing=True, + ) + adapted_count += count + self.cur_adapter_name[module_name] = merged_name + self.cur_adapter_path[module_name] = ",".join( + str(p or self.loaded_adapter_paths.get(n, "")) + for n, p in zip(tgt_nicknames, tgt_paths) + ) + self.is_lora_merged[module_name] = True + self.cur_adapter_strength[module_name] = tgt_strengths[0] + # Store full configuration for multi-LoRA support (preserves order and all strengths) + self.cur_adapter_config[module_name] = ( + tgt_nicknames.copy(), + tgt_strengths.copy(), + ) + + logger.info( + "Rank %d: LoRA adapter(s) %s applied to %d layers (targets: %s, strengths: %s)", + rank, + ", ".join(map(str, lora_paths)) if lora_paths else None, + adapted_count, + ", ".join(targets) if len(set(targets)) > 1 else targets[0], + ( + ", ".join(f"{s:.2f}" for s in strengths) + if len(strengths) > 1 + else f"{strengths[0]:.2f}" + ), ) - if all_already_applied: - return - - # Disable layerwise offload if enabled: load all layers to GPU - with self._temporarily_disable_offload(target_modules=target_modules): - # Apply LoRA to target modules (now all layers are on GPU) - adapted_count = 0 - for module_name, lora_layers_dict in target_modules: - count = self._apply_lora_to_layers( - lora_layers_dict, - lora_nickname, - lora_path, - rank, - strength, - ) - adapted_count += count - self.cur_adapter_name[module_name] = lora_nickname - self.cur_adapter_path[module_name] = ( - lora_path or self.loaded_adapter_paths.get(lora_nickname, "") - ) - self.is_lora_merged[module_name] = True - self.cur_adapter_strength[module_name] = strength - - logger.info( - "Rank %d: LoRA adapter %s applied to %d layers (target: %s, strength: %s)", - rank, - lora_path, - adapted_count, - target, - strength, - ) def merge_lora_weights(self, target: str = "all", strength: float = 1.0) -> None: """ @@ -649,7 +784,8 @@ class LoRAPipeline(ComposedPipelineBase): continue self.is_lora_merged[module_name] = False self.cur_adapter_strength.pop(module_name, None) - logger.info("LoRA weights unmerged for %s", module_name) + self.cur_adapter_config.pop(module_name, None) + logger.info("LoRA weights unmerged for %s", module_name) def get_lora_status(self) -> dict[str, Any]: """ diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index a14903ac7..769d5adf3 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -529,6 +529,31 @@ "expected_avg_denoise_ms": 94.15, "expected_median_denoise_ms": 102.03 }, + "zimage_image_t2i_multi_lora": { + "stages_ms": { + "InputValidationStage": 0.04, + "TextEncodingStage": 103.81, + "ConditioningStage": 0.01, + "TimestepPreparationStage": 1.3, + "LatentPreparationStage": 0.11, + "DenoisingStage": 813.7, + "DecodingStage": 34.51 + }, + "denoise_step_ms": { + "0": 30.35, + "1": 74.53, + "2": 99.34, + "3": 100.92, + "4": 99.46, + "5": 100.57, + "6": 99.72, + "7": 100.86, + "8": 103.87 + }, + "expected_e2e_ms": 955.14, + "expected_avg_denoise_ms": 89.96, + "expected_median_denoise_ms": 99.72 + }, "zimage_image_t2i_2_gpus": { "stages_ms": { "InputValidationStage": 0.08, diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 6dcdbaedc..66c29d627 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -556,6 +556,76 @@ Consider updating perf_baselines.json with the snippets below: assert resp.status_code == 200, f"Dynamic set_lora failed: {resp.text}" logger.info("[Dynamic LoRA] set_lora succeeded for %s", case.id) + def _test_multi_lora_e2e( + self, + ctx: ServerContext, + case: DiffusionTestCase, + generate_fn: Callable[[str, openai.Client], str], + first_lora_path: str, + second_lora_path: str, + ) -> None: + """ + Test multiple LoRA adapters with different set_lora input scenarios. + Tests: basic multi-LoRA, different strengths, cached adapters, switch back to single. + """ + base_url = f"http://localhost:{ctx.port}/v1" + client = OpenAI(base_url=base_url, api_key="dummy") + + # Test 1: Basic multi-LoRA with list format + resp = requests.post( + f"{base_url}/set_lora", + json={ + "lora_nickname": ["default", "lora2"], + "lora_path": [first_lora_path, second_lora_path], + "target": "all", + "strength": [1.0, 1.0], + }, + ) + assert ( + resp.status_code == 200 + ), f"set_lora with multiple adapters failed: {resp.text}" + assert generate_fn(case.id, client) is not None + + # Test 2: Different strengths + resp = requests.post( + f"{base_url}/set_lora", + json={ + "lora_nickname": ["default", "lora2"], + "lora_path": [first_lora_path, second_lora_path], + "target": "all", + "strength": [0.8, 0.5], + }, + ) + assert ( + resp.status_code == 200 + ), f"set_lora with different strengths failed: {resp.text}" + assert generate_fn(case.id, client) is not None + + # Test 3: Different targets + requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"}) + resp = requests.post( + f"{base_url}/set_lora", + json={ + "lora_nickname": ["default", "lora2"], + "lora_path": [first_lora_path, second_lora_path], + "target": ["transformer", "transformer_2"], + "strength": [0.8, 0.5], + }, + ) + assert ( + resp.status_code == 200 + ), f"set_lora with cached adapters failed: {resp.text}" + assert generate_fn(case.id, client) is not None + + # Test 4: Switch back to single LoRA + resp = requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"}) + assert ( + resp.status_code == 200 + ), f"set_lora back to single adapter failed: {resp.text}" + assert generate_fn(case.id, client) is not None + + logger.info("[Multi-LoRA] All multi-LoRA tests passed for %s", case.id) + def _test_v1_models_endpoint( self, ctx: ServerContext, case: DiffusionTestCase ) -> None: @@ -681,3 +751,21 @@ Consider updating perf_baselines.json with the snippets below: # LoRA API functionality test with E2E validation (only for LoRA-enabled cases) if case.server_args.lora_path or case.server_args.dynamic_lora_path: self._test_lora_api_functionality(diffusion_server, case, generate_fn) + + # Test dynamic LoRA switching (requires a second LoRA adapter) + if case.server_args.second_lora_path: + self._test_lora_dynamic_switch_e2e( + diffusion_server, + case, + generate_fn, + case.server_args.second_lora_path, + ) + + # Test multi-LoRA functionality + self._test_multi_lora_e2e( + diffusion_server, + case, + generate_fn, + case.server_args.lora_path, + case.server_args.second_lora_path, + ) diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 4a20274d1..562435511 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -155,6 +155,9 @@ class DiffusionServerArgs: dynamic_lora_path: str | None = ( None # LoRA path for dynamic loading test (loaded via set_lora after startup) ) + second_lora_path: str | None = ( + None # Second LoRA adapter path for multi-LoRA testing + ) # misc enable_warmup: bool = False @@ -361,6 +364,17 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [ ), T2I_sampling_params, ), + # Multi-LoRA test case for Z-Image-Turbo + DiffusionTestCase( + "zimage_image_t2i_multi_lora", + DiffusionServerArgs( + model_path="Tongyi-MAI/Z-Image-Turbo", + modality="image", + lora_path="reverentelusarca/elusarca-anime-style-lora-z-image-turbo", + second_lora_path="tarn59/pixel_art_style_lora_z_image_turbo", + ), + T2I_sampling_params, + ), # === Text and Image to Image (TI2I) === DiffusionTestCase( "qwen_image_edit_ti2i",