[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:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user