[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,
|
TextEncoderArchConfig,
|
||||||
TextEncoderConfig,
|
TextEncoderConfig,
|
||||||
)
|
)
|
||||||
from sglang.multimodal_gen.utils import StoreBoolean
|
|
||||||
|
|
||||||
|
|
||||||
def _is_transformer_layer(n: str, m) -> bool:
|
def _is_transformer_layer(n: str, m) -> bool:
|
||||||
@@ -95,19 +94,4 @@ class T5Config(TextEncoderConfig):
|
|||||||
def add_cli_args(
|
def add_cli_args(
|
||||||
parser: argparse.ArgumentParser, prefix: str = "t5-config"
|
parser: argparse.ArgumentParser, prefix: str = "t5-config"
|
||||||
) -> argparse.ArgumentParser:
|
) -> 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
|
return parser
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ import addict
|
|||||||
import yaml
|
import yaml
|
||||||
|
|
||||||
from sglang.multimodal_gen import envs
|
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.pipeline_configs.base import PipelineConfig
|
||||||
from sglang.multimodal_gen.configs.quantization import NunchakuSVDQuantArgs
|
from sglang.multimodal_gen.configs.quantization import NunchakuSVDQuantArgs
|
||||||
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
|
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
|
||||||
@@ -353,6 +354,7 @@ class ServerArgs:
|
|||||||
self._adjust_attention_backend()
|
self._adjust_attention_backend()
|
||||||
self._adjust_platform_specific()
|
self._adjust_platform_specific()
|
||||||
self._adjust_autocast()
|
self._adjust_autocast()
|
||||||
|
self.adjust_pipeline_config()
|
||||||
|
|
||||||
def _validate_parameters(self):
|
def _validate_parameters(self):
|
||||||
"""check consistency and raise errors for invalid configs"""
|
"""check consistency and raise errors for invalid configs"""
|
||||||
@@ -391,6 +393,26 @@ class ServerArgs:
|
|||||||
transformer_weights_path=self.nunchaku_config.transformer_weights_path,
|
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):
|
def _adjust_offload(self):
|
||||||
# TODO: to be handled by each platform
|
# TODO: to be handled by each platform
|
||||||
if current_platform.get_device_total_memory() / BYTES_PER_GB < 30:
|
if current_platform.get_device_total_memory() / BYTES_PER_GB < 30:
|
||||||
|
|||||||
Reference in New Issue
Block a user