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

Signed-off-by: Lancer <maruxiang6688@gmail.com>
Co-authored-by: Lancer <maruxiang6688@gmail.com>

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:

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,

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

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)