[diffusion] feat: support sp for image models (#13180)

This commit is contained in:
Mick
2025-11-23 18:11:42 +08:00
committed by GitHub
parent 53fffefd5d
commit d4593964fe
19 changed files with 624 additions and 393 deletions

View File

@@ -793,18 +793,6 @@ class ServerArgs:
return provided_args
def check_server_sp_args(self):
if self.pipeline_config.task_type.is_image_gen():
if (
(self.sp_degree and self.sp_degree > 1)
or (self.ulysses_degree and self.ulysses_degree > 1)
or (self.ring_degree and self.ring_degree > 1)
):
raise ValueError(
"SP is not supported for image generation models for now"
)
self.sp_degree = self.ulysses_degree = self.ring_degree = 1
if self.sp_degree == -1:
# assume we leave all remaining gpus to sp
num_gpus_per_group = self.dp_size * self.tp_size
@@ -861,8 +849,11 @@ class ServerArgs:
def check_server_dp_args(self):
assert self.num_gpus % self.dp_size == 0, f"{self.num_gpus=}, {self.dp_size=}"
assert self.dp_size >= 1, "--dp-size must be natural number"
self.dp_degree = self.num_gpus // self.dp_size
# NOTE: disable temporarily
# self.dp_degree = self.num_gpus // self.dp_size
logger.info(f"Setting dp_degree to: {self.dp_degree}")
if self.dp_size > 1:
raise ValueError("DP is not yet supported")
def check_server_args(self) -> None:
"""Validate inference arguments for consistency"""
@@ -920,18 +911,6 @@ class ServerArgs:
self.pipeline_config.check_pipeline_config()
# Add preprocessing config validation if needed
if self.mode == ExecutionMode.PREPROCESS:
if self.preprocess_config is None:
raise ValueError(
"preprocess_config is not set in ServerArgs when mode is PREPROCESS"
)
if self.preprocess_config.model_path == "":
self.preprocess_config.model_path = self.model_path
if not self.pipeline_config.vae_config.load_encoder:
self.pipeline_config.vae_config.load_encoder = True
self.preprocess_config.check_preprocess_config()
# parallelism
self.check_server_dp_args()
# allocate all remaining gpus for sp-size