From 8038806557f349898384cf6a8350f23c17d7dfb7 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 12 Jun 2026 13:13:18 +0800 Subject: [PATCH] [diffusion] optimize: optimize flux1 tensor parallel sharding (#27826) --- .../runtime/layers/attention/layer.py | 16 +- .../runtime/models/dits/flux.py | 210 +++++++++++++++--- .../runtime/utils/model_overlay.py | 64 +++++- .../sglang/multimodal_gen/test/test_utils.py | 77 +++++-- .../test/unit/test_transformer_quant.py | 5 + 5 files changed, 313 insertions(+), 59 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 8fb2c6b85..427400550 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -359,20 +359,28 @@ class LocalAttention(nn.Module): mask = mask[:, None, :, :] mask = (mask - 1.0) * torch.finfo(q_.dtype).max + if q_.shape[1] != k_.shape[1]: + repeat_factor = q_.shape[1] // k_.shape[1] + k_ = k_.repeat_interleave(repeat_factor, dim=1) + v_ = v_.repeat_interleave(repeat_factor, dim=1) + sdpa_context = ( sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS) if self.allow_cudnn_sdp and q_.device.type == "cuda" else nullcontext() ) + attn_kwargs = { + "attn_mask": mask, + "dropout_p": 0.0, + "is_causal": False, + "scale": self.softmax_scale, + } with sdpa_context: return torch.nn.functional.scaled_dot_product_attention( q_, k_, v_, - attn_mask=mask, - dropout_p=0.0, - is_causal=False, - scale=self.softmax_scale, + **attn_kwargs, ).transpose(1, 2) output = self.attn_impl.forward(q, k, v, attn_metadata=ctx_attn_metadata) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 24cc25e06..2196c14da 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -28,6 +28,7 @@ from diffusers.models.normalization import ( from torch.nn import LayerNorm as LayerNorm from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig +from sglang.multimodal_gen.runtime.distributed import divide, get_tp_world_size from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import ( RMSNorm, @@ -36,6 +37,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import ( from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, + RowParallelLinear, ) from sglang.multimodal_gen.runtime.layers.mlp import FeedForward from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( @@ -209,6 +211,75 @@ def _get_qkv_projections( return query, key, value, encoder_query, encoder_key, encoder_value +class FluxGELU(nn.Module): + def __init__( + self, + dim: int, + inner_dim: int, + bias: bool = True, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.proj = ColumnParallelLinear( + dim, + inner_dim, + bias=bias, + gather_output=False, + quant_config=quant_config, + prefix=f"{prefix}.proj" if prefix else "proj", + ) + self.gelu = nn.GELU(approximate="tanh") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.proj(hidden_states) + return self.gelu(hidden_states) + + +class FluxParallelFeedForward(nn.Module): + def __init__( + self, + dim: int, + dim_out: Optional[int] = None, + mult: int = 4, + inner_dim: Optional[int] = None, + bias: bool = True, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + if inner_dim is None: + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + self.net = nn.ModuleList( + [ + FluxGELU( + dim, + inner_dim, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.net.0" if prefix else "net.0", + ), + nn.Dropout(0.0), + RowParallelLinear( + inner_dim, + dim_out, + bias=bias, + input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.net.2" if prefix else "net.2", + ), + ] + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.net[0](hidden_states) + hidden_states = self.net[1](hidden_states) + hidden_states, _ = self.net[2](hidden_states) + return hidden_states + + class FluxAttention(torch.nn.Module, AttentionModuleMixin): def __init__( self, @@ -238,6 +309,11 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): self.context_pre_only = context_pre_only self.pre_only = pre_only self.heads = out_dim // dim_head if out_dim is not None else num_heads + self.tp_size = get_tp_world_size() + self.shard_qkv = self.tp_size > 1 and not isinstance( + quant_config, NunchakuConfig + ) + self.local_heads = divide(self.heads, self.tp_size) self.added_kv_proj_dim = added_kv_proj_dim self.added_proj_bias = added_proj_bias @@ -252,7 +328,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): query_dim, [self.inner_dim] * 3, bias=bias, - gather_output=True, + gather_output=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.to_qkv" if prefix else "to_qkv", ) @@ -261,7 +337,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): query_dim, self.inner_dim, bias=bias, - gather_output=True, + gather_output=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.to_q" if prefix else "to_q", ) @@ -269,7 +345,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): query_dim, self.inner_dim, bias=bias, - gather_output=True, + gather_output=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.to_k" if prefix else "to_k", ) @@ -277,18 +353,24 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): query_dim, self.inner_dim, bias=bias, - gather_output=True, + gather_output=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.to_v" if prefix else "to_v", ) if not self.pre_only: self.to_out = torch.nn.ModuleList([]) + out_proj_cls = RowParallelLinear if self.shard_qkv else ColumnParallelLinear + out_proj_kwargs = ( + {"input_is_parallel": True} + if self.shard_qkv + else {"gather_output": True} + ) self.to_out.append( - ColumnParallelLinear( + out_proj_cls( self.inner_dim, self.out_dim, bias=out_bias, - gather_output=True, + **out_proj_kwargs, quant_config=quant_config, prefix=f"{prefix}.to_out.0" if prefix else "", ) @@ -304,7 +386,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): added_kv_proj_dim, [self.inner_dim] * 3, bias=added_proj_bias, - gather_output=True, + gather_output=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.to_added_qkv" if prefix else "to_added_qkv", ) @@ -313,7 +395,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, - gather_output=True, + gather_output=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.add_q_proj" if prefix else "add_q_proj", ) @@ -321,7 +403,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, - gather_output=True, + gather_output=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.add_k_proj" if prefix else "add_k_proj", ) @@ -329,21 +411,29 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, - gather_output=True, + gather_output=not self.shard_qkv, quant_config=quant_config, prefix=f"{prefix}.add_v_proj" if prefix else "add_v_proj", ) - self.to_add_out = ColumnParallelLinear( + add_out_proj_cls = ( + RowParallelLinear if self.shard_qkv else ColumnParallelLinear + ) + add_out_proj_kwargs = ( + {"input_is_parallel": True} + if self.shard_qkv + else {"gather_output": True} + ) + self.to_add_out = add_out_proj_cls( self.inner_dim, query_dim, bias=out_bias, - gather_output=True, + **add_out_proj_kwargs, quant_config=quant_config, prefix=f"{prefix}.to_add_out" if prefix else "", ) self.attn = USPAttention( - num_heads=num_heads, + num_heads=self.local_heads if self.shard_qkv else num_heads, head_size=self.head_dim, dropout_rate=0, softmax_scale=None, @@ -366,9 +456,10 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): encoder_value, ) = _get_qkv_projections(self, x, encoder_hidden_states) - query = query.unflatten(-1, (self.heads, -1)) - key = key.unflatten(-1, (self.heads, -1)) - value = value.unflatten(-1, (self.heads, -1)) + num_heads = self.local_heads if self.shard_qkv else self.heads + query = query.unflatten(-1, (num_heads, -1)) + key = key.unflatten(-1, (num_heads, -1)) + value = value.unflatten(-1, (num_heads, -1)) cos_sin_cache = None if freqs_cis is not None: cos, sin = freqs_cis @@ -381,9 +472,9 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): ) if self.added_kv_proj_dim is not None: - encoder_query = encoder_query.unflatten(-1, (self.heads, -1)) - encoder_key = encoder_key.unflatten(-1, (self.heads, -1)) - encoder_value = encoder_value.unflatten(-1, (self.heads, -1)) + encoder_query = encoder_query.unflatten(-1, (num_heads, -1)) + encoder_key = encoder_key.unflatten(-1, (num_heads, -1)) + encoder_value = encoder_value.unflatten(-1, (num_heads, -1)) text_seq_len = encoder_query.shape[1] encoder_query, encoder_key = apply_qk_norm_with_optional_rope( @@ -466,6 +557,9 @@ class FluxSingleTransformerBlock(nn.Module): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) self.use_nunchaku_structure = isinstance(quant_config, NunchakuConfig) + self.tp_size = get_tp_world_size() + self.local_mlp_hidden_dim = divide(self.mlp_hidden_dim, self.tp_size) + self.local_dim = divide(dim, self.tp_size) self.norm = AdaLayerNormZeroSingle(dim) @@ -502,23 +596,34 @@ class FluxSingleTransformerBlock(nn.Module): if is_nunchaku_available(): self.norm = NunchakuAdaLayerNormZeroSingle(self.norm, scale_shift=0) else: + shard_single_block = self.tp_size > 1 self.proj_mlp = ColumnParallelLinear( dim, self.mlp_hidden_dim, bias=True, - gather_output=True, + gather_output=not shard_single_block, quant_config=quant_config, prefix=f"{prefix}.proj_mlp" if prefix else "proj_mlp", ) self.act_mlp = nn.GELU(approximate="tanh") - self.proj_out = ColumnParallelLinear( + proj_out_cls = ( + RowParallelLinear if shard_single_block else ColumnParallelLinear + ) + proj_out_kwargs = ( + {"input_is_parallel": True} + if shard_single_block + else {"gather_output": True} + ) + self.proj_out = proj_out_cls( dim + self.mlp_hidden_dim, dim, bias=True, - gather_output=True, + **proj_out_kwargs, quant_config=quant_config, prefix=f"{prefix}.proj_out" if prefix else "proj_out", ) + if shard_single_block: + self._patch_proj_out_weight_loader() self.attn = FluxAttention( query_dim=dim, dim_head=attention_head_dim, @@ -531,6 +636,30 @@ class FluxSingleTransformerBlock(nn.Module): prefix=f"{prefix}.attn" if prefix else "attn", ) + def _patch_proj_out_weight_loader(self) -> None: + dim, mlp_dim = self.local_dim, self.local_mlp_hidden_dim + tp_rank = self.proj_out.tp_rank + + def _loader(param, loaded_weight): + input_dim = getattr(param, "input_dim", None) + if input_dim is not None: + # checkpoint columns are [attn_full | mlp_full], while TP consumes [attn_shard | mlp_shard] + attn_cols = loaded_weight.narrow(input_dim, tp_rank * dim, dim) + mlp_cols = loaded_weight.narrow( + input_dim, + self.tp_size * dim + tp_rank * mlp_dim, + mlp_dim, + ) + param.data.copy_(torch.cat([attn_cols, mlp_cols], dim=input_dim)) + else: + param.data.copy_(loaded_weight) + + self.proj_out.weight_loader = _loader + if hasattr(self.proj_out.weight, "_weight_loader"): + self.proj_out.weight._weight_loader = _loader + else: + self.proj_out.weight.weight_loader = _loader + def forward( self, hidden_states: torch.Tensor, @@ -634,13 +763,16 @@ class FluxTransformerBlock(nn.Module): and is_nunchaku_available() ) self.use_nunchaku_structure = nunchaku_enabled - self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") - self.ff_context = FeedForward( - dim=dim, - dim_out=dim, - activation_fn="gelu-approximate", - ) + self.tp_size = get_tp_world_size() if nunchaku_enabled: + self.ff = FeedForward( + dim=dim, dim_out=dim, activation_fn="gelu-approximate" + ) + self.ff_context = FeedForward( + dim=dim, + dim_out=dim, + activation_fn="gelu-approximate", + ) nunchaku_kwargs = { "precision": quant_config.precision, "rank": quant_config.rank, @@ -652,6 +784,28 @@ class FluxTransformerBlock(nn.Module): self.norm1_context = NunchakuAdaLayerNormZero( self.norm1_context, scale_shift=0 ) + elif self.tp_size > 1: + self.ff = FluxParallelFeedForward( + dim=dim, + dim_out=dim, + quant_config=quant_config, + prefix=f"{prefix}.ff" if prefix else "ff", + ) + self.ff_context = FluxParallelFeedForward( + dim=dim, + dim_out=dim, + quant_config=quant_config, + prefix=f"{prefix}.ff_context" if prefix else "ff_context", + ) + else: + self.ff = FeedForward( + dim=dim, dim_out=dim, activation_fn="gelu-approximate" + ) + self.ff_context = FeedForward( + dim=dim, + dim_out=dim, + activation_fn="gelu-approximate", + ) def forward( self, diff --git a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py index fc36551fc..566ea3d88 100644 --- a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py +++ b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py @@ -51,6 +51,16 @@ MODEL_OVERLAY_METADATA_PATTERNS = [ "**/*.txt", ] +_MATERIALIZED_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth", ".pt") +_MATERIALIZED_CONFIG_ONLY_COMPONENTS = { + "feature_extractor", + "image_processor", + "processor", + "scheduler", + "tokenizer", + "tokenizer_2", +} + _MODEL_OVERLAY_REGISTRY_CACHE: dict[str, dict[str, Any]] | None = None @@ -196,6 +206,52 @@ def load_model_index_from_dir(model_dir: str) -> dict[str, Any]: return config +def _component_has_weight_file(component_dir: str) -> bool: + for root, _, file_names in os.walk(component_dir): + if any( + file_name.endswith(_MATERIALIZED_WEIGHT_SUFFIXES) + and os.path.isfile(os.path.join(root, file_name)) + for file_name in file_names + ): + return True + return False + + +def _materialized_overlay_has_component_weights(model_dir: str) -> bool: + model_index = load_model_index_from_dir(model_dir) + for component_name, entry in model_index.items(): + if ( + component_name.startswith("_") + or component_name == "pipeline_name" + or component_name in _MATERIALIZED_CONFIG_ONLY_COMPONENTS + or not isinstance(entry, list) + ): + continue + component_dir = os.path.join(model_dir, component_name) + if not os.path.isdir(component_dir) or not _component_has_weight_file( + component_dir + ): + logger.warning( + "Materialized overlay cache for %s is missing weights for component %s", + model_dir, + component_name, + ) + return False + return True + + +def _materialized_overlay_cache_complete( + final_dir: str, + marker_path: str, + verify_diffusers_model_complete_fn: Callable[[str], bool], +) -> bool: + return ( + verify_diffusers_model_complete_fn(final_dir) + and os.path.exists(marker_path) + and _materialized_overlay_has_component_weights(final_dir) + ) + + def _ensure_dir(path: str) -> None: os.makedirs(path, exist_ok=True) @@ -503,15 +559,17 @@ def materialize_overlay_model( safe_name = source_model_id.replace("/", "__") final_dir = os.path.join(cache_root, f"{safe_name}-{cache_key}") marker_path = os.path.join(final_dir, ".sglang_overlay_materialized.json") - if verify_diffusers_model_complete_fn(final_dir) and os.path.exists(marker_path): + if _materialized_overlay_cache_complete( + final_dir, marker_path, verify_diffusers_model_complete_fn + ): return final_dir lock_name = ( f"overlay-materialize::{source_model_id}::{overlay_repo_id}::{overlay_revision}" ) with get_lock(lock_name).acquire(poll_interval=2): - if verify_diffusers_model_complete_fn(final_dir) and os.path.exists( - marker_path + if _materialized_overlay_cache_complete( + final_dir, marker_path, verify_diffusers_model_complete_fn ): return final_dir diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 41343218b..9bec42542 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -995,29 +995,56 @@ def _remote_consistency_gt_candidates( return [(filename, f"{base_url}/{filename}") for filename in filenames] +def _is_ascend_consistency_case(case_id: str) -> bool: + return "npu" in case_id + + def _remote_file_exists(url: str) -> bool: - for method in ("head", "get"): - try: - if method == "head": - resp = requests.head(url, timeout=10, allow_redirects=True) - else: - resp = requests.get( - url, - timeout=10, - allow_redirects=True, - headers={"Range": "bytes=0-0"}, - stream=True, - ) + for _ in range(3): + for method in ("head", "get"): try: - if resp.status_code in (200, 206): - return True - if resp.status_code not in (403, 405, 429) and resp.status_code < 500: - return False + if method == "head": + resp = requests.head(url, timeout=30, allow_redirects=True) + else: + resp = requests.get( + url, + timeout=30, + allow_redirects=True, + headers={"Range": "bytes=0-0"}, + stream=True, + ) + try: + if resp.status_code in (200, 206): + return True + if ( + resp.status_code not in (403, 405, 429) + and resp.status_code < 500 + ): + return False + finally: + resp.close() + except requests.RequestException: + pass + return False + + +def _load_remote_gt_image(url: str) -> np.ndarray: + last_error: Exception | None = None + for _ in range(3): + try: + resp = requests.get(url, timeout=60) + try: + if resp.status_code == 200: + image = Image.open(io.BytesIO(resp.content)).convert("RGB") + return np.array(image) + last_error = FileNotFoundError(f"GT image not found: {url}") + if resp.status_code not in (403, 429) and resp.status_code < 500: + break finally: resp.close() - except requests.RequestException: - pass - return False + except requests.RequestException as exc: + last_error = exc + raise FileNotFoundError(f"GT image not found: {url}") from last_error def _find_remote_consistency_gt_files( @@ -1026,7 +1053,12 @@ def _find_remote_consistency_gt_files( is_video: bool, output_format: str | None = None, ) -> list[tuple[str, str]]: - if case_id in SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES: + if _is_ascend_consistency_case(case_id): + bases = ( + SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE_ASCEND, + SGL_TEST_FILES_CONSISTENCY_GT_BASE, + ) + elif case_id in SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_CASES: bases = SGL_TEST_FILES_CONSISTENCY_GT_BASES else: # Avoid accidentally comparing non-comparable CI cases against official GT. @@ -1115,10 +1147,7 @@ def load_consistency_gt( f"GT image not found for {case_id}. Tried: {', '.join(filenames)}" ) for _, url in remote_files: - resp = requests.get(url, timeout=30) - if resp.status_code != 200: - raise FileNotFoundError(f"GT image not found: {url}") - images.append(np.array(Image.open(io.BytesIO(resp.content)).convert("RGB"))) + images.append(_load_remote_gt_image(url)) source_dir = remote_files[0][1].rsplit("/", 1)[0] logger.info(f"Loaded {len(images)} GT images for {case_id} from {source_dir}") diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 6e00922e0..a34563c42 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -390,8 +390,13 @@ class TestTransformerQuantHelpers(unittest.TestCase): "sglang.multimodal_gen.runtime.layers.attention.selector.get_global_server_args", return_value=SimpleNamespace(attention_backend=None), ) + @patch( + "sglang.multimodal_gen.runtime.models.dits.flux.get_tp_world_size", + return_value=1, + ) def test_flux_single_transformer_block_modelopt_excludes_use_full_prefix( self, + _mock_tp_world_size, _mock_server_args, _mock_ring_world_size, _mock_tp_group,