From 198c8ecf98402080cb0e89023f09f8153fde56bb Mon Sep 17 00:00:00 2001 From: Prozac614 Date: Fri, 12 Dec 2025 13:08:57 +0800 Subject: [PATCH] [diffusion] fix: support applying different LoRA adapters to different transformers in multi-transformer pipelines (#14839) --- .../entrypoints/diffusion_generator.py | 83 +++-- .../runtime/entrypoints/openai/common_api.py | 51 ++- .../runtime/entrypoints/openai/utils.py | 5 +- .../runtime/layers/lora/linear.py | 3 +- .../runtime/managers/gpu_worker.py | 25 +- .../runtime/managers/scheduler.py | 12 +- .../runtime/pipelines_core/lora_pipeline.py | 290 ++++++++++++++---- 7 files changed, 349 insertions(+), 120 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 3dae916de..83b4409ee 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/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, 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 2dcbeb77c..2f1501672 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py +++ b/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", ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index ec7a3cb1a..8607f06ad 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index 898b90d4f..c507e9ea2 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 7fa2014f5..0482af743 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 35da11dcf..e2fe164fb 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/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]): 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 1e6876f26..a9270ae8a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/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)