[Diffusion] Open t5 encoder parallel folding for wan2.2 and mova video (#18493)

Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Xiaoyu Zhang
2026-03-05 10:18:00 +08:00
committed by GitHub
parent e555a6c171
commit 9795b4cd5b
2 changed files with 22 additions and 16 deletions

View File

@@ -8,7 +8,6 @@ from sglang.multimodal_gen.configs.models.encoders.base import (
TextEncoderArchConfig,
TextEncoderConfig,
)
from sglang.multimodal_gen.utils import StoreBoolean
def _is_transformer_layer(n: str, m) -> bool:
@@ -95,19 +94,4 @@ class T5Config(TextEncoderConfig):
def add_cli_args(
parser: argparse.ArgumentParser, prefix: str = "t5-config"
) -> argparse.ArgumentParser:
parser.add_argument(
f"--{prefix}.parallel-folding",
action=StoreBoolean,
dest=f"{prefix.replace('-', '_')}.parallel_folding",
default=T5Config.parallel_folding,
help="Whether to use parallel folding for T5",
)
parser.add_argument(
f"--{prefix}.parallel-folding-mode",
type=str,
choices=["sp", "ulysses", "ring"],
dest=f"{prefix.replace('-', '_')}.parallel_folding_mode",
default=T5Config.parallel_folding_mode,
help="Parallel folding mode for T5",
)
return parser

View File

@@ -21,6 +21,7 @@ import addict
import yaml
from sglang.multimodal_gen import envs
from sglang.multimodal_gen.configs.models.encoders import T5Config
from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
from sglang.multimodal_gen.configs.quantization import NunchakuSVDQuantArgs
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
@@ -353,6 +354,7 @@ class ServerArgs:
self._adjust_attention_backend()
self._adjust_platform_specific()
self._adjust_autocast()
self.adjust_pipeline_config()
def _validate_parameters(self):
"""check consistency and raise errors for invalid configs"""
@@ -391,6 +393,26 @@ class ServerArgs:
transformer_weights_path=self.nunchaku_config.transformer_weights_path,
)
def adjust_pipeline_config(self):
# enable parallel folding when SP is enabled
if self.tp_size != 1 or self.sp_degree <= 1:
return
enabled = False
for text_encoder_config in self.pipeline_config.text_encoder_configs:
if isinstance(text_encoder_config, T5Config):
text_encoder_config.parallel_folding = True
enabled = True
text_encoder_config.parallel_folding_mode = "sp"
if enabled:
logger.info(
"Enabled T5 text encoder parallel folding (mode=sp) for %s (tp_size=%s, sp_degree=%s).",
self.__class__.__name__,
self.tp_size,
self.sp_degree,
)
def _adjust_offload(self):
# TODO: to be handled by each platform
if current_platform.get_device_total_memory() / BYTES_PER_GB < 30: