已合并
[新需求]: Layer adapter + Quant method adapter #18
CXY-Katrina创建于 2025年12月31日
[新需求]: Layer adapter + Quant method adapter #18
已合并
共 4 个文件变更+988-0
| @@ -0,0 +1,200 @@ | |||
| 1 | +#!/usr/bin/env python | ||
| 2 | +# coding=utf-8 | ||
| 3 | +# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | ||
| 4 | +# MindIE is licensed under Mulan PSL v2. | ||
| 5 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 6 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 7 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 8 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 9 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 10 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 11 | +# See the Mulan PSL v2 for more details. | ||
| 12 | + | ||
| 13 | +from abc import ABC, abstractmethod | ||
| 14 | +import torch | ||
| 15 | + | ||
| 16 | +from atb_llm.utils.data.quant_method_adapter import ( | ||
| 17 | + QuantizationConfig, | ||
| 18 | +) | ||
| 19 | +from atb_llm.utils.quantize.pack_type import TransposeType | ||
| 20 | +from atb_llm.utils.quantize.quant_type import LinearTypeV2, QuantType | ||
| 21 | +from mindie_llm.runtime.layers.linear.linear import ( | ||
| 22 | + RowParallelLinear as RowParallelLinearAdaptee, | ||
| 23 | + ColumnParallelLinear as ColumnParallelLinearAdaptee, | ||
| 24 | + MergedColumnParallelLinear as MergedColumnParallelLinearAdaptee, | ||
| 25 | + QKVParallelLinear as QKVParallelLinearAdaptee, | ||
| 26 | +) | ||
| 27 | +from mindie_llm.runtime.layers.embedding.embedding import ( | ||
| 28 | + VocabParallelEmbedding as VocabParallelEmbeddingAdaptee, | ||
| 29 | + ParallelLMHead as ParallelLMHeadAdaptee, | ||
| 30 | +) | ||
| 31 | +from mindie_llm.runtime.layers.normalization import RMSNorm as RMSNormAdaptee | ||
| 32 | +from mindie_llm.runtime.layers.quantization.unquantized import ( | ||
| 33 | + UnquantizedEmbeddingMethod, | ||
| 34 | + UnquantizedNormMethod, | ||
| 35 | + UnquantizedLinearMethod, | ||
| 36 | +) | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +class LayerSupportAtbGraph(ABC): | ||
| 40 | + _PLACEHOLDER = torch.tensor([1], dtype=torch.get_default_dtype(), device='npu') | ||
| 41 | + | ||
| 42 | + @abstractmethod | ||
| 43 | + def get_weights_for_atb_graph(self, padding: bool = True) -> list[torch.Tensor]: | ||
| 44 | + pass | ||
| 45 | + | ||
| 46 | + | ||
| 47 | +class LinearLayerSupportAtbGraph(LayerSupportAtbGraph): | ||
| 48 | + @abstractmethod | ||
| 49 | + def get_weights_for_atb_graph( | ||
| 50 | + self, padding: bool = True, | ||
| 51 | + is_swiglu_quant_enabled: bool = False, | ||
| 52 | + quant_type: QuantType = None | ||
| 53 | + ) -> list[torch.Tensor]: | ||
| 54 | + pass | ||
| 55 | + | ||
| 56 | + @abstractmethod | ||
| 57 | + def get_linear_descs(self) -> list[LinearTypeV2]: | ||
| 58 | + pass | ||
| 59 | + | ||
| 60 | + @abstractmethod | ||
| 61 | + def get_weight_transpose_type(self) -> list[TransposeType]: | ||
| 62 | + pass | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +class RowParallelLinear(RowParallelLinearAdaptee, LayerSupportAtbGraph): | ||
| 66 | + def __init__(self, *args, **kwargs): | ||
| 67 | + kwargs["quant_config"] = QuantizationConfig( | ||
| 68 | + kwargs.get("quant_config"), UnquantizedLinearMethod) | ||
| 69 | + super().__init__(*args, **kwargs) | ||
| 70 | + | ||
| 71 | + def get_weights_for_atb_graph( | ||
| 72 | + self, padding: bool = True, | ||
| 73 | + is_swiglu_quant_enabled: bool = False, | ||
| 74 | + quant_type: QuantType = None | ||
| 75 | + ) -> list[torch.Tensor]: | ||
| 76 | + return self.quant_method.get_weights_for_atb_graph( | ||
| 77 | + self, padding=padding, is_swiglu_quant_enabled=is_swiglu_quant_enabled, | ||
| 78 | + quant_type=quant_type) | ||
| 79 | + | ||
| 80 | + def get_linear_descs(self) -> list[LinearTypeV2]: | ||
| 81 | + return [self.quant_method.get_linear_descs(self)] | ||
| 82 | + | ||
| 83 | + def get_weight_transpose_type(self) -> list[TransposeType]: | ||
| 84 | + return [self.quant_method.get_weight_transpose_type(self)] | ||
| 85 | + | ||
| 86 | + | ||
| 87 | +class ColumnParallelLinear(ColumnParallelLinearAdaptee, LayerSupportAtbGraph): | ||
| 88 | + def __init__(self, *args, **kwargs): | ||
| 89 | + kwargs["quant_config"] = QuantizationConfig( | ||
| 90 | + kwargs.get("quant_config"), UnquantizedLinearMethod) | ||
| 91 | + super().__init__(*args, **kwargs) | ||
| 92 | + | ||
| 93 | + def get_weights_for_atb_graph( | ||
| 94 | + self, padding: bool = True, | ||
| 95 | + is_swiglu_quant_enabled: bool = False, | ||
| 96 | + quant_type: QuantType = None | ||
| 97 | + ) -> list[torch.Tensor]: | ||
| 98 | + if is_swiglu_quant_enabled: | ||
| 99 | + raise ValueError("Cannot set `is_swiglu_quant_enabled` to True in `ColumnParallelLinear`.") | ||
| 100 | + return self.quant_method.get_weights_for_atb_graph(self, padding=padding, quant_type=quant_type) | ||
| 101 | + | ||
| 102 | + def get_linear_descs(self) -> list[LinearTypeV2]: | ||
| 103 | + return [self.quant_method.get_linear_descs(self)] | ||
| 104 | + | ||
| 105 | + def get_weight_transpose_type(self) -> list[TransposeType]: | ||
| 106 | + return [self.quant_method.get_weight_transpose_type(self)] | ||
| 107 | + | ||
| 108 | + | ||
| 109 | +class MergedColumnParallelLinear(MergedColumnParallelLinearAdaptee, ColumnParallelLinear): | ||
| 110 | + def __init__(self, *args, **kwargs): | ||
| 111 | + super().__init__(*args, **kwargs) | ||
| 112 | + | ||
| 113 | + def get_weights_for_atb_graph( | ||
| 114 | + self, padding: bool = True, | ||
| 115 | + is_swiglu_quant_enabled: bool = False, | ||
| 116 | + quant_type: QuantType = None | ||
| 117 | + ) -> list[torch.Tensor]: | ||
| 118 | + if is_swiglu_quant_enabled: | ||
| 119 | + raise ValueError("Cannot set `is_swiglu_quant_enabled` to True in `MergedColumnParallelLinear`.") | ||
| 120 | + weights = super().get_weights_for_atb_graph(padding=padding, quant_type=quant_type) | ||
| 121 | + weights.extend([self._PLACEHOLDER] * 6) | ||
| 122 | + return weights | ||
| 123 | + | ||
| 124 | + def get_linear_descs(self) -> list[LinearTypeV2]: | ||
| 125 | + linear_descs = super().get_linear_descs() | ||
| 126 | + linear_descs.append(LinearTypeV2.INVALID) | ||
| 127 | + return linear_descs | ||
| 128 | + | ||
| 129 | + def get_weight_transpose_type(self) -> list[TransposeType]: | ||
| 130 | + transpose_type = super().get_weight_transpose_type() | ||
| 131 | + transpose_type.append(TransposeType.INVALID) | ||
| 132 | + return transpose_type | ||
| 133 | + | ||
| 134 | + | ||
| 135 | +class QKVParallelLinear(QKVParallelLinearAdaptee, ColumnParallelLinear): | ||
| 136 | + def __init__(self, *args, **kwargs): | ||
| 137 | + super().__init__(*args, **kwargs) | ||
| 138 | + | ||
| 139 | + def get_weights_for_atb_graph( | ||
| 140 | + self, padding: bool = True, | ||
| 141 | + is_swiglu_quant_enabled: bool = False, | ||
| 142 | + quant_type: QuantType = None | ||
| 143 | + ) -> list[torch.Tensor]: | ||
| 144 | + if is_swiglu_quant_enabled: | ||
| 145 | + raise ValueError("Cannot set `is_swiglu_quant_enabled` to True in `QKVParallelLinear`.") | ||
| 146 | + weights = self.quant_method.get_weights_for_atb_graph(self, padding=padding, quant_type=quant_type) | ||
| 147 | + weights.extend([self._PLACEHOLDER] * 12) | ||
| 148 | + return weights | ||
| 149 | + | ||
| 150 | + def get_linear_descs(self) -> list[LinearTypeV2]: | ||
| 151 | + linear_descs = super().get_linear_descs() | ||
| 152 | + linear_descs.extend([LinearTypeV2.INVALID] * 2) | ||
| 153 | + return linear_descs | ||
| 154 | + | ||
| 155 | + def get_weight_transpose_type(self) -> list[TransposeType]: | ||
| 156 | + transpose_type = super().get_weight_transpose_type() | ||
| 157 | + transpose_type.extend([TransposeType.INVALID] * 2) | ||
| 158 | + return transpose_type | ||
| 159 | + | ||
| 160 | + | ||
| 161 | +class VocabParallelEmbedding(VocabParallelEmbeddingAdaptee, LayerSupportAtbGraph): | ||
| 162 | + def __init__(self, *args, **kwargs): | ||
| 163 | + kwargs["quant_config"] = QuantizationConfig( | ||
| 164 | + kwargs.get("quant_config"), UnquantizedEmbeddingMethod) | ||
| 165 | + super().__init__(*args, **kwargs) | ||
| 166 | + | ||
| 167 | + def get_weights_for_atb_graph(self, padding: bool = True): | ||
| 168 | + weights = self.quant_method.get_weights_for_atb_graph(self, padding=padding) | ||
| 169 | + return weights | ||
| 170 | + | ||
| 171 | + | ||
| 172 | +class ParallelLMHead(ParallelLMHeadAdaptee, LayerSupportAtbGraph): | ||
| 173 | + def __init__(self, *args, **kwargs): | ||
| 174 | + kwargs["quant_config"] = QuantizationConfig( | ||
| 175 | + kwargs.get("quant_config"), UnquantizedLinearMethod) | ||
| 176 | + super().__init__(*args, **kwargs) | ||
| 177 | + | ||
| 178 | + def get_weights_for_atb_graph( | ||
| 179 | + self, padding: bool = True, | ||
| 180 | + is_swiglu_quant_enabled: bool = False, | ||
| 181 | + quant_type: QuantType = None | ||
| 182 | + ) -> list[torch.Tensor]: | ||
| 183 | + weights = self.quant_method.get_weights_for_atb_graph(self, padding=padding) | ||
| 184 | + return weights | ||
| 185 | + | ||
| 186 | + def get_linear_descs(self) -> list[LinearTypeV2]: | ||
| 187 | + return [self.quant_method.get_linear_descs(self)] | ||
| 188 | + | ||
| 189 | + def get_weight_transpose_type(self) -> list[TransposeType]: | ||
| 190 | + return [self.quant_method.get_weight_transpose_type(self)] | ||
| 191 | + | ||
| 192 | + | ||
| 193 | +class RMSNorm(RMSNormAdaptee, LayerSupportAtbGraph): | ||
| 194 | + def __init__(self, *args, **kwargs): | ||
| 195 | + kwargs["quant_config"] = QuantizationConfig( | ||
| 196 | + kwargs.get("quant_config"), UnquantizedNormMethod) | ||
| 197 | + super().__init__(*args, **kwargs) | ||
| 198 | + | ||
| 199 | + def get_weights_for_atb_graph(self, padding: bool = True): | ||
| 200 | + return self.quant_method.get_weights_for_atb_graph(self, padding=padding) | ||
| @@ -0,0 +1,329 @@ | |||
| 1 | +#!/usr/bin/env python | ||
| 2 | +# coding=utf-8 | ||
| 3 | +# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | ||
| 4 | +# MindIE is licensed under Mulan PSL v2. | ||
| 5 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 6 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 7 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 8 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 9 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 10 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 11 | +# See the Mulan PSL v2 for more details. | ||
| 12 | + | ||
| 13 | +from abc import ABC, abstractmethod | ||
| 14 | + | ||
| 15 | +import torch | ||
| 16 | +import torch_npu | ||
| 17 | +from torch import nn | ||
| 18 | + | ||
| 19 | +from atb_llm.utils.quantize.pack_type import TransposeType | ||
| 20 | +from atb_llm.utils.quantize.quant_type import LinearTypeV2, QuantType | ||
| 21 | +from atb_llm.utils.env import ENV | ||
| 22 | +from atb_llm.utils.initial import NPUSocInfo | ||
| 23 | +from mindie_llm.runtime.layers.linear.linear import LinearBase | ||
| 24 | +from mindie_llm.runtime.layers.quantization.quantization_method_base import QuantizationMethodBase | ||
| 25 | +from mindie_llm.runtime.layers.quantization.ms_model_slim.quant_type import InferenceMode | ||
| 26 | +from mindie_llm.runtime.layers.quantization.unquantized import ( | ||
| 27 | + UnquantizedLinearMethod as UnquantizedLinearMethodAdaptee, | ||
| 28 | + UnquantizedEmbeddingMethod as UnquantizedEmbeddingMethodAdaptee, | ||
| 29 | + UnquantizedNormMethod as UnquantizedNormMethodAdaptee, | ||
| 30 | +) | ||
| 31 | +from mindie_llm.runtime.layers.quantization.ms_model_slim.quantization_config import \ | ||
| 32 | + QuantizationConfig as QuantizationConfigAdaptee | ||
| 33 | +from mindie_llm.runtime.layers.quantization.ms_model_slim.w8a8 import ( | ||
| 34 | + W8A8PerTensorLinearMethod as W8A8PerTensorLinearMethodAdaptee, | ||
| 35 | + W8A8PerTokenLinearMethod as W8A8PerTokenLinearMethodAdaptee, | ||
| 36 | + W8A8MixLinearMethod as W8A8MixLinearMethodAdaptee, | ||
| 37 | +) | ||
| 38 | +from mindie_llm.runtime.layers.quantization.ms_model_slim.anti_outlier import \ | ||
| 39 | + AntiOutlierNormMethod as AntiOutlierNormMethodAdaptee | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +class QuantizationConfig: | ||
| 43 | + def __init__(self, adaptee: QuantizationConfigAdaptee, default_quant_method_cls): | ||
| 44 | + self.adaptee = adaptee | ||
| 45 | + self.default_quant_method_cls = default_quant_method_cls | ||
| 46 | + | ||
| 47 | + @staticmethod | ||
| 48 | + def _get_method_adpater(quant_method: QuantizationMethodBase): | ||
| 49 | + quant_method_cls_adapter_map = { | ||
| 50 | + UnquantizedLinearMethodAdaptee: UnquantizedLinearMethod, | ||
| 51 | + UnquantizedEmbeddingMethodAdaptee: UnquantizedEmbeddingMethod, | ||
| 52 | + UnquantizedNormMethodAdaptee: UnquantizedNormMethod, | ||
| 53 | + W8A8PerTensorLinearMethodAdaptee: W8A8PerTensorLinearMethod, | ||
| 54 | + W8A8PerTokenLinearMethodAdaptee: W8A8PerTokenLinearMethod, | ||
| 55 | + W8A8MixLinearMethodAdaptee: W8A8MixLinearMethod, | ||
| 56 | + AntiOutlierNormMethodAdaptee: AntiOutlierNormMethod, | ||
| 57 | + } | ||
| 58 | + for quant_method_cls, adapter_cls in quant_method_cls_adapter_map.items(): | ||
| 59 | + if isinstance(quant_method, quant_method_cls): | ||
| 60 | + return adapter_cls(quant_method) | ||
| 61 | + raise NotImplementedError(f"Cannot found the adapter class for `{quant_method}`") | ||
| 62 | + | ||
| 63 | + def get_quant_type_by_weight_name(self, *args, **kwargs) -> QuantizationMethodBase | None: | ||
| 64 | + if self.adaptee is None: | ||
| 65 | + raise ValueError(f"NoneType adaptee doesn't support `get_quant_type_by_weight_name`.") | ||
| 66 | + return self.adaptee.get_quant_type_by_weight_name(*args, **kwargs) | ||
| 67 | + | ||
| 68 | + def get_quant_method( | ||
| 69 | + self, layer: torch.nn.Module, prefix: str | list[str] = "" | ||
| 70 | + ) -> QuantizationMethodBase | None: | ||
| 71 | + if self.adaptee is not None: | ||
| 72 | + quant_method = self.adaptee.get_quant_method(layer, prefix) | ||
| 73 | + else: | ||
| 74 | + quant_method = self.default_quant_method_cls() | ||
| 75 | + return self._get_method_adpater(quant_method) | ||
| 76 | + | ||
| 77 | + | ||
| 78 | +class MethodSupportAtbGraph(ABC): | ||
| 79 | + _PLACEHOLDER = torch.tensor([1], dtype=torch.get_default_dtype(), device='npu') | ||
| 80 | + | ||
| 81 | + def __init__(self, adaptee: QuantizationMethodBase): | ||
| 82 | + self.adaptee = adaptee | ||
| 83 | + | ||
| 84 | + def create_weights(self, *args, **kwargs): | ||
| 85 | + self.adaptee.create_weights(*args, **kwargs) | ||
| 86 | + | ||
| 87 | + def apply(self, *args, **kwargs): | ||
| 88 | + self.adaptee.apply(*args, **kwargs) | ||
| 89 | + | ||
| 90 | + def process_weights_after_loading(self, *args, **kwargs): | ||
| 91 | + self.adaptee.process_weights_after_loading(*args, **kwargs) | ||
| 92 | + | ||
| 93 | + @abstractmethod | ||
| 94 | + def get_weights_for_atb_graph(self, layer: torch.nn.Module, padding: bool = True) -> list[torch.Tensor]: | ||
| 95 | + pass | ||
| 96 | + | ||
| 97 | + | ||
| 98 | +class LinearMethodSupportAtbGraph(MethodSupportAtbGraph): | ||
| 99 | + _soc_info = None | ||
| 100 | + | ||
| 101 | + def __init__(self, adaptee: QuantizationMethodBase): | ||
| 102 | + super().__init__(adaptee) | ||
| 103 | + | ||
| 104 | + @classmethod | ||
| 105 | + def set_soc_info(cls, soc_info: NPUSocInfo): | ||
| 106 | + cls._soc_info = soc_info | ||
| 107 | + | ||
| 108 | + @abstractmethod | ||
| 109 | + def get_weights_for_atb_graph( | ||
| 110 | + self, layer: torch.nn.Module, padding: bool = True, | ||
| 111 | + is_swiglu_quant_enabled: bool = False, | ||
| 112 | + quant_type: QuantType = None | ||
| 113 | + ) -> list[torch.Tensor]: | ||
| 114 | + pass | ||
| 115 | + | ||
| 116 | + @abstractmethod | ||
| 117 | + def get_linear_descs(self, layer: LinearBase) -> list[LinearTypeV2]: | ||
| 118 | + pass | ||
| 119 | + | ||
| 120 | + @abstractmethod | ||
| 121 | + def get_weight_transpose_type(self, layer: LinearBase) -> list[TransposeType]: | ||
| 122 | + pass | ||
| 123 | + | ||
| 124 | + def _check_transpose(self, weight_shape): | ||
| 125 | + if self._soc_info is None: | ||
| 126 | + raise ValueError("``NPUSocInfo` is not set in `LinearMethodSupportAtbGraph`.") | ||
| 127 | + | ||
| 128 | + if self._soc_info.need_nz or not ENV.auto_transpose_enable: | ||
| 129 | + return TransposeType.TRANSPOSE | ||
| 130 | + | ||
| 131 | + if self._soc_info.matmul_nd_nz: | ||
| 132 | + # transpose weights to [k, n] when using nz format | ||
| 133 | + return TransposeType.NOT_TRANSPOSE | ||
| 134 | + | ||
| 135 | + is_k_divisible = weight_shape[-1] % 256 == 0 | ||
| 136 | + is_n_divisible = weight_shape[-2] % 256 == 0 | ||
| 137 | + if not is_k_divisible and is_n_divisible and len(weight_shape) != 3: | ||
| 138 | + return TransposeType.NOT_TRANSPOSE | ||
| 139 | + return TransposeType.TRANSPOSE | ||
| 140 | + | ||
| 141 | + | ||
| 142 | +class UnquantizedLinearMethod(LinearMethodSupportAtbGraph): | ||
| 143 | + def __init__(self, *args, **kwargs): | ||
| 144 | + super().__init__(*args, **kwargs) | ||
| 145 | + self._weight_transpose_type = TransposeType.TRANSPOSE | ||
| 146 | + | ||
| 147 | + def get_weights_for_atb_graph( | ||
| 148 | + self, layer: torch.nn.Module, padding: bool = True, | ||
| 149 | + is_swiglu_quant_enabled: bool = False, | ||
| 150 | + quant_type: QuantType = None | ||
| 151 | + ) -> list[torch.Tensor]: | ||
| 152 | + weight_list = [layer.weight.data] | ||
| 153 | + if layer.bias is not None: | ||
| 154 | + weight_list.append(layer.bias.data) | ||
| 155 | + if not padding: | ||
| 156 | + return weight_list | ||
| 157 | + if layer.bias is None: | ||
| 158 | + weight_list.append(self._PLACEHOLDER) | ||
| 159 | + weight_list.extend([self._PLACEHOLDER] * 4) | ||
| 160 | + return weight_list | ||
| 161 | + | ||
| 162 | + def get_linear_descs(self, layer: LinearBase) -> list[LinearTypeV2]: | ||
| 163 | + if layer.weight.data.dtype == torch.bfloat16: | ||
| 164 | + return LinearTypeV2.BFLOAT16 | ||
| 165 | + else: | ||
| 166 | + return LinearTypeV2.FLOAT16 | ||
| 167 | + | ||
| 168 | + def get_weight_transpose_type(self, layer: LinearBase) -> list[TransposeType]: | ||
| 169 | + return self._weight_transpose_type | ||
| 170 | + | ||
| 171 | + def process_weights_after_loading(self, layer: nn.Module) -> None: | ||
| 172 | + weight_shape = layer.weight.data.shape | ||
| 173 | + self._weight_transpose_type = self._check_transpose(weight_shape) | ||
| 174 | + if self._weight_transpose_type == TransposeType.NOT_TRANSPOSE: | ||
| 175 | + layer.weight.data = layer.weight.data.transpose(-1, -2).contiguous() | ||
| 176 | + if self._soc_info.need_nz or self._soc_info.matmul_nd_nz: | ||
| 177 | + layer.weight.data = torch_npu.npu_format_cast_(layer.weight.data, 29) | ||
| 178 | + | ||
| 179 | + | ||
| 180 | +class W8A8PerTensorLinearMethod(LinearMethodSupportAtbGraph): | ||
| 181 | + def __init__(self, *args, **kwargs): | ||
| 182 | + super().__init__(*args, **kwargs) | ||
| 183 | + | ||
| 184 | + def get_weights_for_atb_graph( | ||
| 185 | + self, layer: torch.nn.Module, padding: bool = True, | ||
| 186 | + is_swiglu_quant_enabled: bool = False, | ||
| 187 | + quant_type: QuantType = None | ||
| 188 | + ) -> list[torch.Tensor]: | ||
| 189 | + if is_swiglu_quant_enabled: | ||
| 190 | + layer.input_offset.data = layer.input_offset.data.to(torch.float32) | ||
| 191 | + layer.input_scale.data = layer.input_scale.data.to(torch.float32) | ||
| 192 | + weight_list = [ | ||
| 193 | + layer.weight.data, layer.quant_bias.data, layer.deq_scale.data, | ||
| 194 | + layer.input_offset.data, layer.input_scale.data | ||
| 195 | + ] | ||
| 196 | + if not padding: | ||
| 197 | + return weight_list | ||
| 198 | + weight_list.append(self._PLACEHOLDER) | ||
| 199 | + return weight_list | ||
| 200 | + | ||
| 201 | + def get_linear_descs(self, layer: LinearBase) -> list[LinearTypeV2]: | ||
| 202 | + return LinearTypeV2.W8A8 | ||
| 203 | + | ||
| 204 | + def get_weight_transpose_type(self, layer: LinearBase) -> list[TransposeType]: | ||
| 205 | + return self._weight_transpose_type | ||
| 206 | + | ||
| 207 | + def process_weights_after_loading(self, layer: nn.Module) -> None: | ||
| 208 | + weight_shape = layer.weight.data.shape | ||
| 209 | + self._weight_transpose_type = self._check_transpose(weight_shape) | ||
| 210 | + if self._weight_transpose_type == TransposeType.NOT_TRANSPOSE: | ||
| 211 | + layer.weight.data = layer.weight.data.transpose(-1, -2).contiguous() | ||
| 212 | + if self._soc_info.need_nz or self._soc_info.matmul_nd_nz: | ||
| 213 | + layer.weight.data = torch_npu.npu_format_cast_(layer.weight.data, 29) | ||
| 214 | + | ||
| 215 | + | ||
| 216 | +class W8A8PerTokenLinearMethod(LinearMethodSupportAtbGraph): | ||
| 217 | + def __init__(self, *args, **kwargs): | ||
| 218 | + super().__init__(*args, **kwargs) | ||
| 219 | + | ||
| 220 | + def get_weights_for_atb_graph( | ||
| 221 | + self, layer: torch.nn.Module, padding: bool = True, | ||
| 222 | + is_swiglu_quant_enabled: bool = False, | ||
| 223 | + quant_type: QuantType = None | ||
| 224 | + ) -> list[torch.Tensor]: | ||
| 225 | + weight_list = [layer.weight.data] | ||
| 226 | + if layer.bias is not None: | ||
| 227 | + weight_list.append(layer.bias.data) | ||
| 228 | + else: | ||
| 229 | + weight_list.append(self._PLACEHOLDER) | ||
| 230 | + weight_list.append(self._PLACEHOLDER) | ||
| 231 | + weight_list.extend([layer.weight_offset.data, layer.weight_scale.data]) | ||
| 232 | + if not padding: | ||
| 233 | + return weight_list | ||
| 234 | + weight_list.append(self._PLACEHOLDER) | ||
| 235 | + return weight_list | ||
| 236 | + | ||
| 237 | + def get_linear_descs(self, layer: LinearBase) -> list[LinearTypeV2]: | ||
| 238 | + return LinearTypeV2.W8A8_DYNAMIC | ||
| 239 | + | ||
| 240 | + def get_weight_transpose_type(self, layer: LinearBase) -> list[TransposeType]: | ||
| 241 | + return self._weight_transpose_type | ||
| 242 | + | ||
| 243 | + def process_weights_after_loading(self, layer: nn.Module) -> None: | ||
| 244 | + weight_shape = layer.weight.data.shape | ||
| 245 | + self._weight_transpose_type = self._check_transpose(weight_shape) | ||
| 246 | + if self._weight_transpose_type == TransposeType.NOT_TRANSPOSE: | ||
| 247 | + layer.weight.data = layer.weight.data.transpose(-1, -2).contiguous() | ||
| 248 | + layer.weight_scale.data = layer.weight_scale.data.flatten() | ||
| 249 | + layer.weight_offset.data = -layer.weight_offset.data.flatten() | ||
| 250 | + if self._soc_info.need_nz or self._soc_info.matmul_nd_nz: | ||
| 251 | + layer.weight.data = torch_npu.npu_format_cast_(layer.weight.data, 29) | ||
| 252 | + | ||
| 253 | + | ||
| 254 | +class W8A8MixLinearMethod(LinearMethodSupportAtbGraph): | ||
| 255 | + def __init__(self, *args, **kwargs): | ||
| 256 | + super().__init__(*args, **kwargs) | ||
| 257 | + prefill_quant_method = self.adaptee.quant_method.get(InferenceMode.PREFILL) | ||
| 258 | + if prefill_quant_method is not None: | ||
| 259 | + self.adaptee.quant_method[InferenceMode.PREFILL] = \ | ||
| 260 | + W8A8PerTokenLinearMethod(self.adaptee.quant_method[InferenceMode.PREFILL]) | ||
| 261 | + decode_quant_method = self.adaptee.quant_method.get(InferenceMode.DECODE) | ||
| 262 | + if decode_quant_method is not None: | ||
| 263 | + self.adaptee.quant_method[InferenceMode.DECODE] = \ | ||
| 264 | + W8A8PerTensorLinearMethod(self.adaptee.quant_method[InferenceMode.DECODE]) | ||
| 265 | + | ||
| 266 | + def get_weights_for_atb_graph( | ||
| 267 | + self, layer: torch.nn.Module, padding: bool = True, | ||
| 268 | + is_swiglu_quant_enabled: bool = False, | ||
| 269 | + quant_type: QuantType = None | ||
| 270 | + ) -> list[torch.Tensor]: | ||
| 271 | + if quant_type == QuantType.W8A8_DYNAMIC: | ||
| 272 | + quant_method = self.adaptee.quant_method[InferenceMode.PREFILL] | ||
| 273 | + else: | ||
| 274 | + quant_method = self.adaptee.quant_method[InferenceMode.DECODE] | ||
| 275 | + return quant_method.get_weights_for_atb_graph( | ||
| 276 | + layer=layer, padding=padding, | ||
| 277 | + is_swiglu_quant_enabled=is_swiglu_quant_enabled | ||
| 278 | + ) | ||
| 279 | + | ||
| 280 | + def get_linear_descs(self, layer: LinearBase) -> list[LinearTypeV2]: | ||
| 281 | + return LinearTypeV2.W8A8_PDMIX | ||
| 282 | + | ||
| 283 | + def get_weight_transpose_type(self, layer: LinearBase) -> list[TransposeType]: | ||
| 284 | + return self._weight_transpose_type | ||
| 285 | + | ||
| 286 | + def process_weights_after_loading(self, layer: nn.Module) -> None: | ||
| 287 | + weight_shape = layer.weight.data.shape | ||
| 288 | + self._weight_transpose_type = self._check_transpose(weight_shape) | ||
| 289 | + if self._weight_transpose_type == TransposeType.NOT_TRANSPOSE: | ||
| 290 | + layer.weight.data = layer.weight.data.transpose(-1, -2).contiguous() | ||
| 291 | + layer.weight_scale.data = layer.weight_scale.data.flatten() | ||
| 292 | + layer.weight_offset.data = -layer.weight_offset.data.flatten() | ||
| 293 | + if self._soc_info.need_nz or self._soc_info.matmul_nd_nz: | ||
| 294 | + layer.weight.data = torch_npu.npu_format_cast_(layer.weight.data, 29) | ||
| 295 | + | ||
| 296 | + | ||
| 297 | +class UnquantizedEmbeddingMethod(MethodSupportAtbGraph): | ||
| 298 | + def __init__(self, *args, **kwargs): | ||
| 299 | + super().__init__(*args, **kwargs) | ||
| 300 | + | ||
| 301 | + def get_weights_for_atb_graph(self, layer: torch.nn.Module, padding: bool = True) -> list[torch.Tensor]: | ||
| 302 | + if not padding: | ||
| 303 | + raise NotImplementedError( | ||
| 304 | + f"The `padding` parameter for `get_weights_for_atb_graph` " | ||
| 305 | + f"in `UnquantizedEmbeddingMethod` cannot be set to False.") | ||
| 306 | + return [layer.weight.data] | ||
| 307 | + | ||
| 308 | + | ||
| 309 | +class UnquantizedNormMethod(MethodSupportAtbGraph): | ||
| 310 | + def __init__(self, *args, **kwargs): | ||
| 311 | + super().__init__(*args, **kwargs) | ||
| 312 | + | ||
| 313 | + def get_weights_for_atb_graph(self, layer: torch.nn.Module, padding: bool = True) -> list[torch.Tensor]: | ||
| 314 | + if not padding: | ||
| 315 | + return [layer.weight.data] | ||
| 316 | + return [layer.weight.data, self._PLACEHOLDER, self._PLACEHOLDER, self._PLACEHOLDER] | ||
| 317 | + | ||
| 318 | + | ||
| 319 | +class AntiOutlierNormMethod(MethodSupportAtbGraph): | ||
| 320 | + def __init__(self, *args, **kwargs): | ||
| 321 | + super().__init__(*args, **kwargs) | ||
| 322 | + | ||
| 323 | + def get_weights_for_atb_graph(self, layer: torch.nn.Module, padding: bool = True) -> list[torch.Tensor]: | ||
| 324 | + if not padding: | ||
| 325 | + weights = [layer.weight.data] | ||
| 326 | + if not torch.all(layer.bias.data == 0): | ||
| 327 | + weights.append(layer.bias.data) | ||
| 328 | + return weights | ||
| 329 | + return [layer.weight.data, layer.bias.data, self._PLACEHOLDER, self._PLACEHOLDER] | ||
| @@ -0,0 +1,155 @@ | |||
| 1 | +#!/usr/bin/env python | ||
| 2 | +# coding=utf-8 | ||
| 3 | +# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | ||
| 4 | +# MindIE is licensed under Mulan PSL v2. | ||
| 5 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 6 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 7 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 8 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 9 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 10 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 11 | +# See the Mulan PSL v2 for more details. | ||
| 12 | + | ||
| 13 | +import unittest | ||
| 14 | +from unittest.mock import MagicMock, patch | ||
| 15 | + | ||
| 16 | +import torch | ||
| 17 | + | ||
| 18 | +from atb_llm.utils.data.layer_adapter import ( | ||
| 19 | + RowParallelLinear, | ||
| 20 | + ColumnParallelLinear, | ||
| 21 | + MergedColumnParallelLinear, | ||
| 22 | + QKVParallelLinear, | ||
| 23 | + VocabParallelEmbedding, | ||
| 24 | + ParallelLMHead, | ||
| 25 | + RMSNorm, | ||
| 26 | +) | ||
| 27 | +from atb_llm.utils.quantize.quant_type import LinearTypeV2 | ||
| 28 | +from atb_llm.utils.quantize.pack_type import TransposeType | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +def make_fake_layer(has_bias=True, weight_shape=(64, 64), dtype=torch.float16): | ||
| 32 | + layer = MagicMock() | ||
| 33 | + layer.weight.data = torch.randn(weight_shape, dtype=dtype, device='cpu') | ||
| 34 | + if has_bias: | ||
| 35 | + layer.bias.data = torch.randn(weight_shape[0], dtype=dtype, device='cpu') | ||
| 36 | + else: | ||
| 37 | + layer.bias = None | ||
| 38 | + layer.quant_method = MagicMock() | ||
| 39 | + layer.quant_method.get_weights_for_atb_graph.return_value = [torch.tensor([1.0], device='cpu')] | ||
| 40 | + desc = LinearTypeV2.BFLOAT16 if dtype == torch.bfloat16 else LinearTypeV2.FLOAT16 | ||
| 41 | + layer.quant_method.get_linear_descs.return_value = desc | ||
| 42 | + layer.quant_method.get_weight_transpose_type.return_value = TransposeType.NOT_TRANSPOSE | ||
| 43 | + layer._PLACEHOLDER = torch.tensor([999], dtype=torch.float32, device='cpu') | ||
| 44 | + return layer | ||
| 45 | + | ||
| 46 | + | ||
| 47 | +class TestLayerAdapter(unittest.TestCase): | ||
| 48 | + | ||
| 49 | + def test_row_parallel_linear(self): | ||
| 50 | + layer = make_fake_layer() | ||
| 51 | + weights = RowParallelLinear.get_weights_for_atb_graph(layer, padding=True, is_swiglu_quant_enabled=True) | ||
| 52 | + self.assertEqual(len(weights), 1) | ||
| 53 | + descs = RowParallelLinear.get_linear_descs(layer) | ||
| 54 | + trans = RowParallelLinear.get_weight_transpose_type(layer) | ||
| 55 | + self.assertEqual(descs, [LinearTypeV2.FLOAT16]) | ||
| 56 | + self.assertEqual(trans, [TransposeType.NOT_TRANSPOSE]) | ||
| 57 | + | ||
| 58 | + def test_row_parallel_linear_no_bias(self): | ||
| 59 | + layer = make_fake_layer(has_bias=False) | ||
| 60 | + weights = RowParallelLinear.get_weights_for_atb_graph(layer, padding=True) | ||
| 61 | + # When bias is None, it should append _PLACEHOLDER after weight | ||
| 62 | + self.assertEqual(len(weights), 1) | ||
| 63 | + | ||
| 64 | + def test_column_parallel_linear(self): | ||
| 65 | + layer = make_fake_layer() | ||
| 66 | + weights = ColumnParallelLinear.get_weights_for_atb_graph(layer, padding=True) | ||
| 67 | + self.assertEqual(len(weights), 1) | ||
| 68 | + | ||
| 69 | + with self.assertRaises(ValueError) as cm: | ||
| 70 | + ColumnParallelLinear.get_weights_for_atb_graph(layer, is_swiglu_quant_enabled=True) | ||
| 71 | + self.assertIn("Cannot set `is_swiglu_quant_enabled` to True", str(cm.exception)) | ||
| 72 | + | ||
| 73 | + descs = ColumnParallelLinear.get_linear_descs(layer) | ||
| 74 | + trans = ColumnParallelLinear.get_weight_transpose_type(layer) | ||
| 75 | + self.assertEqual(descs, [LinearTypeV2.FLOAT16]) | ||
| 76 | + self.assertEqual(trans, [TransposeType.NOT_TRANSPOSE]) | ||
| 77 | + | ||
| 78 | + def test_merged_column_parallel_linear(self): | ||
| 79 | + with patch('mindie_llm.runtime.layers.linear.linear.MergedColumnParallelLinear.__init__', return_value=None): | ||
| 80 | + layer = MergedColumnParallelLinear(input_size=64, output_sizes=[32, 32]) | ||
| 81 | + | ||
| 82 | + layer.weight = MagicMock() | ||
| 83 | + layer.weight.data = torch.randn(64, 64, dtype=torch.float16, device='cpu') | ||
| 84 | + layer.bias = None | ||
| 85 | + layer.quant_method = MagicMock() | ||
| 86 | + layer.quant_method.get_weights_for_atb_graph.return_value = [torch.tensor([1.0], device='cpu')] | ||
| 87 | + layer.quant_method.get_linear_descs.return_value = LinearTypeV2.FLOAT16 | ||
| 88 | + layer.quant_method.get_weight_transpose_type.return_value = TransposeType.NOT_TRANSPOSE | ||
| 89 | + layer._PLACEHOLDER = torch.tensor([888], dtype=torch.float32, device='cpu') | ||
| 90 | + | ||
| 91 | + weights = layer.get_weights_for_atb_graph(padding=True) | ||
| 92 | + self.assertEqual(len(weights), 7) | ||
| 93 | + for w in weights[-6:]: | ||
| 94 | + self.assertEqual(w.item(), 888) | ||
| 95 | + | ||
| 96 | + descs = MergedColumnParallelLinear.get_linear_descs(layer) | ||
| 97 | + trans = MergedColumnParallelLinear.get_weight_transpose_type(layer) | ||
| 98 | + self.assertListEqual(descs, [LinearTypeV2.FLOAT16, LinearTypeV2.INVALID]) | ||
| 99 | + self.assertListEqual(trans, [TransposeType.NOT_TRANSPOSE, TransposeType.INVALID]) | ||
| 100 | + | ||
| 101 | + def test_qkv_parallel_linear(self): | ||
| 102 | + with patch('mindie_llm.runtime.layers.linear.linear.QKVParallelLinear.__init__', return_value=None): | ||
| 103 | + layer = QKVParallelLinear(num_heads=2, num_kv_heads=1, head_size=32) | ||
| 104 | + | ||
| 105 | + layer.weight = MagicMock() | ||
| 106 | + layer.weight.data = torch.randn(96, 64, dtype=torch.float16, device='cpu') # q+k+v | ||
| 107 | + layer.bias = None | ||
| 108 | + layer.quant_method = MagicMock() | ||
| 109 | + layer.quant_method.get_weights_for_atb_graph.return_value = [torch.tensor([1.0], device='cpu')] | ||
| 110 | + layer.quant_method.get_linear_descs.return_value = LinearTypeV2.FLOAT16 | ||
| 111 | + layer.quant_method.get_weight_transpose_type.return_value = TransposeType.NOT_TRANSPOSE | ||
| 112 | + layer._PLACEHOLDER = torch.tensor([777], dtype=torch.float32, device='cpu') | ||
| 113 | + | ||
| 114 | + weights = layer.get_weights_for_atb_graph(padding=True) | ||
| 115 | + self.assertEqual(len(weights), 13) | ||
| 116 | + for w in weights[-12:]: | ||
| 117 | + self.assertEqual(w.item(), 777) | ||
| 118 | + | ||
| 119 | + with self.assertRaises(ValueError): | ||
| 120 | + layer.get_weights_for_atb_graph(is_swiglu_quant_enabled=True) | ||
| 121 | + | ||
| 122 | + descs = layer.get_linear_descs() | ||
| 123 | + trans = layer.get_weight_transpose_type() | ||
| 124 | + self.assertListEqual(descs, [LinearTypeV2.FLOAT16, LinearTypeV2.INVALID, LinearTypeV2.INVALID]) | ||
| 125 | + self.assertListEqual(trans, [TransposeType.NOT_TRANSPOSE, TransposeType.INVALID, TransposeType.INVALID]) | ||
| 126 | + | ||
| 127 | + def test_vocab_parallel_embedding(self): | ||
| 128 | + layer = MagicMock() | ||
| 129 | + layer.quant_method.get_weights_for_atb_graph.return_value = [torch.tensor([[1.0]])] | ||
| 130 | + layer._PLACEHOLDER = torch.tensor([0]) | ||
| 131 | + weights = VocabParallelEmbedding.get_weights_for_atb_graph(layer, padding=True) | ||
| 132 | + expected = [torch.tensor([[1.0]])] | ||
| 133 | + self.assertEqual(len(weights), len(expected)) | ||
| 134 | + self.assertTrue(torch.equal(weights[0], expected[0])) | ||
| 135 | + | ||
| 136 | + def test_parallel_lm_head(self): | ||
| 137 | + layer = make_fake_layer() | ||
| 138 | + weights = ParallelLMHead.get_weights_for_atb_graph(layer, padding=True) | ||
| 139 | + self.assertEqual(len(weights), 1) | ||
| 140 | + descs = ParallelLMHead.get_linear_descs(layer) | ||
| 141 | + trans = ParallelLMHead.get_weight_transpose_type(layer) | ||
| 142 | + self.assertEqual(descs, [LinearTypeV2.FLOAT16]) | ||
| 143 | + self.assertEqual(trans, [TransposeType.NOT_TRANSPOSE]) | ||
| 144 | + | ||
| 145 | + def test_rms_norm(self): | ||
| 146 | + layer = MagicMock() | ||
| 147 | + layer.quant_method.get_weights_for_atb_graph.return_value = [torch.tensor([0.5])] | ||
| 148 | + weights = RMSNorm.get_weights_for_atb_graph(layer, padding=True) | ||
| 149 | + expected = [torch.tensor([0.5])] | ||
| 150 | + self.assertEqual(len(weights), len(expected)) | ||
| 151 | + self.assertTrue(torch.equal(weights[0], expected[0])) | ||
| 152 | + | ||
| 153 | + | ||
| 154 | +if __name__ == '__main__': | ||
| 155 | + unittest.main() | ||
| @@ -0,0 +1,304 @@ | |||
| 1 | +#!/usr/bin/env python | ||
| 2 | +# coding=utf-8 | ||
| 3 | +# Copyright (c) Huawei Technologies Co., Ltd. 2025-2026. All rights reserved. | ||
| 4 | +# MindIE is licensed under Mulan PSL v2. | ||
| 5 | +# You can use this software according to the terms and conditions of the Mulan PSL v2. | ||
| 6 | +# You may obtain a copy of Mulan PSL v2 at: | ||
| 7 | +# http://license.coscl.org.cn/MulanPSL2 | ||
| 8 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, | ||
| 9 | +# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | ||
| 10 | +# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | ||
| 11 | +# See the Mulan PSL v2 for more details. | ||
| 12 | + | ||
| 13 | +import sys | ||
| 14 | +import unittest | ||
| 15 | +from unittest.mock import Mock, patch | ||
| 16 | + | ||
| 17 | +import torch | ||
| 18 | + | ||
| 19 | +from atb_llm.utils.data.quant_method_adapter import ( | ||
| 20 | + UnquantizedLinearMethod, | ||
| 21 | + W8A8PerTensorLinearMethod, | ||
| 22 | + W8A8PerTokenLinearMethod, | ||
| 23 | + W8A8MixLinearMethod, | ||
| 24 | + UnquantizedEmbeddingMethod, | ||
| 25 | + UnquantizedNormMethod, | ||
| 26 | + AntiOutlierNormMethod, | ||
| 27 | + LinearMethodSupportAtbGraph, | ||
| 28 | +) | ||
| 29 | +from atb_llm.utils.quantize.quant_type import LinearTypeV2, QuantType | ||
| 30 | +from atb_llm.utils.quantize.pack_type import TransposeType | ||
| 31 | +from mindie_llm.runtime.layers.quantization.ms_model_slim.quant_type import InferenceMode | ||
| 32 | + | ||
| 33 | + | ||
| 34 | +# Mock torch_npu to avoid NPU dependency | ||
| 35 | +class MockTorchNpu: | ||
| 36 | + @staticmethod | ||
| 37 | + def npu_format_cast_(x, fmt): | ||
| 38 | + return x # no-op | ||
| 39 | + | ||
| 40 | + | ||
| 41 | +# Mock ENV and NPUSocInfo | ||
| 42 | +class MockNPUSocInfo: | ||
| 43 | + def __init__(self): | ||
| 44 | + self.need_nz = False | ||
| 45 | + self.matmul_nd_nz = False | ||
| 46 | + | ||
| 47 | + | ||
| 48 | +class MockENV: | ||
| 49 | + auto_transpose_enable = True | ||
| 50 | + | ||
| 51 | + | ||
| 52 | +def make_fake_layer( | ||
| 53 | + weight_shape=(64, 128), | ||
| 54 | + bias_shape=(64,), | ||
| 55 | + has_bias=True, | ||
| 56 | + weight_dtype=torch.float16, | ||
| 57 | + use_int8_weight=False | ||
| 58 | +): | ||
| 59 | + layer = Mock() | ||
| 60 | + if use_int8_weight: | ||
| 61 | + layer.weight.data = torch.randint(0, 256, weight_shape, dtype=torch.uint8, device='cpu') | ||
| 62 | + else: | ||
| 63 | + layer.weight.data = torch.randn(weight_shape, dtype=weight_dtype, device='cpu') | ||
| 64 | + if has_bias: | ||
| 65 | + layer.bias = Mock() | ||
| 66 | + layer.bias.data = torch.randn(bias_shape, dtype=torch.float32, device='cpu') | ||
| 67 | + else: | ||
| 68 | + layer.bias = None | ||
| 69 | + # For W8A8PerToken | ||
| 70 | + layer.weight_offset = Mock() | ||
| 71 | + layer.weight_offset.data = torch.tensor([10.0], device='cpu') | ||
| 72 | + layer.weight_scale = Mock() | ||
| 73 | + layer.weight_scale.data = torch.tensor([0.5], device='cpu') | ||
| 74 | + # For W8A8PerTensor | ||
| 75 | + layer.quant_bias = Mock() | ||
| 76 | + layer.quant_bias.data = torch.randn(bias_shape, device='cpu') | ||
| 77 | + layer.deq_scale = Mock() | ||
| 78 | + layer.deq_scale.data = torch.tensor(0.1, device='cpu') | ||
| 79 | + layer.input_offset = Mock() | ||
| 80 | + layer.input_offset.data = torch.tensor(128.0, device='cpu') | ||
| 81 | + layer.input_scale = Mock() | ||
| 82 | + layer.input_scale.data = torch.tensor(0.01, device='cpu') | ||
| 83 | + return layer | ||
| 84 | + | ||
| 85 | + | ||
| 86 | +class TestQuantMethodAdapter(unittest.TestCase): | ||
| 87 | + | ||
| 88 | + @patch.dict("sys.modules", {"torch_npu": MockTorchNpu()}) | ||
| 89 | + @patch("atb_llm.utils.data.quant_method_adapter.ENV", MockENV()) | ||
| 90 | + def test_unquantized_linear_method(self): | ||
| 91 | + LinearMethodSupportAtbGraph.set_soc_info(MockNPUSocInfo()) | ||
| 92 | + | ||
| 93 | + adaptee = Mock() | ||
| 94 | + method = UnquantizedLinearMethod(adaptee) | ||
| 95 | + method._PLACEHOLDER = torch.tensor([999.0], device='cpu') | ||
| 96 | + | ||
| 97 | + # Test with bias, padding=True | ||
| 98 | + layer = make_fake_layer(has_bias=True) | ||
| 99 | + weights = method.get_weights_for_atb_graph(layer, padding=True) | ||
| 100 | + self.assertEqual(len(weights), 6) | ||
| 101 | + | ||
| 102 | + # Test padding=False | ||
| 103 | + weights = method.get_weights_for_atb_graph(layer, padding=False) | ||
| 104 | + self.assertEqual(len(weights), 2) | ||
| 105 | + | ||
| 106 | + # Test dtype | ||
| 107 | + descs = method.get_linear_descs(layer) | ||
| 108 | + self.assertEqual(descs, LinearTypeV2.FLOAT16) | ||
| 109 | + | ||
| 110 | + layer_bf16 = make_fake_layer(weight_dtype=torch.bfloat16) | ||
| 111 | + descs_bf16 = method.get_linear_descs(layer_bf16) | ||
| 112 | + self.assertEqual(descs_bf16, LinearTypeV2.BFLOAT16) | ||
| 113 | + | ||
| 114 | + # Test process | ||
| 115 | + method.process_weights_after_loading(layer) | ||
| 116 | + | ||
| 117 | + LinearMethodSupportAtbGraph._soc_info = None | ||
| 118 | + | ||
| 119 | + @patch.dict("sys.modules", {"torch_npu": MockTorchNpu()}) | ||
| 120 | + @patch("atb_llm.utils.data.quant_method_adapter.ENV", MockENV()) | ||
| 121 | + def test_w8a8_per_tensor(self): | ||
| 122 | + LinearMethodSupportAtbGraph.set_soc_info(MockNPUSocInfo()) | ||
| 123 | + | ||
| 124 | + adaptee = Mock() | ||
| 125 | + method = W8A8PerTensorLinearMethod(adaptee) | ||
| 126 | + method._PLACEHOLDER = torch.tensor([888.0], device='cpu') | ||
| 127 | + | ||
| 128 | + layer = make_fake_layer(use_int8_weight=True, has_bias=True) | ||
| 129 | + weights = method.get_weights_for_atb_graph(layer, padding=True) | ||
| 130 | + self.assertEqual(len(weights), 6) | ||
| 131 | + | ||
| 132 | + weights = method.get_weights_for_atb_graph(layer, padding=False) | ||
| 133 | + self.assertEqual(len(weights), 5) | ||
| 134 | + | ||
| 135 | + self.assertEqual(method.get_linear_descs(layer), LinearTypeV2.W8A8) | ||
| 136 | + method.process_weights_after_loading(layer) | ||
| 137 | + | ||
| 138 | + LinearMethodSupportAtbGraph._soc_info = None | ||
| 139 | + | ||
| 140 | + @patch.dict("sys.modules", {"torch_npu": MockTorchNpu()}) | ||
| 141 | + @patch("atb_llm.utils.data.quant_method_adapter.ENV", MockENV()) | ||
| 142 | + def test_w8a8_per_token(self): | ||
| 143 | + LinearMethodSupportAtbGraph.set_soc_info(MockNPUSocInfo()) | ||
| 144 | + | ||
| 145 | + adaptee = Mock() | ||
| 146 | + method = W8A8PerTokenLinearMethod(adaptee) | ||
| 147 | + method._PLACEHOLDER = torch.tensor([777.0], device='cpu') | ||
| 148 | + | ||
| 149 | + layer = make_fake_layer(use_int8_weight=True, has_bias=False) | ||
| 150 | + weights = method.get_weights_for_atb_graph(layer, padding=True) | ||
| 151 | + self.assertEqual(len(weights), 6) | ||
| 152 | + | ||
| 153 | + self.assertEqual(method.get_linear_descs(layer), LinearTypeV2.W8A8_DYNAMIC) | ||
| 154 | + method.process_weights_after_loading(layer) | ||
| 155 | + | ||
| 156 | + LinearMethodSupportAtbGraph._soc_info = None | ||
| 157 | + | ||
| 158 | + @patch.dict("sys.modules", {"torch_npu": MockTorchNpu()}) | ||
| 159 | + @patch("atb_llm.utils.data.quant_method_adapter.ENV", MockENV()) | ||
| 160 | + def test_w8a8_mix_linear_method(self): | ||
| 161 | + LinearMethodSupportAtbGraph.set_soc_info(MockNPUSocInfo()) | ||
| 162 | + | ||
| 163 | + prefill_adaptee = Mock() | ||
| 164 | + decode_adaptee = Mock() | ||
| 165 | + adaptee = Mock() | ||
| 166 | + adaptee.quant_method = { | ||
| 167 | + InferenceMode.PREFILL: prefill_adaptee, | ||
| 168 | + InferenceMode.DECODE: decode_adaptee, | ||
| 169 | + } | ||
| 170 | + | ||
| 171 | + with patch("atb_llm.utils.data.quant_method_adapter.W8A8PerTokenLinearMethod") as mock_prefill, \ | ||
| 172 | + patch("atb_llm.utils.data.quant_method_adapter.W8A8PerTensorLinearMethod") as mock_decode: | ||
| 173 | + mock_prefill_instance = Mock() | ||
| 174 | + mock_prefill_instance.get_weights_for_atb_graph.return_value = ["prefill"] | ||
| 175 | + mock_prefill.return_value = mock_prefill_instance | ||
| 176 | + | ||
| 177 | + mock_decode_instance = Mock() | ||
| 178 | + mock_decode_instance.get_weights_for_atb_graph.return_value = ["decode"] | ||
| 179 | + mock_decode.return_value = mock_decode_instance | ||
| 180 | + | ||
| 181 | + method = W8A8MixLinearMethod(adaptee) | ||
| 182 | + | ||
| 183 | + layer = Mock() | ||
| 184 | + layer.weight = Mock() | ||
| 185 | + layer.weight.data = torch.randint(0, 256, (64, 128), dtype=torch.uint8, device='cpu') | ||
| 186 | + layer.weight_scale = Mock() | ||
| 187 | + layer.weight_scale.data = torch.tensor([0.5, 0.6], device='cpu') | ||
| 188 | + layer.weight_offset = Mock() | ||
| 189 | + layer.weight_offset.data = torch.tensor([10.0, 12.0], device='cpu') | ||
| 190 | + | ||
| 191 | + weights = method.get_weights_for_atb_graph(layer, quant_type=QuantType.W8A8_DYNAMIC) | ||
| 192 | + self.assertEqual(weights, ["prefill"]) | ||
| 193 | + | ||
| 194 | + weights = method.get_weights_for_atb_graph(layer, quant_type=QuantType.W8A8) | ||
| 195 | + self.assertEqual(weights, ["decode"]) | ||
| 196 | + | ||
| 197 | + self.assertEqual(method.get_linear_descs(layer), LinearTypeV2.W8A8_PDMIX) | ||
| 198 | + method.process_weights_after_loading(layer) | ||
| 199 | + | ||
| 200 | + LinearMethodSupportAtbGraph._soc_info = None | ||
| 201 | + | ||
| 202 | + def test_unquantized_embedding_method(self): | ||
| 203 | + adaptee = Mock() | ||
| 204 | + method = UnquantizedEmbeddingMethod(adaptee) | ||
| 205 | + method._PLACEHOLDER = torch.tensor([666.0], device='cpu') | ||
| 206 | + | ||
| 207 | + layer = Mock() | ||
| 208 | + layer.weight.data = torch.randn(100, 64, device='cpu') | ||
| 209 | + | ||
| 210 | + weights = method.get_weights_for_atb_graph(layer, padding=True) | ||
| 211 | + self.assertEqual(len(weights), 1) | ||
| 212 | + self.assertTrue(torch.equal(weights[0], layer.weight.data)) | ||
| 213 | + | ||
| 214 | + with self.assertRaises(NotImplementedError) as cm: | ||
| 215 | + method.get_weights_for_atb_graph(layer, padding=False) | ||
| 216 | + self.assertIn("cannot be set to False", str(cm.exception)) | ||
| 217 | + | ||
| 218 | + def test_unquantized_norm_method(self): | ||
| 219 | + adaptee = Mock() | ||
| 220 | + method = UnquantizedNormMethod(adaptee) | ||
| 221 | + method._PLACEHOLDER = torch.tensor([555.0], device='cpu') | ||
| 222 | + | ||
| 223 | + layer = Mock() | ||
| 224 | + layer.weight.data = torch.randn(64, device='cpu') | ||
| 225 | + | ||
| 226 | + weights = method.get_weights_for_atb_graph(layer, padding=False) | ||
| 227 | + self.assertEqual(len(weights), 1) | ||
| 228 | + self.assertTrue(torch.equal(weights[0], layer.weight.data)) | ||
| 229 | + | ||
| 230 | + weights = method.get_weights_for_atb_graph(layer, padding=True) | ||
| 231 | + self.assertEqual(len(weights), 4) | ||
| 232 | + self.assertTrue(torch.equal(weights[0], layer.weight.data)) | ||
| 233 | + for w in weights[1:]: | ||
| 234 | + self.assertEqual(w.item(), 555.0) | ||
| 235 | + | ||
| 236 | + def test_anti_outlier_norm_method(self): | ||
| 237 | + adaptee = Mock() | ||
| 238 | + method = AntiOutlierNormMethod(adaptee) | ||
| 239 | + method._PLACEHOLDER = torch.tensor([444.0], device='cpu') | ||
| 240 | + | ||
| 241 | + layer = Mock() | ||
| 242 | + layer.weight.data = torch.randn(64, device='cpu') | ||
| 243 | + layer.bias = Mock() | ||
| 244 | + layer.bias.data = torch.randn(64, device='cpu') | ||
| 245 | + | ||
| 246 | + weights = method.get_weights_for_atb_graph(layer, padding=False) | ||
| 247 | + self.assertEqual(len(weights), 2) | ||
| 248 | + | ||
| 249 | + weights = method.get_weights_for_atb_graph(layer, padding=True) | ||
| 250 | + self.assertEqual(len(weights), 4) | ||
| 251 | + self.assertTrue(torch.equal(weights[0], layer.weight.data)) | ||
| 252 | + self.assertTrue(torch.equal(weights[1], layer.bias.data)) | ||
| 253 | + for w in weights[2:]: | ||
| 254 | + self.assertEqual(w.item(), 444.0) | ||
| 255 | + | ||
| 256 | + layer_zero_bias = Mock() | ||
| 257 | + layer_zero_bias.weight.data = torch.randn(64, device='cpu') | ||
| 258 | + layer_zero_bias.bias = Mock() | ||
| 259 | + layer_zero_bias.bias.data = torch.zeros(64, device='cpu') | ||
| 260 | + weights = method.get_weights_for_atb_graph(layer_zero_bias, padding=False) | ||
| 261 | + self.assertEqual(len(weights), 1) | ||
| 262 | + | ||
| 263 | + @patch.dict("sys.modules", {"torch_npu": MockTorchNpu()}) | ||
| 264 | + @patch("atb_llm.utils.data.quant_method_adapter.ENV", MockENV()) | ||
| 265 | + def test_unquantized_linear_method_full_coverage(self): | ||
| 266 | + """Cover missing branches in UnquantizedLinearMethod""" | ||
| 267 | + LinearMethodSupportAtbGraph.set_soc_info(MockNPUSocInfo()) | ||
| 268 | + | ||
| 269 | + adaptee = Mock() | ||
| 270 | + method = UnquantizedLinearMethod(adaptee) | ||
| 271 | + method._PLACEHOLDER = torch.tensor([1.0], device='cpu') | ||
| 272 | + | ||
| 273 | + # 1. bias=None + padding=True | ||
| 274 | + layer_no_bias = Mock() | ||
| 275 | + layer_no_bias.weight.data = torch.randn(128, 64, dtype=torch.float16, device='cpu') | ||
| 276 | + layer_no_bias.bias = None | ||
| 277 | + weights = method.get_weights_for_atb_graph(layer_no_bias, padding=True) | ||
| 278 | + self.assertEqual(len(weights), 6) | ||
| 279 | + self.assertEqual(weights[1].item(), 1.0) | ||
| 280 | + | ||
| 281 | + # 2. process with NOT_TRANSPOSE | ||
| 282 | + with patch.object(method, '_check_transpose', return_value=TransposeType.NOT_TRANSPOSE): | ||
| 283 | + layer = Mock() | ||
| 284 | + layer.weight.data = torch.randn(128, 64, device='cpu') | ||
| 285 | + original = layer.weight.data.clone() | ||
| 286 | + method.process_weights_after_loading(layer) | ||
| 287 | + self.assertEqual(layer.weight.data.shape, (64, 128)) | ||
| 288 | + self.assertTrue(torch.equal(layer.weight.data, original.t())) | ||
| 289 | + | ||
| 290 | + # 3. get_weight_transpose_type after update | ||
| 291 | + layer = Mock() | ||
| 292 | + layer.weight.data = torch.randn(128, 64, device='cpu') | ||
| 293 | + method._weight_transpose_type = TransposeType.TRANSPOSE | ||
| 294 | + self.assertEqual(method.get_weight_transpose_type(layer), TransposeType.TRANSPOSE) | ||
| 295 | + | ||
| 296 | + with patch.object(method, '_check_transpose', return_value=TransposeType.NOT_TRANSPOSE): | ||
| 297 | + method.process_weights_after_loading(layer) | ||
| 298 | + self.assertEqual(method.get_weight_transpose_type(layer), TransposeType.NOT_TRANSPOSE) | ||
| 299 | + | ||
| 300 | + LinearMethodSupportAtbGraph._soc_info = None | ||
| 301 | + | ||
| 302 | + | ||
| 303 | +if __name__ == '__main__': | ||
| 304 | + unittest.main() | ||