[diffusion] fix: fix fsdp (#18187)

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)

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(

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()

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
 
 

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}"
                 )

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(

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

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

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

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
 
 

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
         }
     }
 }

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:

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.