support w4 for qw_vl and disk offload for qw_2511 (#681)

Co-authored-by: gushiqiao <975033167>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: gushiqiao

lightx2v/common/offload/manager.py

@@ -50,7 +50,10 @@ def init_cuda_buffer(self, blocks_cuda_buffer=None, phases_cuda_buffer=None):
     def init_first_buffer(self, blocks, adapter_block_idx=None):
         with torch_device_module.stream(self.init_stream):
             if hasattr(self, "cpu_buffers"):
-                self.cuda_buffers[0].load_state_dict(self.cpu_buffers[0][0].state_dict(), 0, adapter_block_idx)
+                if self.offload_granularity == "block":
+                    self.cuda_buffers[0].load_state_dict(self.cpu_buffers[0].state_dict(), 0, adapter_block_idx)
+                else:
+                    self.cuda_buffers[0].load_state_dict(self.cpu_buffers[0][0].state_dict(), 0, adapter_block_idx)
             else:
                 if self.offload_granularity == "block":
                     self.cuda_buffers[0].load_state_dict(blocks[0].state_dict(), 0, adapter_block_idx)
@@ -62,8 +65,7 @@ def init_first_buffer(self, blocks, adapter_block_idx=None):
     def prefetch_weights(self, block_idx, blocks, adapter_block_idx=None):
         with torch_device_module.stream(self.cuda_load_stream):
             if hasattr(self, "cpu_buffers"):
-                self.cpu_buffers[1].load_state_dict_from_disk(block_idx, adapter_block_idx)
-                self.cuda_buffers[1].load_state_dict(self.cpu_buffers[1].state_dict(), block_idx, adapter_block_idx)
+                self.cuda_buffers[1].load_state_dict(self.cpu_buffers[0].state_dict(), block_idx, adapter_block_idx)
             else:
                 self.cuda_buffers[1].load_state_dict(blocks[block_idx].state_dict(), block_idx, adapter_block_idx)
 
@@ -110,12 +112,17 @@ def init_lazy_load(self, num_workers=6):
     def start_prefetch_block(self, block_idx, adapter_block_idx=None):
         self.prefetch_block_idx = block_idx
         self.prefetch_futures = []
-        for phase in self.cpu_buffers[1]:
-            future = self.executor.submit(phase.load_state_dict_from_disk, block_idx, adapter_block_idx)
+        if self.offload_granularity == "block":
+            future = self.executor.submit(self.cpu_buffers[1].load_state_dict_from_disk, block_idx, adapter_block_idx)
             self.prefetch_futures.append(future)
+        else:
+            for phase in self.cpu_buffers[1]:
+                future = self.executor.submit(phase.load_state_dict_from_disk, block_idx, adapter_block_idx)
+                self.prefetch_futures.append(future)
 
     def swap_cpu_buffers(self):
-        #  wait_start = time.time()
+        # import time
+        # wait_start = time.time()
         # already_done = all(f.done() for f in self.prefetch_futures)
         for f in self.prefetch_futures:
             f.result()

lightx2v/common/ops/mm/mm_weight.py

@@ -1531,15 +1531,15 @@ def __init__(
         if ops is not None:
             self.act_quant_func = self.act_quant_fp8_perchannel_sym_vllm
         else:
-            self.act_quant_func = self.fp8_quantize_triton
+            self.act_quant_func = fp8_quantize_triton
 
     def apply(self, input_tensor):
         input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor)
         output_tensor = fp8_linear(
             input_tensor_quant,
             self.weight,
             self.bias.float() if self.bias is not None else None,
-            input_tensor_scale,
+            input_tensor_scale.float(),
             self.weight_scale,
             out_dtype=self.infer_dtype,
         )
@@ -1582,15 +1582,15 @@ def __init__(
         if ops is not None:
             self.act_quant_func = self.act_quant_int8_perchannel_sym_vllm
         else:
-            self.act_quant_func = self.int8_quantize_triton
+            self.act_quant_func = int8_quantize_triton
 
     def apply(self, input_tensor):
         input_tensor_quant, input_tensor_scale = self.act_quant_func(input_tensor)
         output_tensor = q8_linear(
             input_tensor_quant,
             self.weight,
             self.bias.float() if self.bias is not None else None,
-            input_tensor_scale,
+            input_tensor_scale.float(),
             self.weight_scale,
             fuse_gelu=False,
             out_dtype=self.infer_dtype,

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

@@ -71,7 +71,20 @@ def __init__(self, config):
         self.load()
 
     def load(self):
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(os.path.join(self.config["model_path"], "text_encoder"), torch_dtype=torch.bfloat16)
+        if self.config.get("qwen25vl_quantized", False):
+            assert self.config["qwen25vl_quant_scheme"] == "int4"
+            if self.config["cpu_offload"]:
+                self.device_map = {
+                    "lm_head": AI_DEVICE,
+                    "model.visual": "cpu",
+                    "model.language_model": "cpu",
+                }
+            else:
+                self.device_map = "auto"
+            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(self.config["qwen25vl_quantized_ckpt"], dtype=torch.bfloat16, device_map=self.device_map, low_cpu_mem_usage=True)
+        else:
+            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(os.path.join(self.config["model_path"], "text_encoder"), torch_dtype=torch.bfloat16)
+
         if not self.cpu_offload:
             self.text_encoder = self.text_encoder.to(AI_DEVICE)
 
@@ -99,7 +112,8 @@ def preprocess_image(self, image):
     @torch.no_grad()
     def infer(self, text, image_list=None):
         if self.cpu_offload:
-            self.text_encoder.to(AI_DEVICE)
+            if not hasattr(self, "device_map") or self.device_map == "auto":
+                self.text_encoder.to(AI_DEVICE)
 
         if image_list is not None:
             condition_image_list = []
@@ -143,7 +157,6 @@ def infer(self, text, image_list=None):
                 image_grid_thw=model_inputs.image_grid_thw,
                 output_hidden_states=True,
             )
-
             image_info = {
                 "condition_image_list": condition_image_list,
                 "vae_image_list": vae_image_list,
@@ -183,7 +196,8 @@ def infer(self, text, image_list=None):
         prompt_embeds_mask = prompt_embeds_mask.view(1 * 1, seq_len)
 
         if self.cpu_offload:
-            self.text_encoder.to(torch.device("cpu"))
+            if not hasattr(self, "device_map") or self.device_map == "auto":
+                self.text_encoder.to(torch.device("cpu"))
             torch_device_module.empty_cache()
             gc.collect()
 

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)
@@ -10,35 +12,125 @@
 class QwenImageOffloadTransformerInfer(QwenImageTransformerInfer):
     def __init__(self, config):
         super().__init__(config)
-        self.phases_num = 3
         self.num_blocks = config["num_layers"]
+        self.phases_num = 4
         if self.config.get("cpu_offload", False):
             if "offload_ratio" in self.config:
                 self.offload_ratio = self.config["offload_ratio"]
             else:
                 self.offload_ratio = 1
             offload_granularity = self.config.get("offload_granularity", "block")
             if offload_granularity == "block":
-                if not self.config.get("lazy_load", False):
-                    self.infer_func = self.infer_with_blocks_offload
-                else:
-                    assert NotImplementedError
-
-            if offload_granularity != "model":
+                self.infer_func = self.infer_with_blocks_offload
                 self.offload_manager = WeightAsyncStreamManager(offload_granularity=offload_granularity)
-            else:
-                assert NotImplementedError
+            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_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()
 
-    def infer_with_blocks_offload(self, block_weights, hidden_states, encoder_hidden_states, temb_img_silu, temb_txt_silu, image_rotary_emb, modulate_index):
+                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_query, img_key, img_value, img_gate1, img_mod2 = self.infer_img_qkv(
+                            img_attn_phase=self.offload_manager.cuda_buffers[phase_idx],
+                            hidden_states=hidden_states,
+                            temb_img_silu=temb_img_silu,
+                            img_freqs=image_rotary_emb[0],
+                            modulate_index=modulate_index,
+                        )
+                    elif phase_idx == 1:
+                        txt_query, txt_key, txt_value, seq_txt, txt_gate1, txt_mod2 = self.infer_txt_qkv(
+                            txt_attn_phase=self.offload_manager.cuda_buffers[phase_idx],
+                            encoder_hidden_states=encoder_hidden_states,
+                            temb_txt_silu=temb_txt_silu,
+                            txt_freqs=image_rotary_emb[1],
+                        )
+                    elif phase_idx == 2:
+                        hidden_states, encoder_hidden_states = self.infer_cross_attn(
+                            cross_attn_phase=self.offload_manager.cuda_buffers[phase_idx],
+                            seq_txt=seq_txt,
+                            img_query=img_query,
+                            img_key=img_key,
+                            img_value=img_value,
+                            txt_query=txt_query,
+                            txt_key=txt_key,
+                            txt_value=txt_value,
+                            img_gate1=img_gate1,
+                            txt_gate1=txt_gate1,
+                            hidden_states=hidden_states,
+                            encoder_hidden_states=encoder_hidden_states,
+                        )
+
+                    elif phase_idx == 3:
+                        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
+
+            if self.lazy_load:
+                next_prefetch = (block_idx + 1) % self.num_blocks
+                self.offload_manager.start_prefetch_block(next_prefetch)
+
             if block_idx == 0:
-                self.offload_manager.init_first_buffer(block_weights.blocks)
-            if block_idx + 1 < self.num_blocks:
-                self.offload_manager.prefetch_weights(block_idx + 1, 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, 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,