support w4 for qw_vl and disk offload for qw_2511

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/models/input_encoders/hf/qwen25/qwen25_vlforconditionalgeneration.py

@@ -71,7 +71,22 @@ 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.embed_tokens": AI_DEVICE,
+                    "model.norm": 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 +114,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 +159,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 +198,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

@@ -10,7 +10,6 @@
 class QwenImageOffloadTransformerInfer(QwenImageTransformerInfer):
     def __init__(self, config):
         super().__init__(config)
-        self.phases_num = 3
         self.num_blocks = config["num_layers"]
         if self.config.get("cpu_offload", False):
             if "offload_ratio" in self.config:
@@ -19,23 +18,28 @@ def __init__(self, config):
                 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
+
+            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):
         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)
+
+            if self.lazy_load:
+                self.offload_manager.swap_cpu_buffers()
+            self.offload_manager.prefetch_weights((block_idx + 1) % self.num_blocks, block_weights.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],

lightx2v/models/networks/qwen_image/model.py

@@ -36,7 +36,10 @@ def __init__(self, config):
             transformer_config = json.load(f)
             self.in_channels = transformer_config["in_channels"]
         self.attention_kwargs = {}
-
+        self.remove_keys = []
+        self.lazy_load = self.config.get("lazy_load", False)
+        if self.lazy_load:
+            self.remove_keys.extend(["blocks."])
         self.dit_quantized = self.config.get("dit_quantized", False)
 
         if self.config["seq_parallel"]:
@@ -75,10 +78,7 @@ def _init_weights(self, weight_dict=None):
                     weight_dict = self._load_ckpt(unified_dtype, sensitive_layer)
                 else:
                     # Load quantized weights
-                    if not self.config.get("lazy_load", False):
-                        weight_dict = self._load_quant_ckpt(unified_dtype, sensitive_layer)
-                    else:
-                        weight_dict = self._load_quant_split_ckpt(unified_dtype, sensitive_layer)
+                    weight_dict = self._load_quant_ckpt(unified_dtype, sensitive_layer)
 
             if self.config.get("device_mesh") is not None and self.config.get("load_from_rank0", False):
                 weight_dict = self._load_weights_from_rank0(weight_dict, is_weight_loader)
@@ -89,7 +89,10 @@ def _init_weights(self, weight_dict=None):
 
         # Initialize weight containers
         self.pre_weight = self.pre_weight_class(self.config)
-        self.transformer_weights = self.transformer_weight_class(self.config)
+        if self.lazy_load:
+            self.transformer_weights = self.transformer_weight_class(self.config, self.lazy_load_path)
+        else:
+            self.transformer_weights = self.transformer_weight_class(self.config)
         self.post_weight = self.post_weight_class(self.config)
         if not self._should_init_empty_model():
             self._apply_weights()
@@ -150,8 +153,18 @@ def _load_ckpt(self, unified_dtype, sensitive_layer):
             safetensors_path = self.model_path
 
         if os.path.isdir(safetensors_path):
-            safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
+            if self.lazy_load:
+                self.lazy_load_path = safetensors_path
+                non_block_file = os.path.join(safetensors_path, "non_block.safetensors")
+                if os.path.exists(non_block_file):
+                    safetensors_files = [non_block_file]
+                else:
+                    raise ValueError(f"Non-block file not found in {safetensors_path}. Please check the model path.")
+            else:
+                safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
         else:
+            if self.lazy_load:
+                self.lazy_load_path = safetensors_path
             safetensors_files = [safetensors_path]
 
         weight_dict = {}
@@ -171,8 +184,18 @@ def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
             safetensors_path = self.model_path
 
         if os.path.isdir(safetensors_path):
-            safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
+            if self.lazy_load:
+                self.lazy_load_path = safetensors_path
+                non_block_file = os.path.join(safetensors_path, "non_block.safetensors")
+                if os.path.exists(non_block_file):
+                    safetensors_files = [non_block_file]
+                else:
+                    raise ValueError(f"Non-block file not found in {safetensors_path}. Please check the model path.")
+            else:
+                safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
         else:
+            if self.lazy_load:
+                self.lazy_load_path = safetensors_path
             safetensors_files = [safetensors_path]
             safetensors_path = os.path.dirname(safetensors_path)
 
@@ -204,28 +227,6 @@ def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
 
         return weight_dict
 
-    def _load_quant_split_ckpt(self, unified_dtype, sensitive_layer):  # Need rewrite
-        lazy_load_model_path = self.dit_quantized_ckpt
-        logger.info(f"Loading splited quant model from {lazy_load_model_path}")
-        pre_post_weight_dict = {}
-
-        safetensor_path = os.path.join(lazy_load_model_path, "non_block.safetensors")
-        with safe_open(safetensor_path, framework="pt", device="cpu") as f:
-            for k in f.keys():
-                if f.get_tensor(k).dtype in [
-                    torch.float16,
-                    torch.bfloat16,
-                    torch.float,
-                ]:
-                    if unified_dtype or all(s not in k for s in sensitive_layer):
-                        pre_post_weight_dict[k] = f.get_tensor(k).to(GET_DTYPE()).to(self.device)
-                    else:
-                        pre_post_weight_dict[k] = f.get_tensor(k).to(GET_SENSITIVE_DTYPE()).to(self.device)
-                else:
-                    pre_post_weight_dict[k] = f.get_tensor(k).to(self.device)
-
-        return pre_post_weight_dict
-
     def _load_weights_from_rank0(self, weight_dict, is_weight_loader):
         logger.info("Loading distributed weights")
         global_src_rank = 0
@@ -291,6 +292,8 @@ def _init_infer(self):
         self.post_infer = self.post_infer_class(self.config)
         if hasattr(self.transformer_infer, "offload_manager"):
             self.transformer_infer.offload_manager.init_cuda_buffer(self.transformer_weights.offload_block_cuda_buffers, self.transformer_weights.offload_phase_cuda_buffers)
+            if self.lazy_load:
+                self.transformer_infer.offload_manager.init_cpu_buffer(self.transformer_weights.offload_block_cpu_buffers, self.transformer_weights.offload_phase_cpu_buffers)
 
     def to_cpu(self):
         self.pre_weight.to_cpu()

lightx2v/models/networks/qwen_image/weights/post_weights.py

@@ -10,19 +10,12 @@ def __init__(self, config):
         super().__init__()
         self.task = config["task"]
         self.config = config
-        self.lazy_load = self.config.get("lazy_load", False)
-        if self.lazy_load:
-            assert NotImplementedError
-        self.lazy_load_file = False
-
         # norm_out
         self.add_module(
             "norm_out_linear",
             MM_WEIGHT_REGISTER["Default"](
                 "norm_out.linear.weight",
                 "norm_out.linear.bias",
-                self.lazy_load,
-                self.lazy_load_file,
             ),
         )
         self.add_module("norm_out", LN_WEIGHT_REGISTER["Default"](eps=1e-6))
@@ -33,8 +26,6 @@ def __init__(self, config):
             MM_WEIGHT_REGISTER["Default"](
                 "proj_out.weight",
                 "proj_out.bias",
-                self.lazy_load,
-                self.lazy_load_file,
             ),
         )