feat: fit the 18B model on a 24 GB GPU + 64 GB RAM box (staged T5, eager FP8, FP8-only deepcopy)

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:

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
 
 
@@ -203,11 +203,48 @@ 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)
 
+    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))
 
@@ -216,10 +253,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()
@@ -231,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):
@@ -267,19 +298,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