[diffusion] fix: fix fsdp (#18187)

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