[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>
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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}"
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user