[diffusion] fix: make lora compatible with layerwise-offload (#16298)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -43,7 +43,6 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
     get_hf_config,
 )
-from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
@@ -738,15 +737,6 @@ class TransformerLoader(ComponentLoader):
 
         model = model.eval()
 
-        if server_args.dit_layerwise_offload:
-            # enable layerwise offload if possible
-            if isinstance(model, OffloadableDiTMixin):
-                model.configure_layerwise_offload(server_args)
-            else:
-                logger.info(
-                    "Disabling layerwise offload since current model does not support this feature"
-                )
-
         return model
 
 

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

@@ -28,6 +28,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
 from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import (
     configure_logger,
     globally_suppress_loggers,
@@ -88,6 +89,24 @@ class GPUWorker:
 
         self.pipeline = build_pipeline(self.server_args)
 
+        # apply layerwise offload after lora is applied while building LoRAPipeline
+        # otherwise empty offloaded weights could fail lora converting
+        if self.server_args.dit_layerwise_offload:
+            # enable layerwise offload if possible
+            for dit in filter(
+                None,
+                [
+                    self.pipeline.get_module("transformer"),
+                    self.pipeline.get_module("transformer_2"),
+                ],
+            ):
+                if isinstance(dit, OffloadableDiTMixin):
+                    dit.configure_layerwise_offload(self.server_args)
+                else:
+                    logger.info(
+                        f"Module {type(dit).__name__} does not support layerwise offload. Skipping."
+                    )
+
         logger.info(
             f"Worker {self.rank}: Initialized device, model, and distributed environment."
         )