From 7c56d1bc7b83a0a2cbf00f300233d27151be63e9 Mon Sep 17 00:00:00 2001 From: yangohuang <16739522+yangohuang@users.noreply.github.com> Date: Wed, 22 Jul 2026 19:20:09 +0800 Subject: [PATCH 1/3] feat: pre-encode prompts and release T5 before loading the DiT umT5-xxl (~11 GB) stays resident for the whole run even though every prompt is known upfront. On 64 GB-RAM hosts this collides with the 18B DiT (~36 GB bf16) during loading and the process gets OOM-killed. Encode all prompts (including edit_prompt) in a staging pass, free the T5 encoder, then load the DiT with low_cpu_mem_usage=True. Peak host memory drops by ~11 GB; outputs are unchanged for every input json. --- generate.py | 35 +++++++++++++++++++++++------------ 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/generate.py b/generate.py index 96fb83b..afdbd70 100644 --- a/generate.py +++ b/generate.py @@ -203,7 +203,27 @@ def generate(args): } for layer_id in range(40)} for i in range(len(timesteps) - 1) } - wan_i2v_model = WanModel.from_pretrained(args.ckpt_dir, torch_dtype=torch.bfloat16, low_cpu_mem_usage=False) + # Pre-encode all prompts with T5, then release it BEFORE loading the DiT. + # umT5-xxl (~11 GB) and the 18B DiT (~36 GB bf16) co-resident on the host + # will trigger the OOM killer on 64 GB-RAM machines. + with open(args.input_json, 'r', encoding='utf-8') as f: + _pre_input_data = json.load(f) + text_encoder = T5EncoderModel(text_len=512, dtype=torch.bfloat16, device='cpu' if args.t5_cpu else device, + checkpoint_path=os.path.join(args.ckpt_dir, 'models_t5_umt5-xxl-enc-bf16.pth'), + tokenizer_path=os.path.join(args.ckpt_dir, 'google/umt5-xxl')) + precomputed_ctx = [] + for _data in _pre_input_data: + _ctx = [text_encoder(texts=_data['prompt'], device='cpu' if args.t5_cpu else device)[0].to(device, dtype=torch.bfloat16)] + _edit = { + k: text_encoder(texts=v, device='cpu' if args.t5_cpu else device)[0].to(device, dtype=torch.bfloat16) + for k, v in _data.get('edit_prompt', {}).items() + } + precomputed_ctx.append((_ctx, _edit)) + text_encoder.model = None + del text_encoder + torch_gc() + + wan_i2v_model = WanModel.from_pretrained(args.ckpt_dir, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True) wan_i2v_model = wan_i2v_model.to(dtype=torch.bfloat16) for n in range(40): wan_i2v_model.blocks[n].self_attn.init_kvidx(frame_len, world_size) @@ -216,10 +236,6 @@ def generate(args): tokenizer_path=os.path.join(args.ckpt_dir, 'xlm-roberta-large'), dtype=torch.bfloat16, device=device) clip.model = clip.model.to(device, dtype=torch.bfloat16) - text_encoder = T5EncoderModel(text_len=512, dtype=torch.bfloat16, device='cpu' if args.t5_cpu else device, - checkpoint_path=os.path.join(args.ckpt_dir, 'models_t5_umt5-xxl-enc-bf16.pth'), - tokenizer_path=os.path.join(args.ckpt_dir, 'google/umt5-xxl')) - audio_encoder = Wav2Vec2Model.from_pretrained( args.wav2vec_dir, local_files_only=True, torch_dtype=torch.bfloat16 ).to(device, dtype=torch.bfloat16).eval() @@ -267,19 +283,14 @@ def filter_fn(name, module): with open(args.input_json, 'r', encoding='utf-8') as f: input_data = json.load(f) - for data in input_data: + for _item_idx, data in enumerate(input_data): image_path = data['cond_image'] audio_path = data['cond_audio'] out_path = os.path.basename(image_path).split('.')[0] + '_' + os.path.basename(audio_path).split('.')[0] + '.mp4' prompt = data['prompt'] edit_prompts = data.get('edit_prompt', {}) - context = [text_encoder(texts=prompt, device='cpu' if args.t5_cpu else device)[0].to(device, dtype=torch.bfloat16)] - if edit_prompts: - edit_prompts = { - k: text_encoder(texts=v, device='cpu' if args.t5_cpu else device)[0].to(device, dtype=torch.bfloat16) - for k, v in edit_prompts.items() - } + context, edit_prompts = precomputed_ctx[_item_idx] image = Image.open(image_path).convert("RGB") cond_image = transform(image).unsqueeze(1).unsqueeze(0).to(device, torch.bfloat16) # 1 C 1 H W From 2b96986fb28ea42617ce0e52f91b495d74555a74 Mon Sep 17 00:00:00 2001 From: yangohuang <16739522+yangohuang@users.noreply.github.com> Date: Wed, 22 Jul 2026 19:20:28 +0800 Subject: [PATCH 2/3] fix: support FP8-only source in FP8Linear.__deepcopy__ WanBlockOffloadManager deep-copies blocks[0] for its staging buffers. If fp8 weights were materialized with fp16_weight_storage='discard' before enable_block_offload() is called, __deepcopy__ raises RuntimeError because no fp16 source remains, even though the fp8 buffers are sufficient to reconstruct the module. Rebuild from a zero placeholder in that case and clone the fp8 buffers (the existing clone path below already handles them), then drop the placeholder copies. --- fp8_gemm.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/fp8_gemm.py b/fp8_gemm.py index 568762b..cfd8daf 100644 --- a/fp8_gemm.py +++ b/fp8_gemm.py @@ -141,8 +141,15 @@ def __deepcopy__(self, memo): elif self._fp16_weight_cpu is not None: src_weight = self._fp16_weight_cpu.detach() src_bias = self._fp16_bias_cpu.detach() if self._fp16_bias_cpu is not None else None + elif self._fp8_weight is not None: + # FP8-only source (fp16 discarded): build the module from a zero + # placeholder; the real fp8 buffers are cloned below and the + # placeholder copies are dropped before returning. + k_in, n_out = self._fp8_weight.shape # _fp8_weight is the [K, N] transposed view + src_weight = torch.zeros(n_out, k_in, dtype=torch.bfloat16) + src_bias = self.bias.detach() if self.bias is not None else None else: - raise RuntimeError("FP8Linear cannot be deep-copied without an FP16 weight source.") + raise RuntimeError("FP8Linear cannot be deep-copied without an FP16 or FP8 weight source.") linear = nn.Linear( in_features=src_weight.shape[1], @@ -170,6 +177,9 @@ def __deepcopy__(self, memo): cloned._weight_cache_device = self._weight_cache_device cloned._last_weight_version = self._last_weight_version + if self.linear is None and self._fp16_weight_cpu is None: + cloned._fp16_weight_cpu = None + cloned._fp16_bias_cpu = None return cloned def invalidate_weight_cache(self) -> None: From e6d04513b2c248b694708e83be1ecc41ea273272 Mon Sep 17 00:00:00 2001 From: yangohuang <16739522+yangohuang@users.noreply.github.com> Date: Wed, 22 Jul 2026 19:20:53 +0800 Subject: [PATCH 3/3] feat: eager per-block FP8 materialization under --fp8_gemm --block_offload With block offloading, weights live on the host and only fp8 buffers need to travel to the GPU staging blocks each step. Keeping the bf16 originals (~36 GB for the 18B DiT) on the host until lazy first-forward quantization leaves a long window where VAE/CLIP/wav2vec loads push a 64 GB-RAM machine into the OOM killer. Quantize block by block on GPU right after the DiT loads, discarding each bf16 copy immediately: host steady-state drops to ~18 GB and the per-step H2D copy volume is halved. Gated on --block_offload so the datacenter path (no offload) keeps its existing lazy behavior. Requires the FP8-only __deepcopy__ path from the previous commit, since enable_block_offload() deep-copies blocks[0] after the bf16 weights are gone. --- generate.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/generate.py b/generate.py index afdbd70..6517e8e 100644 --- a/generate.py +++ b/generate.py @@ -24,7 +24,7 @@ from src.audio_analysis.wav2vec2 import Wav2Vec2Model from diffusers.utils import export_to_video -from fp8_gemm import FP8GemmOptions, enable_fp8_gemm +from fp8_gemm import FP8GemmOptions, FP8Linear, enable_fp8_gemm from fp4_gemm import FP4GemmOptions, enable_fp4_gemm @@ -228,6 +228,23 @@ def generate(args): for n in range(40): wan_i2v_model.blocks[n].self_attn.init_kvidx(frame_len, world_size) + if args.fp8_gemm: + enable_fp8_gemm(wan_i2v_model, options=FP8GemmOptions()) + if args.block_offload: + # On the consumer/offload path, quantize to FP8 immediately after the + # DiT loads, block by block on GPU, discarding each bf16 copy as we go. + # Host steady-state drops from ~36 GB to ~18 GB before the VAE/CLIP/ + # wav2vec loads, and the per-step host-to-device copy volume of the + # offload double-buffer is halved. Without --block_offload the wrap + # stays lazy, exactly as before. + _quant_dev = torch.device(f"cuda:{device}") + for _blk in wan_i2v_model.blocks: + for _m in _blk.modules(): + if isinstance(_m, FP8Linear): + _m.materialize_fp8_weight(_quant_dev) + _blk.to('cpu') + torch_gc() + vae = LightVAE(vae_path=os.path.join(args.ckpt_dir, 'Wan2.1_VAE.pth'), dtype=torch.bfloat16, device=device, use_lightvae=False, parallel=(world_size > 1)) @@ -247,8 +264,6 @@ def generate(args): for name, param in _model.named_parameters(): param.requires_grad = False - if args.fp8_gemm: - enable_fp8_gemm(wan_i2v_model, options=FP8GemmOptions()) if args.fp4_gemm: print("Enabling FP4 GEMM acceleration...") def filter_fn(name, module):