diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 393f15ca9..5e2175be3 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index c024757e3..b31af0dee 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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." )