【bugfix】ascendv1_save_module_preprocess 允许保存新prefix-MindStudio-ModelSlim-AtomGit

已合并

【bugfix】ascendv1_save_module_preprocess 允许保存新prefix #106

zhangz200102创建于 2月3日从zhangz200102/msmodelslim： bugfix/save_module_preprocess 合入到 Ascend/msmodelslim： master

【bugfix】ascendv1_save_module_preprocess 允许保存新prefix #106

已合并

zhangz200102创建于 2月3日

从zhangz200102/msmodelslim bugfix/save_module_preprocess合入到Ascend/msmodelslim master

共 4 个文件变更+38-64

Mmsmodelslim/core/quant_service/modelslim_v1/save/ascendv1.py+1-1

	@@ -552,5 +552,5 @@ class AscendV1Saver(AutoSaverProcessor):
552	def _process_module(self, prefix: str, module: nn.Module):	552	def _process_module(self, prefix: str, module: nn.Module):
553	if isinstance(self.adapter, AscendV1SaveInterface):	553	if isinstance(self.adapter, AscendV1SaveInterface):
554	self.processed_modules.add(module)	554	self.processed_modules.add(module)
555	- module = self.adapter.ascendv1_save_module_preprocess(prefix, module, self.model) or module	555	+ prefix, module = self.adapter.ascendv1_save_module_preprocess(prefix, module, self.model)
556	super()._process_module(prefix=prefix, module=module)	556	super()._process_module(prefix=prefix, module=module)

Mmsmodelslim/core/quant_service/modelslim_v1/save/interface.py+10-4

	@@ -20,7 +20,7 @@ See the Mulan PSL v2 for more details.
20	"""	20	"""
21		21
22	from abc import ABC, abstractmethod	22	from abc import ABC, abstractmethod
23	-from typing import Optional	23	+from typing import Tuple
24	import torch.nn as nn	24	import torch.nn as nn
25		25
26		26
	@@ -33,6 +33,12 @@ class AscendV1SaveInterface(ABC):
33	"""	33	"""
34	pass	34	pass
35		35
36	- def ascendv1_save_module_preprocess(self, prefix: str, module: nn.Module, model: nn.Module) -> Optional[nn.Module]:	36	+ def ascendv1_save_module_preprocess(self, prefix: str, module: nn.Module, model: nn.Module) -> Tuple[str, nn.Module]:
37	- pass	37	+ """
38	-	38	+ 在保存模块前，对模块进行预处理，返回新的前缀和模块
		39	+ @param prefix: 模块的前缀路径
		40	+ @param module: 待处理的模块
		41	+ @param model: 模型
		42	+ @return: 返回(prefix, module)
		43	+ """
		44	+ pass

Mmsmodelslim/model/qwen3_next/model_adapter.py+3-3

	@@ -113,10 +113,10 @@ class Qwen3NextModelAdapter(TransformersModel,
113	])	113	])
114	return adapter_config	114	return adapter_config
115		115
116	- def ascendv1_save_module_preprocess(self, prefix: str, module: nn.Module, model: nn.Module) -> Optional[nn.Module]:	116	+ def ascendv1_save_module_preprocess(self, prefix: str, module: nn.Module, model: nn.Module) -> Tuple[str, nn.Module]:
117	if 'input_layernorm' in prefix and module.__class__.__name__ == 'Qwen3RMSNorm':	117	if 'input_layernorm' in prefix and module.__class__.__name__ == 'Qwen3RMSNorm':
118	new_module = Qwen3NextRMSNorm(module.weight.shape[0], module.variance_epsilon)	118	new_module = Qwen3NextRMSNorm(module.weight.shape[0], module.variance_epsilon)
119	new_module.weight.data = module.weight.data - 1	119	new_module.weight.data = module.weight.data - 1
120	model.set_submodule(prefix, new_module)	120	model.set_submodule(prefix, new_module)
121	- return new_module	121	+ return prefix, new_module
122	- return None	122	+ return prefix, module

Mtest/cases/model/qwen3_next/test_model_adapter_qwen3_next.py+24-56

 from pathlib import Path
 from unittest.mock import MagicMock, patch
++import torch
 import torch.nn as nn
 from msmodelslim.core.base.protocol import ProcessRequest
             )
             # 创建模拟的Qwen3RMSNorm模块
++            test_prefix = "model.layers.0.input_layernorm"
++            original_weight_data = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
++            expected_weight_data = original_weight_data - 1  # 期望的weight值：原始值 - 1
++
             mock_module = MagicMock()
             mock_module.__class__.__name__ = 'Qwen3RMSNorm'
             mock_module.weight = MagicMock()
++            mock_module.weight.shape = [5]
             mock_module.variance_epsilon = 1e-6
++            mock_module.weight.data = original_weight_data.clone()
+-            mock_module.weight.data.__sub__ = MagicMock(return_value=MagicMock())
+-            mock_module.weight.data.__radd__ = MagicMock(return_value=MagicMock())
             mock_model = MagicMock()
++
++            # Mock Qwen3NextRMSNorm 的创建，使用真实的nn.Module来存储weight
++            with patch('msmodelslim.model.qwen3_next.model_adapter.Qwen3NextRMSNorm') as mock_qwen3_next_rms_norm:
++                # 创建一个真实的新模块来存储weight
++                class MockNewModule(nn.Module):
++                    def __init__(self):
++                        super().__init__()
++                        self.weight = nn.Parameter(torch.zeros(5))
++
++                mock_new_module = MockNewModule()
++                mock_qwen3_next_rms_norm.return_value = mock_new_module
++
++                new_prefix, new_module = adapter.ascendv1_save_module_preprocess(test_prefix, mock_module, mock_model)
++                # 验证prefix没有变化
++                self.assertEqual(new_prefix, test_prefix)
++                # 验证新模块的weight.data是原始weight.data - 1
++                self.assertTrue(torch.allclose(new_module.weight.data, expected_weight_data))
+-            # 验证model.set_submodule被调用
+-            mock_model.set_submodule.assert_called_once()
+-
+-    @patch.dict('sys.modules', {
+-        'transformers.models.qwen3_next.modeling_qwen3_next': MagicMock(),
+-    })
+-    def test_ascendv1_save_module_preprocess_without_input_layernorm(self):
+-        """测试ascendv1_save_module_preprocess方法当prefix不包含input_layernorm时"""
+-        from msmodelslim.model.qwen3_next.model_adapter import Qwen3NextModelAdapter
+-        with patch('msmodelslim.model.qwen3_next.model_adapter.TransformersModel.__init__', return_value=None):
+-            adapter = Qwen3NextModelAdapter(
+-                model_type=self.model_type,
+-                model_path=self.model_path
+-            )
+-
+-            mock_module = MagicMock()
+-            mock_model = MagicMock()
+-
+-            result = adapter.ascendv1_save_module_preprocess("model.layers.0.some_other_module", mock_module,
+-                                                             mock_model)
+-
+-            # 验证返回值为None
+-            self.assertIsNone(result)
+-            # 验证model.set_submodule没有被调用
+-            mock_model.set_submodule.assert_not_called()
+-
+-    @patch.dict('sys.modules', {
+-        'transformers.models.qwen3_next.modeling_qwen3_next': MagicMock(),
+-    })
+-    def test_ascendv1_save_module_preprocess_with_wrong_module_type(self):
+-        """测试ascendv1_save_module_preprocess方法当module不是Qwen3RMSNorm时"""
+-        from msmodelslim.model.qwen3_next.model_adapter import Qwen3NextModelAdapter
+-        with patch('msmodelslim.model.qwen3_next.model_adapter.TransformersModel.__init__', return_value=None):
+-            adapter = Qwen3NextModelAdapter(
+-                model_type=self.model_type,
+-                model_path=self.model_path
+-            )
+-
+-            mock_module = MagicMock()
+-            mock_module.__class__.__name__ = 'SomeOtherModule'
+-            mock_model = MagicMock()
+-
+-            result = adapter.ascendv1_save_module_preprocess("model.layers.0.input_layernorm", mock_module, mock_model)
+-
+-            # 验证返回值为None
+-            self.assertIsNone(result)
+-            # 验证model.set_submodule没有被调用
+-            mock_model.set_submodule.assert_not_called()