From 8999ce754fac64ad9ff7accbbf4299f9060206db Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sat, 20 Dec 2025 22:52:40 +0800 Subject: [PATCH] [diffusion] perf: support zero-cost weight offload and overlap with compute for wan-series (#15511) --- .../runtime/distributed/parallel_state.py | 49 ++++- .../runtime/loader/component_loader.py | 20 ++ .../runtime/models/dits/wanvideo.py | 23 +- .../pipelines_core/stages/denoising.py | 11 + .../multimodal_gen/runtime/server_args.py | 28 +++ .../runtime/utils/layerwise_offload.py | 197 ++++++++++++++++++ 6 files changed, 322 insertions(+), 6 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index ce965f1a7..683b75027 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -32,6 +32,7 @@ If you only need to use the distributed environment without model parallelism, """ import contextlib import os +import time import weakref from collections import namedtuple from collections.abc import Callable @@ -67,8 +68,6 @@ _DP: Optional[GroupCoordinator] = None _DIT: Optional[GroupCoordinator] = None _VAE: Optional[GroupCoordinator] = None -logger = init_logger(__name__) - TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) @@ -433,6 +432,9 @@ def initialize_model_parallel( ring_group=PROCESS_GROUP.RING_PG, ) + if ulysses_degree > 1: + _warmup_ulysses_communication() + global _TP assert _TP is None, "Tensor parallel group is already initialized" _TP = init_parallel_group_coordinator( @@ -945,6 +947,49 @@ def get_ring_parallel_rank(): return get_sp_group().ring_rank +def _warmup_ulysses_communication(): + """ + Warmup NCCL communication for Ulysses all-to-all to avoid first-step latency. + + This function performs a dummy all-to-all operation to initialize NCCL communication + channels, which can take several seconds on the first call. + """ + logger.info("Warming up Ulysses all-to-all communication...") + + try: + import torch.distributed._functional_collectives as ft_c + + ulysses_pg = get_sp_group().ulysses_group + if ulysses_pg is None: + logger.warning("Ulysses group not initialized, skipping warmup") + return + + warmup_start = time.time() + + device = torch.device(f"cuda:{get_world_group().local_rank}") + dummy_tensor = torch.zeros(1024, device=device, dtype=torch.float32) + + output = ft_c.all_to_all_single( + dummy_tensor, + output_split_sizes=None, + input_split_sizes=None, + group=ulysses_pg, + ) + + if isinstance(output, ft_c.AsyncCollectiveTensor): + output = output.wait() + + torch.cuda.synchronize() + + warmup_time = (time.time() - warmup_start) * 1000 + logger.info(f"Ulysses communication warmup completed in {warmup_time:.2f}ms") + + except Exception as e: + logger.warning( + f"Ulysses communication warmup failed: {e}. Continuing without warmup." + ) + + # PP def get_pp_group() -> PipelineGroupCoordinator: assert _PP is not None, "pipeline model parallel group is not initialized" diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 923a4d104..9bc5aa27f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -43,6 +43,9 @@ 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 ( + LayerwiseOffloadManager, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE @@ -709,6 +712,23 @@ class TransformerLoader(ComponentLoader): model = model.eval() + if server_args.dit_layerwise_offload and hasattr(model, "blocks"): + try: + num_layers = len(getattr(model, "blocks")) + except Exception: + num_layers = None + + if isinstance(num_layers, int) and num_layers > 0: + mgr = LayerwiseOffloadManager( + model, + module_list_attr="blocks", + num_layers=num_layers, + enabled=True, + pin_cpu_memory=server_args.pin_cpu_memory, + auto_initialize=True, + ) + setattr(model, "_layerwise_offload_manager", mgr) + return model diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index cb674e491..97044275f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -792,10 +792,25 @@ class WanTransformer3DModel(CachableDiT): if enable_teacache: original_hidden_states = hidden_states.clone() - for block in self.blocks: - hidden_states = block( - hidden_states, encoder_hidden_states, timestep_proj, freqs_cis - ) + offload_mgr = getattr(self, "_layerwise_offload_manager", None) + if offload_mgr is not None and getattr(offload_mgr, "enabled", False): + for i, block in enumerate(self.blocks): + with offload_mgr.layer_scope( + prefetch_layer_idx=i + 1, + release_layer_idx=i, + non_blocking=True, + ): + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + ) + else: + for block in self.blocks: + hidden_states = block( + hidden_states, encoder_hidden_states, timestep_proj, freqs_cis + ) # if teacache is enabled, we need to cache the original hidden states if enable_teacache: self.maybe_cache_states(hidden_states, original_hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 87e0bcf65..96a484b40 100755 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -728,6 +728,17 @@ class DenoisingStage(PipelineStage): latents, batch ) + offload_mgr = getattr(self.transformer, "_layerwise_offload_manager", None) + if offload_mgr is not None and getattr(offload_mgr, "enabled", False): + offload_mgr.release_all() + + if self.transformer_2 is not None: + offload_mgr_2 = getattr( + self.transformer_2, "_layerwise_offload_manager", None + ) + if offload_mgr_2 is not None and getattr(offload_mgr_2, "enabled", False): + offload_mgr_2.release_all() + # Save STA mask search results if needed if ( st_attn_available diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 28e23056e..9b5b24e6e 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -229,6 +229,7 @@ class ServerArgs: # CPU offload parameters dit_cpu_offload: bool = True use_fsdp_inference: bool = False + dit_layerwise_offload: bool = False text_encoder_cpu_offload: bool = True image_encoder_cpu_offload: bool = True vae_cpu_offload: bool = True @@ -507,6 +508,13 @@ class ServerArgs: action=StoreBoolean, help="Use CPU offload for DiT inference. Enable if run out of memory with FSDP.", ) + parser.add_argument( + "--dit-layerwise-offload", + action=StoreBoolean, + default=ServerArgs.dit_layerwise_offload, + help="Enable layerwise CPU offload with async H2D prefetch overlap for supported DiT models (e.g., Wan). " + "Cannot be used together with cache-dit (SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.", + ) parser.add_argument( "--use-fsdp-inference", action=StoreBoolean, @@ -844,6 +852,26 @@ class ServerArgs: """Validate inference arguments for consistency""" if current_platform.is_mps(): self.use_fsdp_inference = False + self.dit_layerwise_offload = False + + if self.dit_layerwise_offload: + if self.use_fsdp_inference: + logger.warning( + "dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference." + ) + self.use_fsdp_inference = False + if self.dit_cpu_offload: + logger.warning( + "dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload." + ) + self.dit_cpu_offload = False + if os.getenv("SGLANG_CACHE_DIT_ENABLED", "").lower() == "true": + raise ValueError( + "dit_layerwise_offload cannot be enabled together with cache-dit. " + "cache-dit may reuse skipped blocks whose weights have been released by layerwise offload, " + "causing shape mismatch errors. " + "Please disable either --dit-layerwise-offload or SGLANG_CACHE_DIT_ENABLED." + ) # autocast if self.disable_autocast is None: diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py new file mode 100644 index 000000000..ffedc3a34 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py @@ -0,0 +1,197 @@ +import re +from contextlib import contextmanager +from typing import Dict, Set + +import torch + + +# Adapted from skywork AI Infra diffusion optimize +class LayerwiseOffloadManager: + """A lightweight layerwise CPU offload manager. + + This utility offloads per-layer parameters/buffers from GPU to CPU, and + supports async H2D prefetch using a dedicated CUDA stream. + + Typical usage: + - Construct the manager with the target model and the list-like module + attribute that represents transformer blocks (e.g. ``blocks``). + - Call :meth:`initialize` once to offload weights and prefetch layer 0. + - During forward, call :meth:`prefetch_layer` for the next layer and + :meth:`release_layer` for the finished layer. + """ + + def __init__( + self, + model: torch.nn.Module, + *, + module_list_attr: str, + num_layers: int, + enabled: bool, + pin_cpu_memory: bool = True, + auto_initialize: bool = False, + ) -> None: + self.model = model + self.module_list_attr = module_list_attr + self.num_layers = num_layers + self.pin_cpu_memory = pin_cpu_memory + + self.enabled = bool(enabled and torch.cuda.is_available()) + self.device = ( + torch.device("cuda", torch.cuda.current_device()) if self.enabled else None + ) + self.copy_stream = torch.cuda.Stream() if self.enabled else None + + self._layer_name_re = re.compile( + rf"(^|\.){re.escape(module_list_attr)}\.(\d+)(\.|$)" + ) + + self._cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {} + self._cpu_dtypes: Dict[int, Dict[str, torch.dtype]] = {} + self._gpu_layers: Dict[int, Set[str]] = {} + + self._named_parameters: Dict[str, torch.nn.Parameter] = {} + self._named_buffers: Dict[str, torch.Tensor] = {} + + if auto_initialize: + self.initialize() + + def _match_layer_idx(self, name: str) -> int | None: + m = self._layer_name_re.search(name) + if not m: + return None + try: + return int(m.group(2)) + except Exception: + return None + + def _offload_tensor(self, name: str, tensor: torch.Tensor, layer_idx: int) -> None: + if layer_idx not in self._cpu_weights: + self._cpu_weights[layer_idx] = {} + self._cpu_dtypes[layer_idx] = {} + + cpu_weight = tensor.detach().to("cpu") + if self.pin_cpu_memory: + cpu_weight = cpu_weight.pin_memory() + self._cpu_weights[layer_idx][name] = cpu_weight + self._cpu_dtypes[layer_idx][name] = tensor.dtype + + if self.device is not None: + tensor.data = torch.empty((1,), device=self.device, dtype=tensor.dtype) + + @torch.compiler.disable + def initialize(self) -> None: + if not self.enabled: + return + + self._named_parameters = dict(self.model.named_parameters()) + self._named_buffers = dict(self.model.named_buffers()) + + for name, param in self._named_parameters.items(): + layer_idx = self._match_layer_idx(name) + if layer_idx is None or layer_idx >= self.num_layers: + continue + self._offload_tensor(name, param, layer_idx) + + for name, buf in self._named_buffers.items(): + layer_idx = self._match_layer_idx(name) + if layer_idx is None or layer_idx >= self.num_layers: + continue + self._offload_tensor(name, buf, layer_idx) + + self.prefetch_layer(0, non_blocking=False) + if self.copy_stream is not None: + torch.cuda.current_stream().wait_stream(self.copy_stream) + + @torch.compiler.disable + def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None: + if not self.enabled or self.device is None or self.copy_stream is None: + return + if layer_idx < 0 or layer_idx >= self.num_layers: + return + if layer_idx in self._gpu_layers: + return + if layer_idx not in self._cpu_weights: + return + + self.copy_stream.wait_stream(torch.cuda.current_stream()) + + param_names: Set[str] = set() + + for name, cpu_weight in self._cpu_weights[layer_idx].items(): + if name in self._named_parameters: + target = self._named_parameters[name] + else: + target = self._named_buffers[name] + + gpu_weight = torch.empty( + cpu_weight.shape, + dtype=self._cpu_dtypes[layer_idx][name], + device=self.device, + ) + with torch.cuda.stream(self.copy_stream): + gpu_weight.copy_(cpu_weight, non_blocking=non_blocking) + target.data = gpu_weight + param_names.add(name) + + self._gpu_layers[layer_idx] = param_names + + @contextmanager + def layer_scope( + self, + *, + prefetch_layer_idx: int | None, + release_layer_idx: int | None, + non_blocking: bool = True, + ): + """A helper context manager to improve readability at call sites. + + It optionally prefetches ``prefetch_layer_idx`` before entering the + context, and waits for the copy stream then releases + ``release_layer_idx`` on exit. + """ + if self.enabled and prefetch_layer_idx is not None: + self.prefetch_layer(prefetch_layer_idx, non_blocking=non_blocking) + try: + yield + finally: + if self.enabled and self.copy_stream is not None: + torch.cuda.current_stream().wait_stream(self.copy_stream) + if self.enabled and release_layer_idx is not None: + self.release_layer(release_layer_idx) + + @torch.compiler.disable + def release_layer(self, layer_idx: int) -> None: + if not self.enabled or self.device is None: + return + if layer_idx <= 0: + return + + param_names = self._gpu_layers.pop(layer_idx, None) + if not param_names: + return + + for name in param_names: + if name in self._named_parameters: + target = self._named_parameters[name] + else: + target = self._named_buffers[name] + target.data = torch.empty((1,), device=self.device, dtype=target.dtype) + + @torch.compiler.disable + def release_all(self) -> None: + 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) + + layer_indices = list(self._gpu_layers.keys()) + for layer_idx in layer_indices: + param_names = self._gpu_layers.pop(layer_idx, None) + if not param_names: + continue + for name in param_names: + if name in self._named_parameters: + target = self._named_parameters[name] + else: + target = self._named_buffers[name] + target.data = torch.empty((1,), device=self.device, dtype=target.dtype)