[codex] Optimize DMD Wan auto residency on high-memory GPUs (#28780)

This commit is contained in:
Xiaoyu Zhang
2026-06-24 14:03:23 +08:00
committed by GitHub
parent 6842335fcf
commit 534ac98eb2
3 changed files with 71 additions and 5 deletions
@@ -108,6 +108,22 @@ class TurboWanT2V480PConfig(WanT2V480PConfig):
)
@dataclass
class TurboWanT2V1_3B480PConfig(TurboWanT2V480PConfig):
"""Configuration for TurboWan T2V 1.3B DMD pipeline."""
def get_model_deployment_config(self) -> ModelDeploymentConfig:
return ModelDeploymentConfig(
auto_dit_layerwise_offload=True,
auto_disable_component_offload_min_available_memory_gb=60,
auto_disable_component_offload_components=(
"text_encoder",
"image_encoder",
"vae",
),
)
@dataclass
class WanT2V720PConfig(WanT2V480PConfig):
"""Base configuration for Wan T2V 14B 720P pipeline architecture."""
@@ -183,6 +199,17 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
default_factory=lambda: [1000, 757, 522]
)
def get_model_deployment_config(self) -> ModelDeploymentConfig:
return ModelDeploymentConfig(
auto_dit_layerwise_offload=True,
auto_disable_component_offload_min_available_memory_gb=60,
auto_disable_component_offload_components=(
"text_encoder",
"image_encoder",
"vae",
),
)
@dataclass
class Wan2_2_TI2V_5B_Config(WanT2V480PConfig, WanI2VCommonConfig):
+2 -1
View File
@@ -88,6 +88,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
FastWan2_1_T2V_480P_Config,
FastWan2_2_TI2V_5B_Config,
TurboWanI2V720Config,
TurboWanT2V1_3B480PConfig,
TurboWanT2V480PConfig,
Wan2_2_I2V_A14B_Config,
Wan2_2_T2V_A14B_Config,
@@ -674,7 +675,7 @@ def _register_configs():
)
register_configs(
sampling_param_cls=WanT2V_1_3B_SamplingParams,
pipeline_config_cls=TurboWanT2V480PConfig,
pipeline_config_cls=TurboWanT2V1_3B480PConfig,
hf_model_paths=[
"IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
],
@@ -95,8 +95,9 @@ class ServerArgsAutoTuner:
return
min_available_gb = self._get_min_available_device_memory_gb()
deployment_config = self._deployment_config()
disable_threshold_gb = (
self._deployment_config().auto_disable_component_offload_min_available_memory_gb
deployment_config.auto_disable_component_offload_min_available_memory_gb
)
if (
min_available_gb is not None
@@ -104,9 +105,7 @@ class ServerArgsAutoTuner:
and min_available_gb >= disable_threshold_gb
):
changed = []
components = (
self._deployment_config().auto_disable_component_offload_components
)
components = deployment_config.auto_disable_component_offload_components
if (
args.layerwise_offload_components is not None
and not args.is_arg_explicitly_set("layerwise_offload_components")
@@ -387,11 +386,50 @@ class ServerArgsAutoTuner:
for component_name, arg_name in DEFAULT_LAYERWISE_COMPONENT_ARG_NAMES
if not args.is_arg_explicitly_set(arg_name)
]
components = self._filter_high_memory_resident_components(components)
if self._should_auto_enable_dit_layerwise_offload():
components.insert(0, LAYERWISE_OFFLOAD_DIT_GROUP)
self._set_default_wan_dit_offload_prefetch_size()
return components
def _filter_high_memory_resident_components(
self, components: list[str]
) -> list[str]:
args = self.server_args
if args.performance_mode != "auto" or current_platform.is_cpu():
return components
deployment_config = self._deployment_config()
threshold_gb = (
deployment_config.auto_disable_component_offload_min_available_memory_gb
)
if threshold_gb is None:
return components
min_available_gb = self._get_min_available_device_memory_gb()
if min_available_gb is None or min_available_gb < threshold_gb:
return components
resident_components = set(
deployment_config.auto_disable_component_offload_components
)
filtered_components = [
component
for component in components
if component not in resident_components
]
skipped_components = [
component for component in components if component in resident_components
]
if skipped_components:
logger.info(
"Keeping default layerwise components resident for %s because minimum available memory on selected GPUs is %.2f GiB: %s",
args.pipeline_config.__class__.__name__,
min_available_gb,
", ".join(skipped_components),
)
return filtered_components
def _should_auto_enable_dit_layerwise_offload(self) -> bool:
args = self.server_args