[diffusion] fix: fix fsdp (sgl-project#18187)

Author: mickqian

python/sglang/multimodal_gen/configs/models/dits/zimage.py

@@ -7,6 +7,17 @@
 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 @@ def forward_cuda(
         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 @@ def forward_cuda(
             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 @@ def __init__(
         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 @@ def forward(
             )
         )
         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 @@ def load_model(
             # 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 @@ def list_loras(self) -> OutputBatch:
   - 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 @@ def __call__(self, ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
 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 @@ def _manage_device_placement(
         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