From 88fcd8535fe1ceff3d28bf7872ba73c0a025403c Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 28 Jan 2026 09:34:27 +0800 Subject: [PATCH] [diffusion] feat: add an arg for controlling the number of prefetched layers in layerwise-offload (#17693) --- .../pipelines_core/stages/denoising.py | 5 +- .../multimodal_gen/runtime/server_args.py | 26 +++++++ .../runtime/utils/layerwise_offload.py | 69 +++++++++++++++---- .../test/scripts/gen_perf_baselines.py | 25 ++++--- .../test/server/test_server_common.py | 5 ++ .../test/server/testcase_configs.py | 2 + 6 files changed, 106 insertions(+), 26 deletions(-) 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 d4bd11387..ce4c8a86f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -742,7 +742,10 @@ class DenoisingStage(PipelineStage): # reset offload managers with prefetching first layer for next forward for dit in filter(None, [self.transformer, self.transformer_2]): if isinstance(dit, OffloadableDiTMixin): - dit.prepare_for_next_denoise() + # release all DiT weights to avoid peak VRAM usage, which may increasing the latency for next req + # TODO: should be make this an option? + for manager in dit.layerwise_offload_managers: + manager.release_all() def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs): """Shard latents for Sequence Parallelism if applicable.""" diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index e56115157..6a6be2064 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -8,6 +8,7 @@ import argparse import dataclasses import inspect import json +import math import os import random import sys @@ -285,6 +286,7 @@ class ServerArgs: # CPU offload parameters dit_cpu_offload: bool | None = None dit_layerwise_offload: bool | None = None + dit_offload_prefetch_size: float = 0.0 text_encoder_cpu_offload: bool | None = None image_encoder_cpu_offload: bool | None = None vae_cpu_offload: bool | None = None @@ -617,6 +619,12 @@ class ServerArgs: 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( + "--dit-offload-prefetch-size", + type=float, + default=ServerArgs.dit_offload_prefetch_size, + help="The size of prefetch for dit-layerwise-offload. If the value is between 0.0 and 1.0, it is treated as a ratio of the total number of layers. If the value is >= 1, it is treated as the absolute number of layers. 0.0 means prefetch 1 layer (lowest memory). Values above 0.5 might have peak memory close to no offload but worse performance.", + ) parser.add_argument( "--use-fsdp-inference", action=StoreBoolean, @@ -949,6 +957,21 @@ class ServerArgs: self.use_fsdp_inference = False self.dit_layerwise_offload = False + if self.dit_offload_prefetch_size > 1 and ( + isinstance(self.dit_offload_prefetch_size, float) + and not self.dit_offload_prefetch_size.is_integer() + ): + self.dit_offload_prefetch_size = int( + math.floor(self.dit_offload_prefetch_size) + ) + logger.info( + f"Invalid --dit-offload-prefetch-size value passed, truncated to: {self.dit_offload_prefetch_size}" + ) + if 0.5 <= self.dit_offload_prefetch_size < 1.0: + logger.info( + f"We do not recommend --dit-offload-prefetch-size to be between 0.5 and 1.0" + ) + if not envs.SGLANG_CACHE_DIT_ENABLED: # TODO: need a better way to tell this if ( @@ -962,6 +985,9 @@ class ServerArgs: self.dit_layerwise_offload = True if self.dit_layerwise_offload: + assert ( + self.dit_offload_prefetch_size >= 0.0 + ), "dit_offload_prefetch_size must be non-negative" if self.use_fsdp_inference: logger.warning( "dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference." diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py index 36d83f472..8af6ad1a6 100644 --- a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py @@ -33,12 +33,13 @@ class LayerwiseOffloadManager: num_layers: int, enabled: bool, pin_cpu_memory: bool = True, + prefetch_size: int = 1, ) -> None: self.model = model self.layers_attr_str = layers_attr_str self.num_layers = num_layers self.pin_cpu_memory = pin_cpu_memory - + self.prefetch_size = min(max(1, prefetch_size), self.num_layers) self.enabled = bool(enabled and torch.cuda.is_available()) if not self.enabled: return @@ -57,6 +58,8 @@ class LayerwiseOffloadManager: self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {} # layer indices that are already in gpu self._gpu_layers: Set[int] = set() + # layer_idx -> torch.cuda.Event for fine-grained sync, to make sure the weight is resident in pre-hook + self._prefetch_events: Dict[int, torch.cuda.Event] = {} self._named_parameters: Dict[str, torch.nn.Parameter] = {} self._named_buffers: Dict[str, torch.Tensor] = {} @@ -127,13 +130,19 @@ class LayerwiseOffloadManager: self._consolidated_cpu_weights[layer_idx][dtype] = cpu_buffer # prefetch the first layer for warm-up - self.prepare_for_next_denoise(non_blocking=False) + self.prepare_for_next_req(non_blocking=False) self.register_forward_hooks() - logger.info("LayerwiseOffloadManager initialized") + logger.info( + f"LayerwiseOffloadManager initialized with num prefetched layer: {self.prefetch_size}, total num layers: {self.num_layers}" + ) - def prepare_for_next_denoise(self, non_blocking=True): - self.prefetch_layer(0, non_blocking=non_blocking) + def prepare_for_next_req(self, non_blocking=True): + """ + Prepare for the next round of denoising loop with prefetching the necessary layers + """ + for i in range(self.prefetch_size): + self.prefetch_layer(i, non_blocking=non_blocking) if not non_blocking and self.copy_stream is not None: torch.cuda.current_stream().wait_stream(self.copy_stream) @@ -147,6 +156,9 @@ class LayerwiseOffloadManager: @torch.compiler.disable def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None: + """ + idempotent + """ 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: @@ -167,6 +179,11 @@ class LayerwiseOffloadManager: gpu_buffer.copy_(cpu_buffer, non_blocking=non_blocking) gpu_buffers[dtype] = gpu_buffer + # record the prefetch event of this layer + event = torch.cuda.Event() + event.record(self.copy_stream) + self._prefetch_events[layer_idx] = event + # restore model's weights by their metadata using gpu buffer for name, meta in self._weight_metadata[layer_idx].items(): dtype = meta["dtype"] @@ -182,8 +199,16 @@ class LayerwiseOffloadManager: @torch.compiler.disable def release_layer(self, layer_idx: int) -> None: + """ + lightweight release layer weights + Basically set the reference count to the gpu weight tensor to zero. The weights on cpu is untouched + """ if not self.enabled or self.device is None: return + + # clear prefetch event, since it's useless and needs to be reset + self._prefetch_events.pop(layer_idx, None) + if layer_idx <= 0: return @@ -259,14 +284,24 @@ class LayerwiseOffloadManager: def make_pre_hook(i): def hook(module, input): - self.prefetch_layer(i + 1, non_blocking=True) + # wait only for the current layer if it's being prefetched + if i == 0: + self.prepare_for_next_req(non_blocking=False) + if i in self._prefetch_events: + torch.cuda.current_stream().wait_event(self._prefetch_events[i]) + + # trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed + if i % self.prefetch_size == 0: + for j in range(i + self.prefetch_size, i + 2 * self.prefetch_size): + layer_to_prefetch = j % self.num_layers + self.prefetch_layer(layer_to_prefetch, non_blocking=True) return hook def make_post_hook(i): def hook(module, input, output): - if self.copy_stream is not None: - torch.cuda.current_stream().wait_stream(self.copy_stream) + # previous, we wait here, until the copy stream for next layer is finished, + # now with any prefetch_size, only wait for the copy stream, when the copy stream is for the next layer self.release_layer(i) return hook @@ -292,7 +327,7 @@ class OffloadableDiTMixin: # the list of names of a DiT's layers/blocks layer_names: List[str] - layerwise_offload_managers: list[LayerwiseOffloadManager] | None = None + layerwise_offload_managers: list[LayerwiseOffloadManager] = [] def configure_layerwise_offload(self, server_args: ServerArgs): self.layerwise_offload_managers = [] @@ -303,12 +338,20 @@ class OffloadableDiTMixin: continue num_layers = len(module_list) + if server_args.dit_offload_prefetch_size < 1.0: + prefetch_size = 1 + int( + round(server_args.dit_offload_prefetch_size * (num_layers - 1)) + ) + else: + prefetch_size = int(server_args.dit_offload_prefetch_size) + manager = LayerwiseOffloadManager( model=self, layers_attr_str=layer_name, num_layers=num_layers, enabled=True, pin_cpu_memory=server_args.pin_cpu_memory, + prefetch_size=prefetch_size, ) self.layerwise_offload_managers.append(manager) @@ -316,11 +359,11 @@ class OffloadableDiTMixin: f"Enabled layerwise offload for {self.__class__.__name__} on modules: {self.layer_names}" ) - def prepare_for_next_denoise(self): + def prepare_for_next_req(self): if self.layerwise_offload_managers is None: return for manager in self.layerwise_offload_managers: - manager.prepare_for_next_denoise(non_blocking=True) + manager.prepare_for_next_req(non_blocking=True) def disable_offload(self) -> None: """Disable layerwise offload: load all layers to GPU and remove hooks.""" @@ -338,7 +381,5 @@ class OffloadableDiTMixin: 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.release_all() manager.register_forward_hooks() diff --git a/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py b/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py index 6d5406559..12473faab 100644 --- a/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py +++ b/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py @@ -53,19 +53,22 @@ def _openai_client(port: int) -> OpenAI: def _build_server_extra_args(case: DiffusionTestCase) -> str: + server_args = case.server_args a = os.environ.get("SGLANG_TEST_SERVE_ARGS", "") - a += f" --num-gpus {case.server_args.num_gpus}" - if case.server_args.tp_size is not None: - a += f" --tp-size {case.server_args.tp_size}" - if case.server_args.ulysses_degree is not None: - a += f" --ulysses-degree {case.server_args.ulysses_degree}" - if case.server_args.dit_layerwise_offload: + a += f" --num-gpus {server_args.num_gpus}" + if server_args.tp_size is not None: + a += f" --tp-size {server_args.tp_size}" + if server_args.ulysses_degree is not None: + a += f" --ulysses-degree {server_args.ulysses_degree}" + if server_args.dit_layerwise_offload: a += " --dit-layerwise-offload true" - if case.server_args.ring_degree is not None: - a += f" --ring-degree {case.server_args.ring_degree}" - if case.server_args.lora_path: - a += f" --lora-path {case.server_args.lora_path}" - if case.server_args.enable_warmup: + if server_args.dit_offload_prefetch_size: + a += f" --dit-offload-prefetch-size {server_args.dit_offload_prefetch_size}" + if server_args.ring_degree is not None: + a += f" --ring-degree {server_args.ring_degree}" + if server_args.lora_path: + a += f" --lora-path {server_args.lora_path}" + if server_args.enable_warmup: a += " --enable-warmup" return a 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 1614aeec8..a29f0606c 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -76,6 +76,11 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: if server_args.dit_layerwise_offload: extra_args += f" --dit-layerwise-offload true" + if server_args.dit_offload_prefetch_size: + extra_args += ( + f" --dit-offload-prefetch-size {server_args.dit_offload_prefetch_size}" + ) + if server_args.text_encoder_cpu_offload: extra_args += f" --text-encoder-cpu-offload" diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index e93ffa7f6..91743050e 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -164,6 +164,7 @@ class DiffusionServerArgs: enable_warmup: bool = False dit_layerwise_offload: bool = False + dit_offload_prefetch_size: int | float | None = None enable_cache_dit: bool = False text_encoder_cpu_offload: bool = False @@ -369,6 +370,7 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [ model_path="black-forest-labs/FLUX.2-dev", modality="image", dit_layerwise_offload=True, + dit_offload_prefetch_size=2, ), T2I_sampling_params, ),