[diffusion] fix: support applying different LoRA adapters to different transformers in multi-transformer pipelines (#14839)

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -82,9 +82,6 @@ class DiffGenerator:
         # The executor is now a client to the Scheduler service
         self.local_scheduler_process: list[mp.Process] | None = None
         self.owns_scheduler_client: bool = False
-        self._current_lora_path: str | None = None
-        self._current_lora_nickname: str | None = None
-        self._is_lora_merged: bool = False
 
     @classmethod
     def from_pretrained(
@@ -344,29 +341,57 @@ class DiffGenerator:
             )
             raise RuntimeError(f"{failure_msg}: {error_msg}")
 
-    def set_lora(self, lora_nickname: str, lora_path: str | None = None) -> None:
-        req = SetLoraReq(lora_nickname=lora_nickname, lora_path=lora_path)
+    def set_lora(
+        self, lora_nickname: str, lora_path: str | None = None, target: str = "all"
+    ) -> None:
+        """
+        Set a LoRA adapter for the specified transformer(s).
+
+        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:
+                - "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
+        """
+        req = SetLoraReq(
+            lora_nickname=lora_nickname, lora_path=lora_path, target=target
+        )
         self._send_lora_request(
             req,
-            f"Successfully set LoRA adapter: {lora_nickname}",
+            f"Successfully set LoRA adapter: {lora_nickname} (target: {target})",
             "Failed to set LoRA adapter",
         )
 
-    def unmerge_lora_weights(self) -> None:
-        req = UnmergeLoraWeightsReq()
+    def unmerge_lora_weights(self, target: str = "all") -> None:
+        """
+        Unmerge LoRA weights from the base model.
+
+        Args:
+            target: Which transformer(s) to unmerge.
+        """
+        req = UnmergeLoraWeightsReq(target=target)
         self._send_lora_request(
             req,
-            "Successfully unmerged LoRA weights",
+            f"Successfully unmerged LoRA weights (target: {target})",
             "Failed to unmerge LoRA weights",
         )
-        self._is_lora_merged = False
 
-    def merge_lora_weights(self) -> None:
-        req = MergeLoraWeightsReq()
+    def merge_lora_weights(self, target: str = "all") -> None:
+        """
+        Merge LoRA weights into the base model.
+
+        Args:
+            target: Which transformer(s) to merge.
+        """
+        req = MergeLoraWeightsReq(target=target)
         self._send_lora_request(
-            req, "Successfully merged LoRA weights", "Failed to merge LoRA weights"
+            req,
+            f"Successfully merged LoRA weights (target: {target})",
+            "Failed to merge LoRA weights",
         )
-        self._is_lora_merged = True
 
     def _ensure_lora_state(
         self,
@@ -374,29 +399,27 @@ class DiffGenerator:
         lora_nickname: str | None = None,
         merge_lora: bool = True,
     ) -> None:
+        """
+        Ensure the LoRA state matches the desired configuration.
+
+        Note: This method does not cache client-side state. The server handles
+        idempotent operations, so redundant calls are safe but may have minor overhead.
+        """
         if lora_path is None:
-            if self._is_lora_merged:
-                self.unmerge_lora_weights()
-            self._current_lora_path = None
-            self._current_lora_nickname = None
-            self._is_lora_merged = False
+            # Unmerge all LoRA weights when no lora_path is provided
+            self.unmerge_lora_weights()
             return
 
         lora_nickname = lora_nickname or self.server_args.lora_nickname
 
-        if self._current_lora_path != lora_path:
-            if self._is_lora_merged:
-                self.unmerge_lora_weights()
-                self._is_lora_merged = False
-            self.set_lora(lora_nickname, lora_path)
-            self._current_lora_path = lora_path
-            self._current_lora_nickname = lora_nickname
-            self._is_lora_merged = False
+        # Set the LoRA adapter (server handles idempotent logic)
+        self.set_lora(lora_nickname, lora_path)
 
-        if merge_lora and not self._is_lora_merged:
+        # Merge or unmerge based on the merge_lora flag
+        if merge_lora:
             self.merge_lora_weights()
-        elif not merge_lora:
-            self._is_lora_merged = False
+        else:
+            self.unmerge_lora_weights()
 
     def generate_with_lora(
         self,

python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py

@@ -37,26 +37,61 @@ async def _handle_lora_request(req: Any, success_msg: str, failure_msg: str):
 async def set_lora(
     lora_nickname: str = Body(..., embed=True),
     lora_path: Optional[str] = Body(None, embed=True),
+    target: str = Body("all", embed=True),
 ):
-    req = SetLoraReq(lora_nickname=lora_nickname, lora_path=lora_path)
+    """
+    Set a LoRA adapter for the specified transformer(s).
+
+    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:
+            - "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
+    """
+    req = SetLoraReq(lora_nickname=lora_nickname, lora_path=lora_path, target=target)
     return await _handle_lora_request(
         req,
-        f"Successfully set LoRA adapter: {lora_nickname}",
+        f"Successfully set LoRA adapter: {lora_nickname} (target: {target})",
         "Failed to set LoRA adapter",
     )
 
 
 @router.post("/merge_lora_weights")
-async def merge_lora_weights():
-    req = MergeLoraWeightsReq()
+async def merge_lora_weights(
+    target: str = Body("all", embed=True),
+):
+    """
+    Merge LoRA weights into the base model.
+
+    Args:
+        target: Which transformer(s) to merge. One of "all", "transformer",
+                "transformer_2", "critic".
+    """
+    req = MergeLoraWeightsReq(target=target)
     return await _handle_lora_request(
-        req, "Successfully merged LoRA weights", "Failed to merge LoRA weights"
+        req,
+        f"Successfully merged LoRA weights (target: {target})",
+        "Failed to merge LoRA weights",
     )
 
 
 @router.post("/unmerge_lora_weights")
-async def unmerge_lora_weights():
-    req = UnmergeLoraWeightsReq()
+async def unmerge_lora_weights(
+    target: str = Body("all", embed=True),
+):
+    """
+    Unmerge LoRA weights from the base model.
+
+    Args:
+        target: Which transformer(s) to unmerge. One of "all", "transformer",
+                "transformer_2", "critic".
+    """
+    req = UnmergeLoraWeightsReq(target=target)
     return await _handle_lora_request(
-        req, "Successfully unmerged LoRA weights", "Failed to unmerge LoRA weights"
+        req,
+        f"Successfully unmerged LoRA weights (target: {target})",
+        "Failed to unmerge LoRA weights",
     )

python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py

@@ -25,16 +25,17 @@ logger = init_logger(__name__)
 class SetLoraReq:
     lora_nickname: str
     lora_path: Optional[str] = None
+    target: str = "all"  # "all", "transformer", "transformer_2", "critic"
 
 
 @dataclasses.dataclass
 class MergeLoraWeightsReq:
-    pass
+    target: str = "all"  # "all", "transformer", "transformer_2", "critic"
 
 
 @dataclasses.dataclass
 class UnmergeLoraWeightsReq:
-    pass
+    target: str = "all"  # "all", "transformer", "transformer_2", "critic"
 
 
 def post_process_sample(

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -51,7 +51,8 @@ class BaseLayerWithLoRA(nn.Module):
         self.cpu_weight = base_layer.weight.to("cpu")
         # indicates adapter weights don't contain this layer
         # (which shouldn't normally happen, but we want to separate it from the case of erroneous merging)
-        self.disable_lora: bool = False
+        # Default to True to prevent using uninitialized weights; set to False when weights are loaded
+        self.disable_lora: bool = True
         self.lora_rank = lora_rank
         self.lora_alpha = lora_alpha
         self.lora_path: str | None = None

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -121,26 +121,39 @@ class GPUWorker:
         finally:
             return output_batch
 
-    def set_lora(self, lora_nickname: str, lora_path: str | None = None) -> None:
+    def set_lora(
+        self, lora_nickname: str, lora_path: str | None = None, target: str = "all"
+    ) -> None:
         """
         Set the LoRA adapter for the pipeline.
+
+        Args:
+            lora_nickname: The nickname of the adapter.
+            lora_path: Path to the LoRA adapter.
+            target: Which transformer(s) to apply the LoRA to.
         """
         assert self.pipeline is not None
-        self.pipeline.set_lora(lora_nickname, lora_path)
+        self.pipeline.set_lora(lora_nickname, lora_path, target)
 
-    def merge_lora_weights(self) -> None:
+    def merge_lora_weights(self, target: str = "all") -> None:
         """
         Merge LoRA weights.
+
+        Args:
+            target: Which transformer(s) to merge.
         """
         assert self.pipeline is not None
-        self.pipeline.merge_lora_weights()
+        self.pipeline.merge_lora_weights(target)
 
-    def unmerge_lora_weights(self) -> None:
+    def unmerge_lora_weights(self, target: str = "all") -> None:
         """
         Unmerge LoRA weights.
+
+        Args:
+            target: Which transformer(s) to unmerge.
         """
         assert self.pipeline is not None
-        self.pipeline.unmerge_lora_weights()
+        self.pipeline.unmerge_lora_weights(target)
 
 
 def run_scheduler_process(

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -78,15 +78,17 @@ class Scheduler:
     def _handle_set_lora(self, reqs: List[Any]):
         # TODO: return set status
         req = reqs[0]
-        self.worker.set_lora(req.lora_nickname, req.lora_path)
+        self.worker.set_lora(req.lora_nickname, req.lora_path, req.target)
         return {"status": "ok"}
 
-    def _handle_merge_lora(self, _reqs: List[Any]):
-        self.worker.merge_lora_weights()
+    def _handle_merge_lora(self, reqs: List[Any]):
+        req = reqs[0]
+        self.worker.merge_lora_weights(req.target)
         return {"status": "ok"}
 
-    def _handle_unmerge_lora(self, _reqs: List[Any]):
-        self.worker.unmerge_lora_weights()
+    def _handle_unmerge_lora(self, reqs: List[Any]):
+        req = reqs[0]
+        self.worker.unmerge_lora_weights(req.target)
         return {"status": "ok"}
 
     def _handle_generation(self, reqs: List[Req]):

python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py

@@ -35,31 +35,50 @@ class LoRAPipeline(ComposedPipelineBase):
     Pipeline that supports injecting LoRA adapters into the diffusion transformer.
     """
 
+    # Type annotations for instance attributes (initialized in __init__)
     # [lora_nickname][target_LoRA_weight_name_in_SGLang_dit] = weight
     # e.g., [jinx][transformer_blocks.0.attn.to_v.lora_A]
-    lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict(
-        dict
-    )  # state dicts of loaded lora adapters
-    loaded_adapter_paths: dict[str, str] = {}  # nickname -> lora_path
-    cur_adapter_name: str = ""
-    cur_adapter_path: str = ""
+    lora_adapters: dict[str, dict[str, torch.Tensor]]
+    loaded_adapter_paths: dict[str, str]  # nickname -> lora_path
+    # Track current adapter per module: {"transformer": "high_lora", "transformer_2": "low_lora"}
+    cur_adapter_name: dict[str, str]
+    cur_adapter_path: dict[str, str]
     # [dit_layer_name] = wrapped_lora_layer
-    lora_layers: dict[str, BaseLayerWithLoRA] = {}
-    lora_layers_critic: dict[str, BaseLayerWithLoRA] = {}
-    lora_layers_transformer_2: dict[str, BaseLayerWithLoRA] = {}
+    lora_layers: dict[str, BaseLayerWithLoRA]
+    lora_layers_critic: dict[str, BaseLayerWithLoRA]
+    lora_layers_transformer_2: dict[str, BaseLayerWithLoRA]
     server_args: ServerArgs
-    exclude_lora_layers: list[str] = []
-    device: torch.device = get_local_torch_device()
-    lora_target_modules: list[str] | None = None
-    lora_path: str | None = None
-    lora_nickname: str = "default"
-    lora_rank: int | None = None
-    lora_alpha: int | None = None
-    lora_initialized: bool = False
-    is_lora_merged: bool = False
+    exclude_lora_layers: list[str]
+    device: torch.device
+    lora_target_modules: list[str] | None
+    lora_path: str | None
+    lora_nickname: str
+    lora_rank: int | None
+    lora_alpha: int | None
+    lora_initialized: bool
+    # Track merge status per module: {"transformer": True, "transformer_2": False}
+    is_lora_merged: dict[str, bool]
+    # Valid target values for set_lora (class constant, immutable)
+    VALID_TARGETS: list[str] = ["all", "transformer", "transformer_2", "critic"]
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
+        # Initialize all mutable instance attributes to avoid sharing across instances
+        self.lora_adapters = defaultdict(dict)
+        self.loaded_adapter_paths = {}
+        self.cur_adapter_name = {}
+        self.cur_adapter_path = {}
+        self.lora_layers = {}
+        self.lora_layers_critic = {}
+        self.lora_layers_transformer_2 = {}
+        self.is_lora_merged = {}
+        self.lora_initialized = False
+        self.lora_rank = None
+        self.lora_alpha = None
+        self.lora_path = None
+        self.lora_nickname = "default"
+
+        # Initialize from server_args
         self.device = get_local_torch_device()
         self.exclude_lora_layers = (
             self.server_args.pipeline_config.dit_config.arch_config.exclude_lora_layers
@@ -80,6 +99,45 @@ class LoRAPipeline(ComposedPipelineBase):
             target_name in module_name for target_name in self.lora_target_modules
         )
 
+    def _get_target_lora_layers(
+        self, target: str
+    ) -> tuple[list[tuple[str, dict[str, BaseLayerWithLoRA]]], str | None]:
+        """
+        Return a list of (module_name, lora_layers_dict) based on the target.
+
+        Args:
+            target: One of "all", "transformer", "transformer_2", "critic".
+
+        Returns:
+            A tuple of (result, error_message):
+            - result: List of tuples (module_name, lora_layers_dict) to operate on.
+            - error_message: Error description if target is invalid or module doesn't exist, None otherwise.
+        """
+        if target == "all":
+            result: list[tuple[str, dict[str, BaseLayerWithLoRA]]] = [
+                ("transformer", self.lora_layers)
+            ]
+            if self.lora_layers_transformer_2:
+                result.append(("transformer_2", self.lora_layers_transformer_2))
+            if self.lora_layers_critic:
+                result.append(("critic", self.lora_layers_critic))
+            return result, None
+        elif target == "transformer":
+            return [("transformer", self.lora_layers)], None
+        elif target == "transformer_2":
+            if not self.lora_layers_transformer_2:
+                return [], "transformer_2 does not exist in this pipeline"
+            return [("transformer_2", self.lora_layers_transformer_2)], None
+        elif target == "critic":
+            if not self.lora_layers_critic:
+                return (
+                    [],
+                    "critic (fake_score_transformer) does not exist in this pipeline",
+                )
+            return [("critic", self.lora_layers_critic)], None
+        else:
+            return [], f"Invalid target: {target}. Valid targets: {self.VALID_TARGETS}"
+
     def convert_module_lora_layers(
         self,
         module: torch.nn.Module,
@@ -210,11 +268,29 @@ class LoRAPipeline(ComposedPipelineBase):
                 layer.disable_lora = True
         return adapted_count
 
-    def is_lora_effective(self):
-        return self.is_lora_merged
+    def is_lora_effective(self, target: str = "all") -> bool:
+        """
+        Check if LoRA is currently effective (merged) for the specified target.
 
-    def is_lora_set(self):
-        return self.lora_initialized and self.cur_adapter_name is not None
+        Args:
+            target: Which transformer to check. "all" returns True if any is merged.
+        """
+        if target == "all":
+            return any(self.is_lora_merged.values())
+        return self.is_lora_merged.get(target, False)
+
+    def is_lora_set(self, target: str = "all") -> bool:
+        """
+        Check if LoRA has been set for the specified target.
+
+        Args:
+            target: Which transformer to check. "all" returns True if any is set.
+        """
+        if not self.lora_initialized:
+            return False
+        if target == "all":
+            return bool(self.cur_adapter_name)
+        return target in self.cur_adapter_name
 
     def load_lora_adapter(self, lora_path: str, lora_nickname: str, rank: int):
         """
@@ -266,25 +342,43 @@ class LoRAPipeline(ComposedPipelineBase):
                     f"Dit target weight name {target_name} already exists in lora_adapters[{lora_nickname}]"
                 )
             self.lora_adapters[lora_nickname][target_name] = weight.to(self.device)
-        self.cur_adapter_path = lora_path
         self.loaded_adapter_paths[lora_nickname] = lora_path
         logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path)
 
     def set_lora(
-        self, lora_nickname: str, lora_path: str | None = None
+        self, lora_nickname: str, lora_path: str | None = None, target: str = "all"
     ):  # type: ignore
         """
-        Load a LoRA adapter into the pipeline and merge it into the transformer.
+        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)
         """
-        if self.is_lora_merged and self.cur_adapter_name != lora_nickname:
+        if target not in self.VALID_TARGETS:
             raise ValueError(
-                f"LoRA '{self.cur_adapter_name}' is currently merged. "
-                "Please call 'unmerge_lora_weights' before setting a new LoRA."
+                f"Invalid target: {target}. Valid targets: {self.VALID_TARGETS}"
             )
 
+        # 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:
             raise ValueError(
                 f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it."
@@ -292,6 +386,11 @@ class LoRAPipeline(ComposedPipelineBase):
         if not self.lora_initialized:
             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)
+
         adapter_updated = False
         rank = dist.get_rank()
 
@@ -306,58 +405,113 @@ class LoRAPipeline(ComposedPipelineBase):
             adapter_updated = True
             self.load_lora_adapter(lora_path, lora_nickname, rank)
 
-        if (
+        # Check if we can skip (same adapter already applied to all target modules)
+        all_already_applied = all(
             not adapter_updated
-            and self.cur_adapter_name == lora_nickname
-            and self.is_lora_merged
-        ):
+            and self.cur_adapter_name.get(module_name) == lora_nickname
+            and self.is_lora_merged.get(module_name, False)
+            for module_name, _ in target_modules
+        )
+        if all_already_applied:
             return
-        self.cur_adapter_name = lora_nickname
 
-        # Merge the new adapter
-        adapted_count = self._apply_lora_to_layers(
-            self.lora_layers, lora_nickname, lora_path, rank
-        )
-        # Apply LoRA to transformer_2 if exists
-        adapted_count += self._apply_lora_to_layers(
-            self.lora_layers_transformer_2, lora_nickname, lora_path, rank
-        )
-        # Apply LoRA to fake_score_transformer (critic) if exists
-        adapted_count += self._apply_lora_to_layers(
-            self.lora_layers_critic, lora_nickname, lora_path, rank
-        )
+        # Apply LoRA to target modules
+        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
+            )
+            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.is_lora_merged = True
         logger.info(
-            "Rank %d: LoRA adapter %s applied to %d layers",
+            "Rank %d: LoRA adapter %s applied to %d layers (target: %s)",
             rank,
             lora_path,
             adapted_count,
+            target,
         )
 
-    def merge_lora_weights(self) -> None:
-        if self.is_lora_merged:
-            logger.warning("LoRA weights are already merged")
+    def merge_lora_weights(self, target: str = "all") -> None:
+        """
+        Merge LoRA weights into the base model for the specified target.
+
+        This operation is idempotent - calling it when LoRA is already merged is safe.
+
+        Args:
+            target: Which transformer(s) to merge. One of "all", "transformer",
+                    "transformer_2", "critic".
+        """
+        target_modules, error = self._get_target_lora_layers(target)
+        if error:
+            logger.warning("merge_lora_weights: %s", error)
+        if not target_modules:
             return
 
-        for name, layer in self.lora_layers.items():
-            layer.merge_lora_weights()
-        for name, layer in self.lora_layers_transformer_2.items():
-            layer.merge_lora_weights()
-        for name, layer in self.lora_layers_critic.items():
-            layer.merge_lora_weights()
-        logger.info("LoRA weights merged")
-        self.is_lora_merged = True
+        for module_name, lora_layers_dict in target_modules:
+            if self.is_lora_merged.get(module_name, False):
+                logger.warning("LoRA weights are already merged for %s", module_name)
+                continue
+            for name, layer in lora_layers_dict.items():
+                # Only re-enable LoRA for layers that actually have LoRA weights
+                has_lora_weights = hasattr(layer, "lora_A") and layer.lora_A is not None
+                if not has_lora_weights:
+                    continue
+                if hasattr(layer, "disable_lora"):
+                    layer.disable_lora = False
+                try:
+                    layer.merge_lora_weights()
+                except Exception as e:
+                    logger.warning("Could not merge layer %s: %s", name, e)
+                    continue
+            self.is_lora_merged[module_name] = True
+            logger.info("LoRA weights merged for %s", module_name)
 
-    def unmerge_lora_weights(self) -> None:
-        if not self.is_lora_merged:
-            logger.warning("LoRA weights are not merged.")
+    def unmerge_lora_weights(self, target: str = "all") -> None:
+        """
+        Unmerge LoRA weights from the base model for the specified target.
+        This also disables LoRA so it won't be computed on-the-fly.
+
+        This operation is idempotent - calling it when LoRA is not merged is safe.
+
+        Args:
+            target: Which transformer(s) to unmerge. One of "all", "transformer",
+                    "transformer_2", "critic".
+        """
+        target_modules, error = self._get_target_lora_layers(target)
+        if error:
+            logger.warning("unmerge_lora_weights: %s", error)
+        if not target_modules:
             return
 
-        for name, layer in self.lora_layers.items():
-            layer.unmerge_lora_weights()
-        for name, layer in self.lora_layers_transformer_2.items():
-            layer.unmerge_lora_weights()
-        for name, layer in self.lora_layers_critic.items():
-            layer.unmerge_lora_weights()
-        self.is_lora_merged = False
+        for module_name, lora_layers_dict in target_modules:
+            if not self.is_lora_merged.get(module_name, False):
+                logger.warning(
+                    "LoRA weights are not merged for %s, skipping", module_name
+                )
+                continue
+            for name, layer in lora_layers_dict.items():
+                # Check layer-level state to avoid raising exception
+                if hasattr(layer, "merged") and not layer.merged:
+                    logger.warning("Layer %s is not merged, skipping", name)
+                    # Still disable LoRA to prevent on-the-fly computation
+                    if hasattr(layer, "disable_lora"):
+                        layer.disable_lora = True
+                    continue
+                try:
+                    layer.unmerge_lora_weights()
+                    # Disable LoRA after unmerge to prevent on-the-fly computation
+                    if hasattr(layer, "disable_lora"):
+                        layer.disable_lora = True
+                except ValueError as e:
+                    logger.warning("Could not unmerge layer %s: %s", name, e)
+                    # Still disable LoRA even if unmerge failed
+                    if hasattr(layer, "disable_lora"):
+                        layer.disable_lora = True
+                    continue
+            self.is_lora_merged[module_name] = False
+            logger.info("LoRA weights unmerged for %s", module_name)