[diffusion] fix: fix LoRA weight merging when using layerwise offload (#16737)

Co-authored-by: niehen6174 <niehen.6174@gmail.com>
Co-authored-by: DavisTao <dwt614707404@163.com>
Co-authored-by: niehen6174 <nihen6174@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

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

@@ -199,13 +199,19 @@ class BaseLayerWithLoRA(nn.Module):
         # avoid precision loss
         if isinstance(self.base_layer.weight, DTensor):
             device = self.base_layer.weight.data.device
-            self.base_layer.weight = nn.Parameter(
+            old_weight = self.base_layer.weight
-                self.cpu_weight.to(device, non_blocking=True)
+            new_weight_data = self.cpu_weight.to(device, non_blocking=True)
-            )
+            self.base_layer.weight = nn.Parameter(new_weight_data)
+            del old_weight
         else:
-            self.base_layer.weight.data = self.cpu_weight.data.to(
+            current_device = self.base_layer.weight.data.device
-                self.base_layer.weight, non_blocking=True
+            cpu_weight_on_device = self.cpu_weight.to(current_device, non_blocking=True)
-            )
+            self.base_layer.weight.data.copy_(cpu_weight_on_device)
+            if (
+                cpu_weight_on_device.data_ptr()
+                != self.base_layer.weight.data.data_ptr()
+            ):
+                del cpu_weight_on_device
 
         self.merged = False
 

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

@@ -4,6 +4,7 @@
 import os
 from collections import defaultdict
 from collections.abc import Hashable
+from contextlib import contextmanager
 from typing import Any
 
 import torch
@@ -143,6 +144,71 @@ class LoRAPipeline(ComposedPipelineBase):
         else:
             return [], f"Invalid target: {target}. Valid targets: {self.VALID_TARGETS}"
 
+    @contextmanager
+    def _temporarily_disable_offload(
+        self,
+        target_modules: list[tuple[str, dict[str, BaseLayerWithLoRA]]] | None = None,
+        target: str | None = None,
+        use_module_names_only: bool = False,
+    ):
+        """
+        Context manager to temporarily disable layerwise offload for the given modules.
+
+        Args:
+            target_modules: List of (module_name, lora_layers_dict) tuples. If None, will be determined from target.
+            target: Target string ("all", "transformer", etc.). Used if target_modules is None.
+            use_module_names_only: If True, determine module names directly from target without requiring
+                                   LoRA initialization. Used for convert_to_lora_layers scenario.
+
+        Yields:
+            List of modules that had offload disabled.
+        """
+        from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
+            OffloadableDiTMixin,
+        )
+
+        module_names = []
+        if target_modules is not None:
+            # Extract module names from target_modules
+            module_names = [module_name for module_name, _ in target_modules]
+        elif target is not None:
+            if use_module_names_only:
+                if target == "all":
+                    module_names = ["transformer", "transformer_2"]
+                elif target in ["transformer", "transformer_2", "critic"]:
+                    module_names = [target]
+            else:
+                target_modules, _ = self._get_target_lora_layers(target)
+                if target_modules:
+                    module_names = [module_name for module_name, _ in target_modules]
+        else:
+            yield []
+            return
+
+        if not module_names:
+            yield []
+            return
+
+        # clear CUDA cache to free up unused memory
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+
+        offload_disabled_modules = []
+        for module_name in module_names:
+            module = self.modules.get(module_name)
+            if module is not None and isinstance(module, OffloadableDiTMixin):
+                if module.layerwise_offload_managers is not None:
+                    module.disable_offload()
+                    offload_disabled_modules.append(module)
+
+        try:
+            yield offload_disabled_modules
+        finally:
+            # Re-enable layerwise offload: sync weights to CPU and restore hooks
+            for module in offload_disabled_modules:
+                module.enable_offload()
+
     def convert_module_lora_layers(
         self,
         module: torch.nn.Module,
@@ -183,6 +249,7 @@ class LoRAPipeline(ComposedPipelineBase):
                 target_lora_layers[name] = lora_layer
                 replace_submodule(self.modules[module_name], name, lora_layer)
                 converted_count += 1
+
         return converted_count
 
     def convert_to_lora_layers(self) -> None:
@@ -412,7 +479,14 @@ class LoRAPipeline(ComposedPipelineBase):
             raise ValueError(
                 f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it."
             )
+
+        # Disable layerwise offload before convert_to_lora_layers to ensure weights are accessible
+        # This is critical because convert_to_lora_layers needs to save cpu_weight from actual weights,
+        # not from offloaded placeholder tensors
         if not self.lora_initialized:
+            with self._temporarily_disable_offload(
+                target="all", use_module_names_only=True
+            ):
                 self.convert_to_lora_layers()
 
         # Re-fetch target_modules after convert_to_lora_layers() to get populated dicts
@@ -445,11 +519,17 @@ class LoRAPipeline(ComposedPipelineBase):
         if all_already_applied:
             return
 
-        # Apply LoRA to target modules
+        # 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
+                    lora_layers_dict,
+                    lora_nickname,
+                    lora_path,
+                    rank,
+                    strength,
                 )
                 adapted_count += count
                 self.cur_adapter_name[module_name] = lora_nickname
@@ -485,6 +565,8 @@ class LoRAPipeline(ComposedPipelineBase):
         if not target_modules:
             return
 
+        # Disable layerwise offload if enabled: load all layers to GPU
+        with self._temporarily_disable_offload(target_modules=target_modules):
             for module_name, lora_layers_dict in target_modules:
                 if self.is_lora_merged.get(module_name, False):
                     # Check if strength is the same - if so, skip (idempotent)
@@ -502,7 +584,9 @@ class LoRAPipeline(ComposedPipelineBase):
                     )
                 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
+                    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"):
@@ -535,12 +619,15 @@ class LoRAPipeline(ComposedPipelineBase):
         if not target_modules:
             return
 
+        # Disable layerwise offload if enabled: load all layers to GPU
+
         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
+            with self._temporarily_disable_offload(target_modules=target_modules):
                 for name, layer in lora_layers_dict.items():
                     # Check layer-level state to avoid raising exception
                     if hasattr(layer, "merged") and not layer.merged:

python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py

@@ -60,6 +60,8 @@ class LayerwiseOffloadManager:
 
         self._named_parameters: Dict[str, torch.nn.Parameter] = {}
         self._named_buffers: Dict[str, torch.Tensor] = {}
+        # Store forward hooks for removal
+        self._forward_hooks: List[Any] = []
 
         self._initialize()
 
@@ -204,6 +206,51 @@ class LayerwiseOffloadManager:
         for layer_idx in list(self._gpu_layers):
             self.release_layer(layer_idx)
 
+    @torch.compiler.disable
+    def load_all_layers(self) -> None:
+        """Load all layers from CPU to GPU."""
+        if not self.enabled or self.device is None:
+            return
+        if self.copy_stream is not None:
+            torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+        for layer_idx in range(self.num_layers):
+            if layer_idx not in self._gpu_layers:
+                self.prefetch_layer(layer_idx, non_blocking=False)
+
+    @torch.compiler.disable
+    def sync_layer_to_cpu(self, layer_idx: int) -> None:
+        """Sync a layer's weights from GPU back to CPU."""
+        if not self.enabled or layer_idx not in self._gpu_layers:
+            return
+        if layer_idx not in self._consolidated_cpu_weights:
+            return
+
+        if self.copy_stream is not None:
+            torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+        # Collect current GPU weights and write back to CPU buffer
+        for name, meta in self._weight_metadata.get(layer_idx, {}).items():
+            target = self.get_target_with_name(name)
+            gpu_weight = target.data.flatten().cpu()
+
+            dtype = meta["dtype"]
+            cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype]
+            offset = meta["offset"]
+            numel = meta["numel"]
+            cpu_buffer[offset : offset + numel].copy_(gpu_weight)
+
+    @torch.compiler.disable
+    def sync_all_layers_to_cpu(self) -> None:
+        """Sync all loaded layers' weights from GPU back to CPU."""
+        if not self.enabled or self.device is None:
+            return
+        if self.copy_stream is not None:
+            torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+        for layer_idx in list(self._gpu_layers):
+            self.sync_layer_to_cpu(layer_idx)
+
     def register_forward_hooks(self) -> None:
         if not self.enabled:
             return
@@ -225,9 +272,17 @@ class LayerwiseOffloadManager:
             return hook
 
         # register prefetch & release hooks for each layer
+        self._forward_hooks.clear()
         for i, layer in enumerate(layers):
-            layer.register_forward_pre_hook(make_pre_hook(i))
+            pre_hook_handle = layer.register_forward_pre_hook(make_pre_hook(i))
-            layer.register_forward_hook(make_post_hook(i))
+            post_hook_handle = layer.register_forward_hook(make_post_hook(i))
+            self._forward_hooks.extend([pre_hook_handle, post_hook_handle])
+
+    def remove_forward_hooks(self) -> None:
+        """Remove all registered forward hooks."""
+        for hook_handle in self._forward_hooks:
+            hook_handle.remove()
+        self._forward_hooks.clear()
 
 
 class OffloadableDiTMixin:
@@ -266,3 +321,24 @@ class OffloadableDiTMixin:
             return
         for manager in self.layerwise_offload_managers:
             manager.prepare_for_next_denoise(non_blocking=True)
+
+    def disable_offload(self) -> None:
+        """Disable layerwise offload: load all layers to GPU and remove hooks."""
+        if self.layerwise_offload_managers is None:
+            return
+        for manager in self.layerwise_offload_managers:
+            if manager.enabled:
+                manager.remove_forward_hooks()
+                manager.load_all_layers()
+
+    def enable_offload(self) -> None:
+        """Re-enable layerwise offload: sync weights to CPU, release layers, and restore hooks."""
+        if self.layerwise_offload_managers is None:
+            return
+        for manager in self.layerwise_offload_managers:
+            if manager.enabled:
+                manager.sync_all_layers_to_cpu()
+                for layer_idx in list(manager._gpu_layers):
+                    if layer_idx > 0:
+                        manager.release_layer(layer_idx)
+                manager.register_forward_hooks()

python/sglang/multimodal_gen/test/server/test_server_common.py

@@ -528,6 +528,34 @@ Consider updating perf_baselines.json with the snippets below:
             "[LoRA Switch E2E] All dynamic switch E2E tests passed for %s", case.id
         )
 
+    def _test_dynamic_lora_loading(
+        self,
+        ctx: ServerContext,
+        case: DiffusionTestCase,
+    ) -> None:
+        """
+        Test dynamic LoRA loading after server startup.
+
+        This test reproduces the LayerwiseOffload + set_lora issue:
+        - Server starts WITHOUT lora_path (LayerwiseOffloadManager initializes first)
+        - Then set_lora is called via API to load LoRA dynamically
+        - This tests the interaction between layerwise offload and dynamic LoRA loading
+        """
+        base_url = f"http://localhost:{ctx.port}/v1"
+        dynamic_lora_path = case.server_args.dynamic_lora_path
+
+        # Call set_lora to load LoRA dynamically after server startup
+        logger.info(
+            "[Dynamic LoRA] Loading LoRA dynamically via set_lora API for %s", case.id
+        )
+        logger.info("[Dynamic LoRA] LoRA path: %s", dynamic_lora_path)
+        resp = requests.post(
+            f"{base_url}/set_lora",
+            json={"lora_nickname": "default", "lora_path": dynamic_lora_path},
+        )
+        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_v1_models_endpoint(
         self, ctx: ServerContext, case: DiffusionTestCase
     ) -> None:
@@ -629,6 +657,11 @@ Consider updating perf_baselines.json with the snippets below:
         - test_diffusion_perf[qwen_image_edit]
         - etc.
         """
+        # Dynamic LoRA loading test - tests LayerwiseOffload + set_lora interaction
+        # Server starts WITHOUT lora_path, then set_lora is called after startup
+        if case.server_args.dynamic_lora_path:
+            self._test_dynamic_lora_loading(diffusion_server, case)
+
         generate_fn = get_generate_fn(
             model_path=case.server_args.model_path,
             modality=case.server_args.modality,
@@ -646,5 +679,5 @@ Consider updating perf_baselines.json with the snippets below:
         self._test_v1_models_endpoint(diffusion_server, case)
 
         # LoRA API functionality test with E2E validation (only for LoRA-enabled cases)
-        if case.server_args.lora_path:
+        if case.server_args.lora_path or case.server_args.dynamic_lora_path:
             self._test_lora_api_functionality(diffusion_server, case, generate_fn)

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -149,7 +149,12 @@ class DiffusionServerArgs:
     ulysses_degree: int | None = None
     ring_degree: int | None = None
     # LoRA
-    lora_path: str | None = None  # LoRA adapter path (HF repo or local path)
+    lora_path: str | None = (
+        None  # LoRA adapter path (HF repo or local path, loaded at startup)
+    )
+    dynamic_lora_path: str | None = (
+        None  # LoRA path for dynamic loading test (loaded via set_lora after startup)
+    )
     # misc
     enable_warmup: bool = False
 
@@ -406,6 +411,8 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
         ),
     ),
     # LoRA test case for single transformer + merge/unmerge API test
+    # Note: Uses dynamic_lora_path instead of lora_path to test LayerwiseOffload + set_lora interaction
+    # Server starts WITHOUT LoRA, then set_lora is called after startup (Wan models auto-enable layerwise offload)
     DiffusionTestCase(
         "wan2_1_t2v_1_3b_lora_1gpu",
         DiffusionServerArgs(
@@ -414,7 +421,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
             warmup=0,
             custom_validator="video",
             num_gpus=1,
-            lora_path="Cseti/Wan-LoRA-Arcane-Jinx-v1",
+            dynamic_lora_path="Cseti/Wan-LoRA-Arcane-Jinx-v1",
         ),
         DiffusionSamplingParams(
             prompt="csetiarcane Nfj1nx with blue hair, a woman walking in a cyberpunk city at night",