[Diffusion] Support layerwise offload for mova (#18272)
This commit is contained in:
@@ -642,7 +642,7 @@ class ServerArgs:
|
|||||||
"--dit-layerwise-offload",
|
"--dit-layerwise-offload",
|
||||||
action=StoreBoolean,
|
action=StoreBoolean,
|
||||||
default=ServerArgs.dit_layerwise_offload,
|
default=ServerArgs.dit_layerwise_offload,
|
||||||
help="Enable layerwise CPU offload with async H2D prefetch overlap for supported DiT models (e.g., Wan). "
|
help="Enable layerwise CPU offload with async H2D prefetch overlap for supported DiT models (e.g., Wan, MOVA). "
|
||||||
"Cannot be used together with cache-dit (SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.",
|
"Cannot be used together with cache-dit (SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
@@ -998,13 +998,14 @@ class ServerArgs:
|
|||||||
|
|
||||||
if not envs.SGLANG_CACHE_DIT_ENABLED:
|
if not envs.SGLANG_CACHE_DIT_ENABLED:
|
||||||
# TODO: need a better way to tell this
|
# TODO: need a better way to tell this
|
||||||
|
pipeline_name_lower = self.pipeline_config.__class__.__name__.lower()
|
||||||
if (
|
if (
|
||||||
"wan" in self.pipeline_config.__class__.__name__.lower()
|
("wan" in pipeline_name_lower or "mova" in pipeline_name_lower)
|
||||||
and self.dit_layerwise_offload is None
|
and self.dit_layerwise_offload is None
|
||||||
and current_platform.enable_dit_layerwise_offload_for_wan_by_default()
|
and current_platform.enable_dit_layerwise_offload_for_wan_by_default()
|
||||||
):
|
):
|
||||||
logger.info(
|
logger.info(
|
||||||
"Automatically enable dit_layerwise_offload for Wan for best performance"
|
f"Automatically enable dit_layerwise_offload for {self.pipeline_config.__class__.__name__} for low memory and performance balance"
|
||||||
)
|
)
|
||||||
self.dit_layerwise_offload = True
|
self.dit_layerwise_offload = True
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user