refactor(wan2.2): move Wan2.2 save logic; update wan_t2v configs

Made-with: Cursor

Author: Charles2530

configs/quantization/video_gen/wan_t2v/awq_w_a.yaml

@@ -2,12 +2,12 @@ base:
     seed: &seed 42
 model:
     type: WanT2V
-    path: /mnt/lm_data_afs/wangzining/charles/lab/llmc/models/Wan2.1-T2V-14B-Diffusers
+    path: /path/to/wan_t2v
     torch_dtype: auto
 calib:
     name: t2v
     download: False
-    path: ./assets/wan_t2v/calib/
+    path: ../assets/wan_t2v/calib/
     sample_steps: 20
     bs: 1
     target_height: 480
@@ -20,7 +20,7 @@ eval:
     type: video_gen
     name: t2v
     download: False
-    path: ./assets/wan_t2v/calib/
+    path: ../assets/wan_t2v/calib/
     bs: 1
     target_height: 480
     target_width: 832
@@ -31,12 +31,12 @@ quant:
     video_gen:
         method: Awq
         weight:
-            bit: 4
+            bit: 6
             symmetric: True
             granularity: per_channel
             group_size: -1
         act:
-            bit: 4
+            bit: 6
             symmetric: True
             granularity: per_token
         special:
@@ -46,4 +46,4 @@ quant:
             clip_sym: True
 save:
     save_lightx2v: True
-    save_path: ../lightx2v/wan_t2v_awq_w_a/x2v/
+    save_path: /path/to/x2v/

configs/quantization/video_gen/wan_t2v/awq_w_a_s.yaml

@@ -2,12 +2,12 @@ base:
     seed: &seed 42
 model:
     type: WanT2V
-    path: /mnt/lm_data_afs/wangzining/charles/lab/llmc/models/Wan2.2-T2V-14B-Diffusers
+    path: /path/to/wan_t2v
     torch_dtype: auto
 calib:
     name: t2v
     download: False
-    path: ./assets/wan_t2v/calib/
+    path: ../assets/wan_t2v/calib/
     sample_steps: 20
     bs: 1
     target_height: 480
@@ -20,30 +20,26 @@ eval:
     type: video_gen
     name: t2v
     download: False
-    path: ./assets/wan_t2v/calib/
+    path: ../assets/wan_t2v/calib/
     bs: 1
     target_height: 480
     target_width: 832
     num_frames: 81
     guidance_scale: 5.0
-    output_video_path: ./output_videos_awq/
+    output_video_path: ./output_videos_sq/
 quant:
     video_gen:
-        method: Awq
+        method: SmoothQuant
         weight:
-            bit: 4
+            bit: 6
             symmetric: True
             granularity: per_channel
-            group_size: -1
         act:
-            bit: 4
+            bit: 6
             symmetric: True
             granularity: per_token
         special:
-            trans: True
-            trans_version: v2
-            weight_clip: True
-            clip_sym: True
+            alpha: 0.7
 save:
     save_lightx2v: True
-    save_path: ../lightx2v/wan_t2v_awq_w_a/x2v/
+    save_path: /path/to/x2v/

configs/quantization/video_gen/wan_t2v/smoothquant_w_a.yaml

@@ -1076,57 +1076,7 @@ def save_model(self, path):
             logger.info('save model done --')
             self.copy_tokenizer(path)
         elif self.config.model.type in ['Wan2T2V']:
-            if getattr(self.model.Pipeline, '_is_wan_official', False):
-                src = getattr(self.model, 'pipeline_model_path', self.model.model_path)
-                self.model.copy_native_checkpoint(src, path)
-
-                self.model.Pipeline.transformer.save_pretrained(
-                    os.path.join(path, 'high_noise_model')
-                )
-                logger.info('save Wan2.2 high_noise_model done --')
-                if (
-                    hasattr(self.model.Pipeline, 'transformer_2')
-                    and self.model.Pipeline.transformer_2 is not None
-                ):
-                    self.model.Pipeline.transformer_2.save_pretrained(
-                        os.path.join(path, 'low_noise_model')
-                    )
-                    logger.info('save Wan2.2 low_noise_model done --')
-                self.model.validate_native_save_structure(path, source_path=src)
-                return
-
-            # Copy the full original pipeline (VAE, text encoder, tokenizer, scheduler, etc.)
-            # so that non-quantized components are preserved.
-            src = getattr(self.model, 'pipeline_model_path', self.model.model_path)
-            copied_from_source = False
-            if isinstance(src, str) and os.path.isdir(src) and os.path.abspath(src) != os.path.abspath(path):
-                if os.path.exists(path):
-                    shutil.rmtree(path)
-                shutil.copytree(src, path)
-                logger.info(f'Copied original pipeline from {src} to {path}')
-                copied_from_source = True
-            if not copied_from_source:
-                if os.path.exists(path):
-                    shutil.rmtree(path)
-                # Fallback for remote repo-id sources: materialize all non-quantized components first.
-                self.model.Pipeline.save_pretrained(path, safe_serialization=True)
-                logger.info(
-                    'save Wan2.2 full pipeline done via Pipeline.save_pretrained '
-                    f'(source={src}) --'
-                )
-            # Overwrite transformer subfolder with quantized weights.
-            self.model.Pipeline.transformer.save_pretrained(
-                os.path.join(path, 'transformer')
-            )
-            logger.info('save Wan2.2 transformer done --')
-            if (
-                hasattr(self.model.Pipeline, 'transformer_2')
-                and self.model.Pipeline.transformer_2 is not None
-            ):
-                self.model.Pipeline.transformer_2.save_pretrained(
-                    os.path.join(path, 'transformer_2')
-                )
-                logger.info('save Wan2.2 transformer_2 done --')
+            self.model.save_wan2_2_pretrained(path)
         else:
             self.model.get_model().save_pretrained(path)
             logger.info('save model done --')

llmc/compression/quantization/base_blockwise_quantization.py

@@ -677,5 +677,68 @@ def validate_native_save_structure(save_path, source_path=None):
             f'top-level entries={sorted(os.listdir(save_path))}'
         )
 
+    def save_wan2_2_pretrained(self, path):
+        """Wan2.2 专用保存：支持官方 native 与非官方 Pipeline 两种布局。
+
+        该逻辑原本位于 llmc/compression/quantization/base_blockwise_quantization.py 的 Wan2T2V 分支。
+        """
+        if int(os.environ.get('RANK', '0')) != 0:
+            return
+
+        if getattr(self.Pipeline, '_is_wan_official', False):
+            src = getattr(self, 'pipeline_model_path', self.model_path)
+            self.copy_native_checkpoint(src, path)
+
+            self.Pipeline.transformer.save_pretrained(
+                os.path.join(path, 'high_noise_model')
+            )
+            logger.info('save Wan2.2 high_noise_model done --')
+            if (
+                hasattr(self.Pipeline, 'transformer_2')
+                and self.Pipeline.transformer_2 is not None
+            ):
+                self.Pipeline.transformer_2.save_pretrained(
+                    os.path.join(path, 'low_noise_model')
+                )
+                logger.info('save Wan2.2 low_noise_model done --')
+
+            self.validate_native_save_structure(path, source_path=src)
+            return
+
+        # Copy the full original pipeline (VAE, text encoder, tokenizer, scheduler, etc.)
+        # so that non-quantized components are preserved.
+        src = getattr(self, 'pipeline_model_path', self.model_path)
+        copied_from_source = False
+        if isinstance(src, str) and os.path.isdir(src) and os.path.abspath(src) != os.path.abspath(path):
+            if os.path.exists(path):
+                shutil.rmtree(path)
+            shutil.copytree(src, path)
+            logger.info(f'Copied original pipeline from {src} to {path}')
+            copied_from_source = True
+
+        if not copied_from_source:
+            if os.path.exists(path):
+                shutil.rmtree(path)
+            # Fallback for remote repo-id sources: materialize all non-quantized components first.
+            self.Pipeline.save_pretrained(path, safe_serialization=True)
+            logger.info(
+                'save Wan2.2 full pipeline done via Pipeline.save_pretrained '
+                f'(source={src}) --'
+            )
+
+        # Overwrite transformer subfolder with quantized weights.
+        self.Pipeline.transformer.save_pretrained(
+            os.path.join(path, 'transformer')
+        )
+        logger.info('save Wan2.2 transformer done --')
+        if (
+            hasattr(self.Pipeline, 'transformer_2')
+            and self.Pipeline.transformer_2 is not None
+        ):
+            self.Pipeline.transformer_2.save_pretrained(
+                os.path.join(path, 'transformer_2')
+            )
+            logger.info('save Wan2.2 transformer_2 done --')
+
     def skip_layer_name(self):
         pass

llmc/models/wan2_2_t2v.py

@@ -31,13 +31,10 @@ def __init__(self, config, device_map=None, use_cache=False):
 
     def build_model(self):
         vae = AutoencoderKLWan.from_pretrained(
-            self.model_path, subfolder='vae', torch_dtype=torch.float32, use_safetensors=True
+            self.model_path, subfolder='vae', torch_dtype=torch.float32
         )
-        # self.Pipeline = WanPipeline.from_pretrained(
-        #     self.model_path, vae=vae, torch_dtype=torch.bfloat16
-        # )
         self.Pipeline = WanPipeline.from_pretrained(
-            self.model_path, vae=vae, torch_dtype=torch.bfloat16, use_safetensors=True
+            self.model_path, vae=vae, torch_dtype=torch.bfloat16
         )
         self.find_llmc_model()
         self.find_blocks()
@@ -64,17 +61,16 @@ def __init__(self, module):
 
             def forward(self, *args, **kwargs):
                 params = list(self.signature.parameters.keys())
-                capture_kwargs = dict(kwargs)
                 for i, arg in enumerate(args):
                     if i > 0:
-                        capture_kwargs[params[i]] = arg
+                        kwargs[params[i]] = arg
                 first_block_input['data'].append(args[0])
-                first_block_input['kwargs'].append(capture_kwargs)
+                first_block_input['kwargs'].append(kwargs)
                 self.step += 1
                 if self.step == sample_steps:
                     raise ValueError
                 else:
-                    return self.module(*args, **kwargs)
+                    return self.module(*args)
 
         return Catcher
 
@@ -166,4 +162,4 @@ def get_layers_except_blocks(self):
         pass
 
     def skip_layer_name(self):
-        pass
+        pass