[codex] Optimize DMD Wan auto residency on high-memory GPUs (#28780)
This commit is contained in:
@@ -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):
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user