From a43284ae1bd12ba3008dd34bf57c3df2f60f7e04 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Sun, 5 Jul 2026 03:17:11 +0000 Subject: [PATCH] [bugfix]: swap cpu-offloaded text encoders to GPU for the forward text_encoder_cpu_offload means params rest on CPU between calls but compute on GPU: FSDP-wrapped encoders (T5/llama/clip) already get this via CPUOffloadPolicy streaming, but encoders without _fsdp_shard_conditions (reason1/Qwen2.5-VL) were parked on CPU and, after the encoder-device input routing added for kandinsky5, silently ran their 7B forward on CPU (crashing outright with flash-attn installed). Move plain CPU-parked encoders to the target device for the forward and back afterwards, mirroring the image-encoder/VAE offload pattern; FSDP (DTensor) encoders keep the existing input routing. Revert the examples to text_encoder_cpu_offload=True now that the default works. Verified e2e: T2V/I2V outputs are metric-identical to GPU-resident runs, with ~12GB lower VRAM during denoising. Co-Authored-By: Claude Fable 5 --- .../inference/basic/basic_kandinsky5_i2v.py | 2 +- .../inference/basic/basic_kandinsky5_t2v.py | 2 +- fastvideo/pipelines/stages/text_encoding.py | 23 +++++++++++++++---- 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/examples/inference/basic/basic_kandinsky5_i2v.py b/examples/inference/basic/basic_kandinsky5_i2v.py index 861ffedac..c30ca2f48 100644 --- a/examples/inference/basic/basic_kandinsky5_i2v.py +++ b/examples/inference/basic/basic_kandinsky5_i2v.py @@ -14,7 +14,7 @@ def main(): use_fsdp_inference=False, dit_cpu_offload=False, vae_cpu_offload=False, - text_encoder_cpu_offload=False, + text_encoder_cpu_offload=True, pin_cpu_memory=True, # image_encoder_cpu_offload=False, ) diff --git a/examples/inference/basic/basic_kandinsky5_t2v.py b/examples/inference/basic/basic_kandinsky5_t2v.py index 45e8aa72b..8042f87fd 100644 --- a/examples/inference/basic/basic_kandinsky5_t2v.py +++ b/examples/inference/basic/basic_kandinsky5_t2v.py @@ -12,7 +12,7 @@ def main(): use_fsdp_inference=False, dit_cpu_offload=False, vae_cpu_offload=False, - text_encoder_cpu_offload=False, + text_encoder_cpu_offload=True, pin_cpu_memory=True, # image_encoder_cpu_offload=False, ) diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py index 95b888f8f..0a5427ee8 100644 --- a/fastvideo/pipelines/stages/text_encoding.py +++ b/fastvideo/pipelines/stages/text_encoding.py @@ -8,6 +8,8 @@ This module contains implementations of prompt encoding stages for diffusion pip import torch from typing import Any +from torch.distributed.tensor import DTensor + from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context @@ -203,10 +205,21 @@ class TextEncodingStage(PipelineStage): encoder_config = encoder_cfgs[i] preprocess_func = preprocess_funcs[i] postprocess_func = postprocess_funcs[i] - encoder_device = next( - (param.device for param in text_encoder.parameters()), - torch.device(target_device), - ) + # cpu_offload semantics: params rest on CPU between calls but the + # forward computes on GPU. FSDP2-wrapped encoders (CPUOffloadPolicy; + # DTensor params) stream themselves per-layer — leave inputs on the + # param device and let FSDP's root pre-forward move them. A plain + # module parked on CPU by text_encoder_cpu_offload is swapped to the + # target device for the forward and back afterwards, mirroring the + # image-encoder/VAE offload pattern. + first_param = next(text_encoder.parameters(), None) + encoder_device = first_param.device if first_param is not None else torch.device(target_device) + moved_for_forward = False + if (first_param is not None and not isinstance(first_param, DTensor) + and encoder_device.type != torch.device(target_device).type): + text_encoder = text_encoder.to(target_device) + encoder_device = torch.device(target_device) + moved_for_forward = True tok_kwargs = dict(encoder_config.tokenizer_kwargs) if max_length is not None: @@ -297,6 +310,8 @@ class TextEncodingStage(PipelineStage): embeds_list.append(prompt_embeds) if return_attention_mask: attn_masks_list.append(attention_mask.to(device=target_device)) + if moved_for_forward and fastvideo_args.text_encoder_cpu_offload: + text_encoder.to("cpu") self._last_audio_embeds = audio_embeds_list if is_ltx2 else None return self.return_embeds(embeds_list, attn_masks_list, return_type, return_attention_mask, indices)