diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index ff45b914b..2f9de9882 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/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( - self.cpu_weight.to(device, non_blocking=True) - ) + old_weight = self.base_layer.weight + 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( - self.base_layer.weight, non_blocking=True - ) + current_device = self.base_layer.weight.data.device + 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 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 f611cf7f9..929be44af 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/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,8 +479,15 @@ 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: - self.convert_to_lora_layers() + 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 target_modules, error = self._get_target_lora_layers(target) @@ -445,28 +519,34 @@ class LoRAPipeline(ComposedPipelineBase): if all_already_applied: return - # 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, strength - ) - 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.cur_adapter_strength[module_name] = strength + # 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, + ) + 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.cur_adapter_strength[module_name] = strength - logger.info( - "Rank %d: LoRA adapter %s applied to %d layers (target: %s, strength: %s)", - rank, - lora_path, - adapted_count, - target, - strength, - ) + logger.info( + "Rank %d: LoRA adapter %s applied to %d layers (target: %s, strength: %s)", + rank, + lora_path, + adapted_count, + target, + strength, + ) def merge_lora_weights(self, target: str = "all", strength: float = 1.0) -> None: """ @@ -485,38 +565,42 @@ class LoRAPipeline(ComposedPipelineBase): if not target_modules: return - 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) - if self.cur_adapter_strength.get(module_name) == strength: - logger.warning( - "LoRA weights are already merged for %s with same strength", + # 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) + if self.cur_adapter_strength.get(module_name) == strength: + logger.warning( + "LoRA weights are already merged for %s with same strength", + module_name, + ) + continue + # Different strength requested - allow re-merge (layer handles unmerge internally) + logger.info( + "Re-merging LoRA weights for %s with new strength %s", module_name, + strength, ) - continue - # Different strength requested - allow re-merge (layer handles unmerge internally) + 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(strength=strength) + except Exception as e: + logger.warning("Could not merge layer %s: %s", name, e) + continue + self.is_lora_merged[module_name] = True + self.cur_adapter_strength[module_name] = strength logger.info( - "Re-merging LoRA weights for %s with new strength %s", - module_name, - strength, + "LoRA weights merged for %s (strength: %s)", module_name, strength ) - 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(strength=strength) - except Exception as e: - logger.warning("Could not merge layer %s: %s", name, e) - continue - self.is_lora_merged[module_name] = True - self.cur_adapter_strength[module_name] = strength - logger.info( - "LoRA weights merged for %s (strength: %s)", module_name, strength - ) def unmerge_lora_weights(self, target: str = "all") -> None: """ @@ -535,34 +619,37 @@ 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 - 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 - self.cur_adapter_strength.pop(module_name, None) - logger.info("LoRA weights unmerged for %s", module_name) + 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: + 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 + self.cur_adapter_strength.pop(module_name, None) + logger.info("LoRA weights unmerged for %s", module_name) def get_lora_status(self) -> dict[str, Any]: """ diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py index 0cb14ce4d..36d83f472 100644 --- a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py +++ b/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)) - layer.register_forward_hook(make_post_hook(i)) + pre_hook_handle = layer.register_forward_pre_hook(make_pre_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() diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 76cec58b1..6dcdbaedc 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 083f2f947..4a20274d1 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/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",