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

[diffusion] fix: optimize text encoder CPU offload initialization to address OOM (#17064)

Browse Source 
Signed-off-by: Lancer <maruxiang6688@gmail.com>
Co-authored-by: Lancer <maruxiang6688@gmail.com>
This commit is contained in:
Lancer
2026-01-15 21:28:57 +08:00
committed by GitHub
parent 6586f44ad4
commit e997995037
4 changed files with 94 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/sglang/multimodal_gen/runtime/loader/component_loader.py
View File

@@ -459,8 +459,14 @@ class TextEncoderLoader(ComponentLoader):
local_torch_device = get_local_torch_device()
should_offload = self.should_offload(server_args, model_config)
if should_offload and not current_platform.is_mps():
model_device = torch.device("cpu")
else:
model_device = local_torch_device
with set_default_torch_dtype(PRECISION_TO_TYPE[dtype]):
with local_torch_device, skip_init_modules():
with model_device, skip_init_modules():
architectures = getattr(model_config, "architectures", [])
model_cls, _ = ModelRegistry.resolve_model_cls(architectures)
enable_image_understanding = (
@@ -478,15 +484,13 @@ class TextEncoderLoader(ComponentLoader):
self._get_all_weights(model, model_path, to_cpu=should_offload)
)
# Explicitly move model to target device after loading weights
model = model.to(local_torch_device)
if should_offload:
# Disable FSDP for MPS as it's not compatible
if current_platform.is_mps():
logger.info(
"Disabling FSDP sharding for MPS platform as it's not compatible"
)
model = model.to(local_torch_device)
else:
mesh = init_device_mesh(
current_platform.device_type,
@@ -502,6 +506,8 @@ class TextEncoderLoader(ComponentLoader):
or getattr(model, "_fsdp_shard_conditions", None),
pin_cpu_memory=server_args.pin_cpu_memory,
)
else:
model = model.to(local_torch_device)
# We only enable strict check for non-quantized models
# that have loaded weights tracking currently.
# if loaded_weights is not None:

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

@@ -804,6 +804,73 @@
"expected_avg_denoise_ms": 260.76,
"expected_median_denoise_ms": 247.84
},
"wan2_1_t2v_1.3b_text_encoder_cpu_offload": {
"stages_ms": {
"InputValidationStage": 0.09,
"TextEncodingStage": 4005.3,
"ConditioningStage": 0.1,
"TimestepPreparationStage": 4.99,
"LatentPreparationStage": 5.49,
"DenoisingStage": 7932.25,
"DecodingStage": 902.53,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 964.7,
"1": 147.48,
"2": 191.94,
"3": 138.0,
"4": 149.09,
"5": 186.5,
"6": 126.5,
"7": 135.06,
"8": 167.13,
"9": 148.28,
"10": 130.8,
"11": 148.97,
"12": 164.98,
"13": 126.23,
"14": 140.33,
"15": 138.65,
"16": 136.43,
"17": 149.42,
"18": 157.08,
"19": 179.23,
"20": 154.54,
"21": 139.79,
"22": 153.94,
"23": 129.32,
"24": 122.91,
"25": 134.7,
"26": 133.6,
"27": 131.25,
"28": 122.71,
"29": 133.91,
"30": 134.5,
"31": 122.44,
"32": 123.12,
"33": 122.11,
"34": 151.54,
"35": 138.77,
"36": 149.28,
"37": 163.16,
"38": 126.12,
"39": 127.19,
"40": 145.81,
"41": 148.05,
"42": 153.0,
"43": 132.07,
"44": 148.57,
"45": 129.11,
"46": 129.98,
"47": 126.94,
"48": 130.1,
"49": 129.64
},
"expected_e2e_ms": 12875.42,
"expected_avg_denoise_ms": 158.3,
"expected_median_denoise_ms": 138.32
},
"wan2_1_t2v_1.3b_cfg_parallel": {
"stages_ms": {
"InputValidationStage": 0.09,

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

@@ -76,6 +76,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
if server_args.dit_layerwise_offload:
extra_args += f" --dit-layerwise-offload true"
if server_args.text_encoder_cpu_offload:
extra_args += f" --text-encoder-cpu-offload"
if server_args.ring_degree is not None:
extra_args += f" --ring-degree {server_args.ring_degree}"

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

@@ -165,6 +165,7 @@ class DiffusionServerArgs:
dit_layerwise_offload: bool = False
enable_cache_dit: bool = False
text_encoder_cpu_offload: bool = False
@dataclass(frozen=True)
@@ -435,6 +436,19 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
prompt=T2V_PROMPT,
),
),
DiffusionTestCase(
"wan2_1_t2v_1.3b_text_encoder_cpu_offload",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
modality="video",
warmup=0,
custom_validator="video",
text_encoder_cpu_offload=True,
),
DiffusionSamplingParams(
prompt=T2V_PROMPT,
),
),
# LoRA test case for single transformer + merge/unmerge API test
# Note: Uses dynamic_lora_path instead of lora_path to test LayerwiseOffload + set_lora interaction
# Server starts WITHOUT LoRA, then set_lora is called after startup (Wan models auto-enable layerwise offload)