support phase offload for qwen-image

Author: gushiqiao

lightx2v/models/input_encoders/hf/qwen25/qwen25_vlforconditionalgeneration.py

@@ -76,8 +76,6 @@ def load(self):
             if self.config["cpu_offload"]:
                 self.device_map = {
                     "lm_head": AI_DEVICE,
-                    "model.embed_tokens": AI_DEVICE,
-                    "model.norm": AI_DEVICE,
                     "model.visual": "cpu",
                     "model.language_model": "cpu",
                 }

lightx2v/models/networks/qwen_image/infer/offload/transformer_infer.py

@@ -1,7 +1,9 @@
 import torch
 
 from lightx2v.common.offload.manager import WeightAsyncStreamManager
-from lightx2v.models.networks.qwen_image.infer.transformer_infer import QwenImageTransformerInfer
+from lightx2v.models.networks.qwen_image.infer.transformer_infer import (
+    QwenImageTransformerInfer,
+)
 from lightx2v_platform.base.global_var import AI_DEVICE
 
 torch_device_module = getattr(torch, AI_DEVICE)
@@ -11,6 +13,7 @@ class QwenImageOffloadTransformerInfer(QwenImageTransformerInfer):
     def __init__(self, config):
         super().__init__(config)
         self.num_blocks = config["num_layers"]
+        self.phases_num = 3
         if self.config.get("cpu_offload", False):
             if "offload_ratio" in self.config:
                 self.offload_ratio = self.config["offload_ratio"]
@@ -20,12 +23,94 @@ def __init__(self, config):
             if offload_granularity == "block":
                 self.infer_func = self.infer_with_blocks_offload
                 self.offload_manager = WeightAsyncStreamManager(offload_granularity=offload_granularity)
+            elif offload_granularity == "phase":
+                self.infer_func = self.infer_with_phases_offload
+                self.offload_manager = WeightAsyncStreamManager(offload_granularity=offload_granularity)
 
             self.lazy_load = self.config.get("lazy_load", False)
             if self.lazy_load:
                 self.offload_manager.init_lazy_load(num_workers=self.config.get("num_disk_workers", 4))
 
-    def infer_with_blocks_offload(self, block_weights, hidden_states, encoder_hidden_states, temb_img_silu, temb_txt_silu, image_rotary_emb, modulate_index):
+    def infer_with_phases_offload(
+        self,
+        blocks,
+        hidden_states,
+        encoder_hidden_states,
+        temb_img_silu,
+        temb_txt_silu,
+        image_rotary_emb,
+        modulate_index,
+    ):
+        for block_idx in range(len(blocks)):
+            self.block_idx = block_idx
+            if self.lazy_load:
+                next_prefetch = (block_idx + 1) % len(blocks)
+                self.offload_manager.start_prefetch_block(next_prefetch)
+
+            for phase_idx in range(self.phases_num):
+                # if self.offload_manager.need_init_first_buffer:
+                if block_idx == 0 and phase_idx == 0:
+                    self.offload_manager.init_first_buffer(blocks)
+
+                next_block_idx = (block_idx + 1) % len(blocks) if phase_idx == self.phases_num - 1 else block_idx
+                next_phase_idx = (phase_idx + 1) % self.phases_num
+                if self.lazy_load:
+                    if phase_idx == self.phases_num - 1:
+                        self.offload_manager.swap_cpu_buffers()
+
+                self.offload_manager.prefetch_phase(next_block_idx, next_phase_idx, blocks)
+                with torch_device_module.stream(self.offload_manager.compute_stream):
+                    if phase_idx == 0:
+                        (
+                            img_modulated,
+                            txt_modulated,
+                            img_gate1,
+                            txt_gate1,
+                            img_mod2,
+                            txt_mod2,
+                        ) = self.infer_modulate(
+                            mod_phase=self.offload_manager.cuda_buffers[phase_idx],
+                            hidden_states=hidden_states,
+                            encoder_hidden_states=encoder_hidden_states,
+                            temb_img_silu=temb_img_silu,
+                            temb_txt_silu=temb_txt_silu,
+                            modulate_index=modulate_index,
+                        )
+                    elif phase_idx == 1:
+                        hidden_states, encoder_hidden_states = self.infer_attn(
+                            attn_phase=self.offload_manager.cuda_buffers[phase_idx],
+                            img_modulated=img_modulated,
+                            txt_modulated=txt_modulated,
+                            img_gate1=img_gate1,
+                            txt_gate1=txt_gate1,
+                            hidden_states=hidden_states,
+                            encoder_hidden_states=encoder_hidden_states,
+                            image_rotary_emb=image_rotary_emb,
+                        )
+                    elif phase_idx == 2:
+                        encoder_hidden_states, hidden_states = self.infer_ffn(
+                            ffn_phase=self.offload_manager.cuda_buffers[phase_idx],
+                            hidden_states=hidden_states,
+                            encoder_hidden_states=encoder_hidden_states,
+                            img_mod2=img_mod2,
+                            txt_mod2=txt_mod2,
+                            modulate_index=modulate_index,
+                        )
+
+                self.offload_manager.swap_phases()
+
+        return hidden_states
+
+    def infer_with_blocks_offload(
+        self,
+        blocks,
+        hidden_states,
+        encoder_hidden_states,
+        temb_img_silu,
+        temb_txt_silu,
+        image_rotary_emb,
+        modulate_index,
+    ):
         for block_idx in range(self.num_blocks):
             self.block_idx = block_idx
 
@@ -34,15 +119,15 @@ def infer_with_blocks_offload(self, block_weights, hidden_states, encoder_hidden
                 self.offload_manager.start_prefetch_block(next_prefetch)
 
             if block_idx == 0:
-                self.offload_manager.init_first_buffer(block_weights.blocks)
+                self.offload_manager.init_first_buffer(blocks)
 
             if self.lazy_load:
                 self.offload_manager.swap_cpu_buffers()
-            self.offload_manager.prefetch_weights((block_idx + 1) % self.num_blocks, block_weights.blocks)
+            self.offload_manager.prefetch_weights((block_idx + 1) % self.num_blocks, blocks)
 
             with torch_device_module.stream(self.offload_manager.compute_stream):
                 encoder_hidden_states, hidden_states = self.infer_block(
-                    block_weight=self.offload_manager.cuda_buffers[0],
+                    block=self.offload_manager.cuda_buffers[0],
                     hidden_states=hidden_states,
                     encoder_hidden_states=encoder_hidden_states,
                     temb_img_silu=temb_img_silu,