From efcdda01766b5e5f3fd94781ad9860bb840d4102 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 10 Feb 2026 20:22:20 +0800 Subject: [PATCH] [diffusion] fix: fix fsdp (#18187) --- .../configs/models/dits/zimage.py | 13 ++++++++++ .../runtime/layers/layernorm.py | 5 ++++ .../runtime/layers/lora/linear.py | 2 +- .../runtime/layers/triton_ops.py | 3 +++ .../component_loaders/text_encoder_loader.py | 2 +- .../runtime/loader/fsdp_load.py | 19 ++++++++++++-- .../runtime/managers/gpu_worker.py | 5 ++-- .../runtime/models/dits/zimage.py | 1 + .../pipelines_core/stages/denoising.py | 4 +++ .../test/scripts/gen_perf_baselines.py | 3 +++ .../test/server/perf_baselines.json | 25 +++++++++++++++++++ .../test/server/test_server_common.py | 3 +++ .../test/server/testcase_configs.py | 25 +++++++++++++++++-- 13 files changed, 102 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/dits/zimage.py b/python/sglang/multimodal_gen/configs/models/dits/zimage.py index bef5e948e..33c50e0cb 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/configs/models/dits/zimage.py @@ -7,6 +7,17 @@ from typing import Tuple from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig +def is_zimage_layer(n: str, m) -> bool: + """Returns if the module should be sharded for Z-Image model.""" + if "layers" in n and str.isdigit(n.split(".")[-1]): + return True + if ("noise_refiner" in n or "context_refiner" in n) and str.isdigit( + n.split(".")[-1] + ): + return True + return False + + @dataclass class ZImageArchConfig(DiTArchConfig): all_patch_size: Tuple[int, ...] = (2,) @@ -26,6 +37,8 @@ class ZImageArchConfig(DiTArchConfig): axes_dims: Tuple[int, int, int] = (32, 48, 48) axes_lens: Tuple[int, int, int] = (1024, 512, 512) + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_layer]) + stacked_params_mapping: list[tuple[str, str, str]] = field( default_factory=lambda: [ # (param_name, shard_name, shard_id) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 0f76516b8..4d550787c 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -81,6 +81,10 @@ class RMSNorm(CustomOp): if x.dtype == torch.float: # fp32 out = self.forward_triton(x, residual) + if residual is not None: + return out[0].view(shape), out[1].view(residual_shape) + out = out.view(shape) + return out elif self.variance_size_override is not None: return self.forward_native(x, residual) elif residual is not None: @@ -94,6 +98,7 @@ class RMSNorm(CustomOp): else: out = rmsnorm(x, self.weight.data, self.variance_epsilon) out = out.view(shape) + return out def forward_native( diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index ebeabef94..a062479d1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -342,7 +342,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): super().__init__(base_layer, lora_rank, lora_alpha) def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor: - return A.to(self.base_layer.weight) + return A def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: tp_rank = get_tp_rank() diff --git a/python/sglang/multimodal_gen/runtime/layers/triton_ops.py b/python/sglang/multimodal_gen/runtime/layers/triton_ops.py index 13d613934..ddcd78229 100644 --- a/python/sglang/multimodal_gen/runtime/layers/triton_ops.py +++ b/python/sglang/multimodal_gen/runtime/layers/triton_ops.py @@ -948,6 +948,9 @@ class LayerNormFn: ) ) y = y.reshape(x_shape_og) + if residual is not None: + residual_out = residual_out.reshape(x_shape_og) + return y, residual_out return y diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 4e275d86b..3aeb6093c 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -279,7 +279,7 @@ class TextEncoderLoader(ComponentLoader): # if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights if weights_not_loaded: - raise ValueError( + logger.warning( "Following model weights were not initialized from " f"checkpoint: {weights_not_loaded}" ) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 6a7ff746a..06f80f281 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -231,10 +231,20 @@ def load_model_from_full_model_state_dict( custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( full_sd_iterator, param_names_mapping ) # type: ignore - for target_param_name, full_tensor in custom_param_sd.items(): + + is_fsdp_model = isinstance(model, FSDPModule) or any( + hasattr(p, "device_mesh") for p in meta_sd.values() + ) + + # sort parameter names to ensure all ranks process parameters in the same order + sorted_param_names = sorted(custom_param_sd.keys()) + + for target_param_name in sorted_param_names: + full_tensor = custom_param_sd[target_param_name] meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: - if strict: + # For FSDP models, ensure all ranks process parameters consistently + if strict or is_fsdp_model: raise ValueError( f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." ) @@ -261,6 +271,9 @@ def load_model_from_full_model_state_dict( sharded_tensor = temp_param.data else: sharded_tensor = full_tensor + + if cpu_offload: + sharded_tensor = sharded_tensor.cpu() else: full_tensor = full_tensor.to(device=device, dtype=param_dtype) sharded_tensor = distribute_tensor( @@ -296,6 +309,8 @@ def load_model_from_full_model_state_dict( sharded_tensor = torch.zeros_like( meta_sharded_param, device=device, dtype=param_dtype ) + if cpu_offload: + sharded_tensor = sharded_tensor.cpu() else: # Initialize with zeros and distribute full_tensor = torch.zeros_like( diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 869068e9b..af537835c 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -349,7 +349,8 @@ OOM detected. Possible solutions: - If the OOM occurs during runtime: 1. Reduce the number of output tokens by lowering resolution or decreasing `--num-frames` 2. Enable SP and/or TP - 3. Enable a sparse-attention backend + 3. Opt for a sparse-attention backend + 4. Enable FSDP by `--use-fsdp-inference` (in a multi-GPU setup) Or, open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose """ @@ -402,7 +403,7 @@ def run_scheduler_process( ) scheduler.event_loop() except torch.OutOfMemoryError as _e: - print(OOM_MSG) + logger.warning(OOM_MSG) raise finally: # Clean up resources to speed up shutdown diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index c3d71c8a1..48ae6da56 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -381,6 +381,7 @@ class RopeEmbedder: class ZImageTransformer2DModel(CachableDiT, OffloadableDiTMixin): _supports_gradient_checkpointing = True _no_split_modules = ["ZImageTransformerBlock"] + _fsdp_shard_conditions = ZImageDitConfig().arch_config._fsdp_shard_conditions param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index e58040b60..702690624 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -846,6 +846,10 @@ class DenoisingStage(PipelineStage): if not server_args.dit_cpu_offload: return + # FSDP manages offloading internally + if server_args.use_fsdp_inference: + return + # Offload the unused model if it's on CUDA if ( model_to_offload is not None diff --git a/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py b/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py index c914e4cd2..45d0c6edc 100644 --- a/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py +++ b/python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py @@ -67,6 +67,9 @@ def _build_server_extra_args(case: DiffusionTestCase) -> str: a += f" --lora-path {server_args.lora_path}" if server_args.warmup: a += " --warmup" + + for extra_arg in server_args.extras: + a += f" {extra_arg}" return a diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index f7317283d..e54158965 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -1997,6 +1997,31 @@ "expected_e2e_ms": 24895.28, "expected_avg_denoise_ms": 596.59, "expected_median_denoise_ms": 599.66 + }, + "fsdp-inference": { + "stages_ms": { + "InputValidationStage": 0.04, + "TextEncodingStage": 128.3, + "ConditioningStage": 0.01, + "TimestepPreparationStage": 1.44, + "LatentPreparationStage": 0.1, + "DenoisingStage": 1569.61, + "DecodingStage": 41.43 + }, + "denoise_step_ms": { + "0": 165.33, + "1": 158.34, + "2": 167.65, + "3": 179.11, + "4": 183.98, + "5": 175.08, + "6": 178.34, + "7": 178.53, + "8": 178.08 + }, + "expected_e2e_ms": 1742.7, + "expected_avg_denoise_ms": 173.83, + "expected_median_denoise_ms": 178.08 } } } diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 7cb1c713d..9b7b6f4eb 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -90,6 +90,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: if server_args.warmup: extra_args += f" --warmup" + for arg in server_args.extras: + extra_args += f" {arg}" + # Build custom environment variables env_vars = {} if server_args.enable_cache_dit: diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index a73bf307c..9379dcc9d 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -21,7 +21,7 @@ from __future__ import annotations import json import os import statistics -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Sequence @@ -151,7 +151,7 @@ class BaselineConfig: return self -@dataclass(frozen=True) +@dataclass class DiffusionServerArgs: """Configuration for a single model/scenario test case.""" @@ -183,6 +183,14 @@ class DiffusionServerArgs: enable_cache_dit: bool = False text_encoder_cpu_offload: bool = False + extras: list[str] = field(default_factory=lambda: []) + + def __post_init__(self): + if self.modality == "image": + self.custom_validator = "image" + elif self.modality == "video": + self.custom_validator = "video" + @dataclass(frozen=True) class DiffusionSamplingParams: @@ -331,6 +339,8 @@ TURBOWAN_I2V_sampling_params = DiffusionSamplingParams( fps=4, ) +DEFAULT_SMALL_MODEL = "Tongyi-MAI/Z-Image-Turbo" + # All test cases with clean default values # To test different models, simply add more DiffusionCase entries ONE_GPU_CASES_A: list[DiffusionTestCase] = [ @@ -644,6 +654,17 @@ TWO_GPU_CASES_A = [ prompt=T2V_PROMPT, ), ), + DiffusionTestCase( + "fsdp-inference", + DiffusionServerArgs( + model_path=DEFAULT_SMALL_MODEL, + modality="image", + num_gpus=2, + warmup=True, + extras=["--use-fsdp-inference"], + ), + T2I_sampling_params, + ), ] # Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.