chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] feat: add an arg for controlling the number of prefetched layers in layerwise-offload (#17693)

Browse Source
This commit is contained in:
Mick
2026-01-28 09:34:27 +08:00
committed by GitHub
parent 1507dc6cdf
commit 88fcd8535f
6 changed files with 106 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

@@ -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."""

26
python/sglang/multimodal_gen/runtime/server_args.py
View File

@@ -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."

69
python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py
View File

@@ -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()

25
python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py
View File

@@ -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

5
python/sglang/multimodal_gen/test/server/test_server_common.py
View File

@@ -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"

2
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -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,
),