chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] fix: fix fsdp (#18187)

Browse Source
This commit is contained in:
Mick
2026-02-10 20:22:20 +08:00
committed by GitHub
parent 49cbb469b4
commit efcdda0176
13 changed files with 102 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/multimodal_gen/configs/models/dits/zimage.py
View File

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

5
python/sglang/multimodal_gen/runtime/layers/layernorm.py
View File

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

2
python/sglang/multimodal_gen/runtime/layers/lora/linear.py
View File

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

3
python/sglang/multimodal_gen/runtime/layers/triton_ops.py
View File

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

2
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
View File

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

19
python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
View File

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

5
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

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

1
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
View File

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

4
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

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

3
python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py
View File

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

25
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

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

3
python/sglang/multimodal_gen/test/server/test_server_common.py
View File

@@ -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:

25
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -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.