From e9979950376aa2137ce9231ff370992b3a6a18f8 Mon Sep 17 00:00:00 2001 From: Lancer <402430575@qq.com> Date: Thu, 15 Jan 2026 21:28:57 +0800 Subject: [PATCH] [diffusion] fix: optimize text encoder CPU offload initialization to address OOM (#17064) Signed-off-by: Lancer Co-authored-by: Lancer --- .../runtime/loader/component_loader.py | 14 ++-- .../test/server/perf_baselines.json | 67 +++++++++++++++++++ .../test/server/test_server_common.py | 3 + .../test/server/testcase_configs.py | 14 ++++ 4 files changed, 94 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index da1dde89c..a704c4297 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -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: diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index ee7ff3c5e..474315b54 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -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, diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 66c29d627..1614aeec8 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -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}" diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 907968c7a..3281ace17 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -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)