[diffusion] Skip automatic Wan/MOVA DiT layerwise offload on high-end GPUs (#21248)
This commit is contained in:
@@ -36,6 +36,10 @@ from sglang.multimodal_gen.runtime.utils.common import (
|
||||
is_valid_ipv6_address,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import (
|
||||
CYAN,
|
||||
GREEN,
|
||||
RED,
|
||||
RESET,
|
||||
_sanitize_for_logging,
|
||||
configure_logger,
|
||||
init_logger,
|
||||
@@ -49,6 +53,18 @@ from sglang.multimodal_gen.utils import (
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Derived from single-H200 benchmarking (~140.4 GiB total) at the maximum
|
||||
# supported 720p workloads with dit_layerwise_offload=False and
|
||||
# num_inference_steps=1:
|
||||
# - Wan-AI/Wan2.2-T2V-A14B-Diffusers, 1280x720, 81 frames:
|
||||
# peak_reserved=108076 MB (~105.5 GiB), peak_allocated=97665 MB (~95.4 GiB)
|
||||
# - OpenMOSS-Team/MOVA-720p, 1280x720, 193 frames:
|
||||
# peak_reserved=130264 MB (~127.2 GiB), peak_allocated=108819 MB (~106.3 GiB)
|
||||
# Also, on H200, enabling dit_layerwise_offload regressed latency noticeably on
|
||||
# our validated Wan/MOVA workloads, so use a 130 GiB cutoff to keep H200-class
|
||||
# GPUs on the faster no-offload default while preserving some headroom.
|
||||
WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB = 130
|
||||
|
||||
|
||||
class Backend(str, Enum):
|
||||
"""
|
||||
@@ -434,15 +450,33 @@ class ServerArgs:
|
||||
if not envs.SGLANG_CACHE_DIT_ENABLED:
|
||||
pipeline_name_lower = self.pipeline_config.__class__.__name__.lower()
|
||||
if (
|
||||
("wan" in pipeline_name_lower or "mova" in pipeline_name_lower)
|
||||
and self.dit_layerwise_offload is None
|
||||
and current_platform.enable_dit_layerwise_offload_for_wan_by_default()
|
||||
):
|
||||
logger.info(
|
||||
f"Automatically enable dit_layerwise_offload for {self.pipeline_config.__class__.__name__} "
|
||||
"for low memory and performance balance"
|
||||
"wan" in pipeline_name_lower or "mova" in pipeline_name_lower
|
||||
) and self.dit_layerwise_offload is None:
|
||||
auto_enable_layerwise_offload = (
|
||||
current_platform.enable_dit_layerwise_offload_for_wan_by_default()
|
||||
)
|
||||
self.dit_layerwise_offload = True
|
||||
if auto_enable_layerwise_offload and current_platform.is_cuda():
|
||||
device_total_memory_gb = (
|
||||
current_platform.get_device_total_memory() / BYTES_PER_GB
|
||||
)
|
||||
if (
|
||||
device_total_memory_gb
|
||||
>= WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB
|
||||
):
|
||||
logger.info(
|
||||
"Skipping automatic dit_layerwise_offload for %s on a high-memory CUDA GPU (e.g. H200/B200/B300-class, %.2f GiB total)",
|
||||
self.pipeline_config.__class__.__name__,
|
||||
device_total_memory_gb,
|
||||
)
|
||||
auto_enable_layerwise_offload = False
|
||||
self.dit_layerwise_offload = False
|
||||
|
||||
if auto_enable_layerwise_offload:
|
||||
logger.info(
|
||||
f"Automatically enable dit_layerwise_offload for {self.pipeline_config.__class__.__name__} "
|
||||
"for low memory and performance balance"
|
||||
)
|
||||
self.dit_layerwise_offload = True
|
||||
|
||||
def _adjust_autocast(self):
|
||||
if self.disable_autocast is None:
|
||||
@@ -1057,6 +1091,18 @@ class ServerArgs:
|
||||
"Please disable either --dit-layerwise-offload or SGLANG_CACHE_DIT_ENABLED."
|
||||
)
|
||||
|
||||
logger.warning(
|
||||
"dit_layerwise_offload is enabled: %slower GPU memory usage%s, but %smay reduce throughput or increase latency%s. "
|
||||
"%sIf you are using multi-GPU deployment and already have enough memory headroom, prefer keeping dit_layerwise_offload disabled.%s "
|
||||
"Please tune this based on your memory headroom and performance target.",
|
||||
GREEN,
|
||||
RESET,
|
||||
RED,
|
||||
RESET,
|
||||
CYAN,
|
||||
RESET,
|
||||
)
|
||||
|
||||
def _validate_parallelism(self):
|
||||
if self.sp_degree > self.num_gpus or self.num_gpus % self.sp_degree != 0:
|
||||
raise ValueError(
|
||||
|
||||
Reference in New Issue
Block a user