[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/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",