[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 @dataclass
class WanT2V720PConfig(WanT2V480PConfig): class WanT2V720PConfig(WanT2V480PConfig):
"""Base configuration for Wan T2V 14B 720P pipeline architecture.""" """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] 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 @dataclass
class Wan2_2_TI2V_5B_Config(WanT2V480PConfig, WanI2VCommonConfig): 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_1_T2V_480P_Config,
FastWan2_2_TI2V_5B_Config, FastWan2_2_TI2V_5B_Config,
TurboWanI2V720Config, TurboWanI2V720Config,
TurboWanT2V1_3B480PConfig,
TurboWanT2V480PConfig, TurboWanT2V480PConfig,
Wan2_2_I2V_A14B_Config, Wan2_2_I2V_A14B_Config,
Wan2_2_T2V_A14B_Config, Wan2_2_T2V_A14B_Config,
@@ -674,7 +675,7 @@ def _register_configs():
) )
register_configs( register_configs(
sampling_param_cls=WanT2V_1_3B_SamplingParams, sampling_param_cls=WanT2V_1_3B_SamplingParams,
pipeline_config_cls=TurboWanT2V480PConfig, pipeline_config_cls=TurboWanT2V1_3B480PConfig,
hf_model_paths=[ hf_model_paths=[
"IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers", "IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
], ],
@@ -95,8 +95,9 @@ class ServerArgsAutoTuner:
return return
min_available_gb = self._get_min_available_device_memory_gb() min_available_gb = self._get_min_available_device_memory_gb()
deployment_config = self._deployment_config()
disable_threshold_gb = ( 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 ( if (
min_available_gb is not None min_available_gb is not None
@@ -104,9 +105,7 @@ class ServerArgsAutoTuner:
and min_available_gb >= disable_threshold_gb and min_available_gb >= disable_threshold_gb
): ):
changed = [] changed = []
components = ( components = deployment_config.auto_disable_component_offload_components
self._deployment_config().auto_disable_component_offload_components
)
if ( if (
args.layerwise_offload_components is not None args.layerwise_offload_components is not None
and not args.is_arg_explicitly_set("layerwise_offload_components") 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 for component_name, arg_name in DEFAULT_LAYERWISE_COMPONENT_ARG_NAMES
if not args.is_arg_explicitly_set(arg_name) 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(): if self._should_auto_enable_dit_layerwise_offload():
components.insert(0, LAYERWISE_OFFLOAD_DIT_GROUP) components.insert(0, LAYERWISE_OFFLOAD_DIT_GROUP)
self._set_default_wan_dit_offload_prefetch_size() self._set_default_wan_dit_offload_prefetch_size()
return components 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: def _should_auto_enable_dit_layerwise_offload(self) -> bool:
args = self.server_args args = self.server_args