已合并
acl graph相关代码cleancode整改 #159
huanglei创建于 1月7日
acl graph相关代码cleancode整改 #159
已合并
共 18 个文件变更+118-105
| @@ -196,5 +196,5 @@ class RMSNorm(RMSNormAdaptee, LayerSupportAtbGraph): | |||
| 196 | kwargs.get("quant_config"), UnquantizedNormMethod) | 196 | kwargs.get("quant_config"), UnquantizedNormMethod) |
| 197 | super().__init__(*args, **kwargs) | 197 | super().__init__(*args, **kwargs) |
| 198 | 198 | ||
| 199 | - def get_weights_for_atb_graph(self, padding: bool = True): | 199 | + def get_weights_for_atb_graph(self, padding: bool = True) -> list[torch.Tensor]: |
| 200 | return self.quant_method.get_weights_for_atb_graph(self, padding=padding) | 200 | return self.quant_method.get_weights_for_atb_graph(self, padding=padding) |
| @@ -45,7 +45,7 @@ class QuantizationConfig: | |||
| 45 | self.default_quant_method_cls = default_quant_method_cls | 45 | self.default_quant_method_cls = default_quant_method_cls |
| 46 | 46 | ||
| 47 | 47 | ||
| 48 | - def _get_method_adpater(quant_method: QuantizationMethodBase): | 48 | + def _get_method_adapter(quant_method: QuantizationMethodBase) -> QuantizationMethodBase | None: |
| 49 | quant_method_cls_adapter_map = { | 49 | quant_method_cls_adapter_map = { |
| 50 | UnquantizedLinearMethodAdaptee: UnquantizedLinearMethod, | 50 | UnquantizedLinearMethodAdaptee: UnquantizedLinearMethod, |
| 51 | UnquantizedEmbeddingMethodAdaptee: UnquantizedEmbeddingMethod, | 51 | UnquantizedEmbeddingMethodAdaptee: UnquantizedEmbeddingMethod, |
| @@ -60,7 +60,7 @@ class QuantizationConfig: | |||
| 60 | return adapter_cls(quant_method) | 60 | return adapter_cls(quant_method) |
| 61 | raise NotImplementedError(f"Cannot found the adapter class for `{quant_method}`") | 61 | raise NotImplementedError(f"Cannot found the adapter class for `{quant_method}`") |
| 62 | 62 | ||
| 63 | - def get_quant_type_by_weight_name(self, *args, **kwargs) -> QuantizationMethodBase | None: | 63 | + def get_quant_type_by_weight_name(self, *args, **kwargs) -> str: |
| 64 | if self.adaptee is None: | 64 | if self.adaptee is None: |
| 65 | raise ValueError(f"NoneType adaptee doesn't support `get_quant_type_by_weight_name`.") | 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) | 66 | return self.adaptee.get_quant_type_by_weight_name(*args, **kwargs) |
| @@ -72,7 +72,7 @@ class QuantizationConfig: | |||
| 72 | quant_method = self.adaptee.get_quant_method(layer, prefix) | 72 | quant_method = self.adaptee.get_quant_method(layer, prefix) |
| 73 | else: | 73 | else: |
| 74 | quant_method = self.default_quant_method_cls() | 74 | quant_method = self.default_quant_method_cls() |
| 75 | - return self._get_method_adpater(quant_method) | 75 | + return self._get_method_adapter(quant_method) |
| 76 | 76 | ||
| 77 | 77 | ||
| 78 | class MethodSupportAtbGraph(ABC): | 78 | class MethodSupportAtbGraph(ABC): |
| @@ -81,13 +81,13 @@ class MethodSupportAtbGraph(ABC): | |||
| 81 | def __init__(self, adaptee: QuantizationMethodBase): | 81 | def __init__(self, adaptee: QuantizationMethodBase): |
| 82 | self.adaptee = adaptee | 82 | self.adaptee = adaptee |
| 83 | 83 | ||
| 84 | - def create_weights(self, *args, **kwargs): | 84 | + def create_weights(self, *args, **kwargs) -> None: |
| 85 | self.adaptee.create_weights(*args, **kwargs) | 85 | self.adaptee.create_weights(*args, **kwargs) |
| 86 | 86 | ||
| 87 | - def apply(self, *args, **kwargs): | 87 | + def apply(self, *args, **kwargs) -> None: |
| 88 | self.adaptee.apply(*args, **kwargs) | 88 | self.adaptee.apply(*args, **kwargs) |
| 89 | 89 | ||
| 90 | - def process_weights_after_loading(self, *args, **kwargs): | 90 | + def process_weights_after_loading(self, *args, **kwargs) -> None: |
| 91 | self.adaptee.process_weights_after_loading(*args, **kwargs) | 91 | self.adaptee.process_weights_after_loading(*args, **kwargs) |
| 92 | 92 | ||
| 93 | 93 | ||
| @@ -121,7 +121,7 @@ class LinearMethodSupportAtbGraph(MethodSupportAtbGraph): | |||
| 121 | def get_weight_transpose_type(self, layer: LinearBase) -> list[TransposeType]: | 121 | def get_weight_transpose_type(self, layer: LinearBase) -> list[TransposeType]: |
| 122 | pass | 122 | pass |
| 123 | 123 | ||
| 124 | - def _check_transpose(self, weight_shape): | 124 | + def _check_transpose(self, weight_shape: torch.Size) -> TransposeType: |
| 125 | if self._soc_info is None: | 125 | if self._soc_info is None: |
| 126 | raise ValueError("``NPUSocInfo` is not set in `LinearMethodSupportAtbGraph`.") | 126 | raise ValueError("``NPUSocInfo` is not set in `LinearMethodSupportAtbGraph`.") |
| 127 | 127 | ||
| @@ -132,8 +132,8 @@ class LinearMethodSupportAtbGraph(MethodSupportAtbGraph): | |||
| 132 | # transpose weights to [k, n] when using nz format | 132 | # transpose weights to [k, n] when using nz format |
| 133 | return TransposeType.NOT_TRANSPOSE | 133 | return TransposeType.NOT_TRANSPOSE |
| 134 | 134 | ||
| 135 | - is_k_divisible = weight_shape[-1] % 256 == 0 | 135 | + is_k_divisible = weight_shape[-1] % 256 == 0 # Input dimension (k) alignment check |
| 136 | - is_n_divisible = weight_shape[-2] % 256 == 0 | 136 | + is_n_divisible = weight_shape[-2] % 256 == 0 # Output dimension (n) alignment check |
| 137 | if not is_k_divisible and is_n_divisible and len(weight_shape) != 3: | 137 | if not is_k_divisible and is_n_divisible and len(weight_shape) != 3: |
| 138 | return TransposeType.NOT_TRANSPOSE | 138 | return TransposeType.NOT_TRANSPOSE |
| 139 | return TransposeType.TRANSPOSE | 139 | return TransposeType.TRANSPOSE |
| @@ -61,7 +61,7 @@ class LLMConfig: | |||
| 61 | "".join(section_reprs) + \ | 61 | "".join(section_reprs) + \ |
| 62 | "\n)" | 62 | "\n)" |
| 63 | 63 | ||
| 64 | - def update(self, config_dict: Dict[str, Any], allow_new_keys: bool = False, current_path: str = ''): | 64 | + def update(self, config_dict: Dict[str, Any], allow_new_keys: bool = False, current_path: str = '') -> None: |
| 65 | """ | 65 | """ |
| 66 | Update configuration values with provided keyword arguments | 66 | Update configuration values with provided keyword arguments |
| 67 | Args: | 67 | Args: |
| @@ -73,7 +73,7 @@ class LLMConfig: | |||
| 73 | return | 73 | return |
| 74 | self._recursive_update(base=self, update=config_dict, allow_new_keys=allow_new_keys, current_path=current_path) | 74 | self._recursive_update(base=self, update=config_dict, allow_new_keys=allow_new_keys, current_path=current_path) |
| 75 | 75 | ||
| 76 | - def merge_models_config(self, model_name: str): | 76 | + def merge_models_config(self, model_name: str) -> None: |
| 77 | """merge model.model_name config to llm config""" | 77 | """merge model.model_name config to llm config""" |
| 78 | if not hasattr(self, 'models'): | 78 | if not hasattr(self, 'models'): |
| 79 | return | 79 | return |
| @@ -123,7 +123,7 @@ class LLMConfig: | |||
| 123 | return current | 123 | return current |
| 124 | 124 | ||
| 125 | def _init_default_config(self): | 125 | def _init_default_config(self): |
| 126 | - """初始化默认配置结构""" | 126 | + """Initialize the default configuration structure""" |
| 127 | self._apply_config(self._DEFAULT_CONFIG) | 127 | self._apply_config(self._DEFAULT_CONFIG) |
| 128 | 128 | ||
| 129 | def _load_config(self) -> None: | 129 | def _load_config(self) -> None: |
| @@ -43,7 +43,7 @@ class MindIELLMConfig: | |||
| 43 | self.soc_info = PlatformInfo() | 43 | self.soc_info = PlatformInfo() |
| 44 | self.quant_config = self._init_quant_config() | 44 | self.quant_config = self._init_quant_config() |
| 45 | 45 | ||
| 46 | - def _init_quant_config(self): | 46 | + def _init_quant_config(self) -> QuantizationConfigBase | None: |
| 47 | # get quant config class | 47 | # get quant config class |
| 48 | # NOTE: Since only `ms_model_slim.QuantizationConfig` is currently supported. | 48 | # NOTE: Since only `ms_model_slim.QuantizationConfig` is currently supported. |
| 49 | # This is a straightforward implementation. A dispatch mechanism | 49 | # This is a straightforward implementation. A dispatch mechanism |
| @@ -31,7 +31,7 @@ class FusedMoEMethodBase(QuantizationMethodBase): | |||
| 31 | intermediate_size_per_partition: int, | 31 | intermediate_size_per_partition: int, |
| 32 | params_dtype: torch.dtype, | 32 | params_dtype: torch.dtype, |
| 33 | **extra_weight_attrs, | 33 | **extra_weight_attrs, |
| 34 | - ): | 34 | + ) -> None: |
| 35 | raise NotImplementedError | 35 | raise NotImplementedError |
| 36 | 36 | ||
| 37 | 37 | ||
| @@ -92,13 +92,13 @@ class LinearBase(CustomLayer): | |||
| 92 | self._post_init() | 92 | self._post_init() |
| 93 | self._create_weights() | 93 | self._create_weights() |
| 94 | 94 | ||
| 95 | - def weight_loader(self, param: BaseParameter, loaded_weight: torch.Tensor, **kwargs): | 95 | + def weight_loader(self, param: BaseParameter, loaded_weight: torch.Tensor, **kwargs) -> None: |
| 96 | param.load_weight(loaded_weight=loaded_weight) | 96 | param.load_weight(loaded_weight=loaded_weight) |
| 97 | 97 | ||
| 98 | - def _post_init(self): | 98 | + def _post_init(self) -> None: |
| 99 | pass | 99 | pass |
| 100 | 100 | ||
| 101 | - def _create_weights(self): | 101 | + def _create_weights(self) -> None: |
| 102 | """Initializes weights and bias based on quantization method.""" | 102 | """Initializes weights and bias based on quantization method.""" |
| 103 | self.quant_method.create_weights( | 103 | self.quant_method.create_weights( |
| 104 | layer=self, | 104 | layer=self, |
| @@ -230,7 +230,7 @@ class RowParallelLinear(LinearBase): | |||
| 230 | parallel_info=parallel_info | 230 | parallel_info=parallel_info |
| 231 | ) | 231 | ) |
| 232 | 232 | ||
| 233 | - def weight_loader(self, param: BaseParameter, loaded_weight: torch.Tensor): | 233 | + def weight_loader(self, param: BaseParameter, loaded_weight: torch.Tensor) -> None: |
| 234 | if isinstance(param, RowParameter): | 234 | if isinstance(param, RowParameter): |
| 235 | param.load_row_parallel_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank) | 235 | param.load_row_parallel_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank) |
| 236 | else: | 236 | else: |
| @@ -272,7 +272,7 @@ class RowParallelLinear(LinearBase): | |||
| 272 | s += f", tp_size={self.tp_size}" | 272 | s += f", tp_size={self.tp_size}" |
| 273 | return s | 273 | return s |
| 274 | 274 | ||
| 275 | - def _post_init(self): | 275 | + def _post_init(self) -> None: |
| 276 | self.input_size_per_partition = even_divide(self.input_size, self.tp_size) | 276 | self.input_size_per_partition = even_divide(self.input_size, self.tp_size) |
| 277 | self.output_partition_sizes = [self.output_size] | 277 | self.output_partition_sizes = [self.output_size] |
| 278 | 278 | ||
| @@ -327,9 +327,7 @@ class ColumnParallelLinear(LinearBase): | |||
| 327 | parallel_info=parallel_info | 327 | parallel_info=parallel_info |
| 328 | ) | 328 | ) |
| 329 | 329 | ||
| 330 | - | 330 | + def weight_loader(self, param: BaseParameter, loaded_weight: torch.Tensor) -> None: |
| 331 | - | ||
| 332 | - def weight_loader(self, param: BaseParameter, loaded_weight: torch.Tensor): | ||
| 333 | if isinstance(param, ColumnParameter): | 331 | if isinstance(param, ColumnParameter): |
| 334 | param.load_column_parallel_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank) | 332 | param.load_column_parallel_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank) |
| 335 | else: | 333 | else: |
| @@ -366,14 +364,14 @@ class ColumnParallelLinear(LinearBase): | |||
| 366 | s += f", gather_output={self.gather_output}" | 364 | s += f", gather_output={self.gather_output}" |
| 367 | return s | 365 | return s |
| 368 | 366 | ||
| 369 | - def _post_init(self): | 367 | + def _post_init(self) -> None: |
| 370 | self.output_partition_sizes = [even_divide(self.output_size, self.tp_size)] | 368 | self.output_partition_sizes = [even_divide(self.output_size, self.tp_size)] |
| 371 | 369 | ||
| 372 | if self.parallel_info is not None: | 370 | if self.parallel_info is not None: |
| 373 | self.tp_rank = self.parallel_info.rank | 371 | self.tp_rank = self.parallel_info.rank |
| 374 | self.tp_size = self.parallel_info.group_size | 372 | self.tp_size = self.parallel_info.group_size |
| 375 | 373 | ||
| 376 | - def _create_weights(self): | 374 | + def _create_weights(self) -> None: |
| 377 | self.quant_method.create_weights( | 375 | self.quant_method.create_weights( |
| 378 | layer=self, | 376 | layer=self, |
| 379 | input_size_per_partition=self.input_size_per_partition, | 377 | input_size_per_partition=self.input_size_per_partition, |
| @@ -437,7 +435,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear): | |||
| 437 | param: BaseParameter, | 435 | param: BaseParameter, |
| 438 | loaded_weight: torch.Tensor, | 436 | loaded_weight: torch.Tensor, |
| 439 | loaded_shard_id: int | None = None, | 437 | loaded_shard_id: int | None = None, |
| 440 | - ): | 438 | + ) -> None: |
| 441 | """ | 439 | """ |
| 442 | Loads weights for a specific shard of the merged linear layer. | 440 | Loads weights for a specific shard of the merged linear layer. |
| 443 | 441 | ||
| @@ -446,11 +444,14 @@ class MergedColumnParallelLinear(ColumnParallelLinear): | |||
| 446 | loaded_weight: The weight tensor to load. | 444 | loaded_weight: The weight tensor to load. |
| 447 | loaded_shard_id: The index of the shard (corresponding to the index in output_sizes). | 445 | loaded_shard_id: The index of the shard (corresponding to the index in output_sizes). |
| 448 | """ | 446 | """ |
| 447 | + | ||
| 448 | + # Validate shard ID | ||
| 449 | if loaded_shard_id >= len(self.output_sizes): | 449 | if loaded_shard_id >= len(self.output_sizes): |
| 450 | raise ValueError( | 450 | raise ValueError( |
| 451 | f"The parameter `loaded_shard_id` {loaded_shard_id} exceeds the valid range of " | 451 | f"The parameter `loaded_shard_id` {loaded_shard_id} exceeds the valid range of " |
| 452 | f"indices for the `output_sizes` array {self.output_sizes} defined in `MergedColumnParallelLinear`.") | 452 | f"indices for the `output_sizes` array {self.output_sizes} defined in `MergedColumnParallelLinear`.") |
| 453 | 453 | ||
| 454 | + # Obtain a tensor slice of size `shard_size` starting from `shard_offset` in self.output_sizes | ||
| 454 | shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size | 455 | shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size |
| 455 | shard_size = self.output_sizes[loaded_shard_id] // self.tp_size | 456 | shard_size = self.output_sizes[loaded_shard_id] // self.tp_size |
| 456 | 457 | ||
| @@ -464,7 +465,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear): | |||
| 464 | else: | 465 | else: |
| 465 | param.load_weight(loaded_weight=loaded_weight) | 466 | param.load_weight(loaded_weight=loaded_weight) |
| 466 | 467 | ||
| 467 | - def _post_init(self): | 468 | + def _post_init(self) -> None: |
| 468 | for output_size in self.output_sizes: | 469 | for output_size in self.output_sizes: |
| 469 | if output_size % self.tp_size != 0: | 470 | if output_size % self.tp_size != 0: |
| 470 | raise ValueError(f"All `output_sizes` {self.output_sizes} in `MergedColumnParallelLinear` " | 471 | raise ValueError(f"All `output_sizes` {self.output_sizes} in `MergedColumnParallelLinear` " |
| @@ -524,6 +525,7 @@ class QKVParallelLinear(ColumnParallelLinear): | |||
| 524 | self.num_kv_heads = even_divide(self.total_num_kv_heads, tp_size) | 525 | self.num_kv_heads = even_divide(self.total_num_kv_heads, tp_size) |
| 525 | self.num_kv_head_replicas = 1 | 526 | self.num_kv_head_replicas = 1 |
| 526 | 527 | ||
| 528 | + # 2: K and V each need their own heads (K has self.num_kv_heads, V has self.num_kv_heads) | ||
| 527 | output_size = ( | 529 | output_size = ( |
| 528 | (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size | 530 | (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size |
| 529 | ) | 531 | ) |
| @@ -547,7 +549,7 @@ class QKVParallelLinear(ColumnParallelLinear): | |||
| 547 | param: BaseParameter, | 549 | param: BaseParameter, |
| 548 | loaded_weight: torch.Tensor, | 550 | loaded_weight: torch.Tensor, |
| 549 | loaded_shard_id: int | None = None, | 551 | loaded_shard_id: int | None = None, |
| 550 | - ): | 552 | + ) -> None: |
| 551 | """ | 553 | """ |
| 552 | Loads weights for Q, K, or V projections. | 554 | Loads weights for Q, K, or V projections. |
| 553 | 555 | ||
| @@ -582,7 +584,7 @@ class QKVParallelLinear(ColumnParallelLinear): | |||
| 582 | self.num_kv_heads * self.head_size, # v | 584 | self.num_kv_heads * self.head_size, # v |
| 583 | ] | 585 | ] |
| 584 | 586 | ||
| 585 | - def _get_shard_offset_mapping(self, loaded_shard_id: str): | 587 | + def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int: |
| 586 | """Returns the offset in the weight matrix for a given shard ID.""" | 588 | """Returns the offset in the weight matrix for a given shard ID.""" |
| 587 | shard_offset_mapping = { | 589 | shard_offset_mapping = { |
| 588 | 0: 0, | 590 | 0: 0, |
| @@ -592,7 +594,7 @@ class QKVParallelLinear(ColumnParallelLinear): | |||
| 592 | } | 594 | } |
| 593 | return shard_offset_mapping.get(loaded_shard_id) | 595 | return shard_offset_mapping.get(loaded_shard_id) |
| 594 | 596 | ||
| 595 | - def _get_shard_size_mapping(self, loaded_shard_id: str): | 597 | + def _get_shard_size_mapping(self, loaded_shard_id: str) -> int: |
| 596 | """Returns the size of the shard for a given shard ID.""" | 598 | """Returns the size of the shard for a given shard ID.""" |
| 597 | shard_size_mapping = { | 599 | shard_size_mapping = { |
| 598 | 0: self.num_heads * self.head_size, | 600 | 0: self.num_heads * self.head_size, |
| @@ -30,7 +30,7 @@ class LinearMethodBase(QuantizationMethodBase): | |||
| 30 | weight_dtype: torch.dtype, | 30 | weight_dtype: torch.dtype, |
| 31 | bias_dtype: torch.dtype, | 31 | bias_dtype: torch.dtype, |
| 32 | **extra_weight_attrs, | 32 | **extra_weight_attrs, |
| 33 | - ): | 33 | + ) -> None: |
| 34 | """Creates weights for the linear layer. | 34 | """Creates weights for the linear layer. |
| 35 | 35 | ||
| 36 | Args: | 36 | Args: |
| @@ -26,7 +26,7 @@ class AntiOutlierNormMethod(QuantizationMethodBase): | |||
| 26 | hidden_size: int, | 26 | hidden_size: int, |
| 27 | params_dtype: torch.dtype, | 27 | params_dtype: torch.dtype, |
| 28 | **extra_weight_attrs, | 28 | **extra_weight_attrs, |
| 29 | - ): | 29 | + ) -> None: |
| 30 | """ | 30 | """ |
| 31 | Args: | 31 | Args: |
| 32 | layer: The layer instance to register weights to | 32 | layer: The layer instance to register weights to |
| @@ -47,7 +47,7 @@ class QuantizationConfig(QuantizationConfigBase): | |||
| 47 | def from_config(cls, config: dict[str, Any]) -> QuantizationConfigBase: | 47 | def from_config(cls, config: dict[str, Any]) -> QuantizationConfigBase: |
| 48 | return cls(config) | 48 | return cls(config) |
| 49 | 49 | ||
| 50 | - def get_quant_type_by_weight_name(self, prefix: str | list[str], suffix: str) -> QuantizationMethodBase | None: | 50 | + def get_quant_type_by_weight_name(self, prefix: str | list[str], suffix: str) -> str: |
| 51 | """ | 51 | """ |
| 52 | Retrieve the quantization type for a specific weight parameter. | 52 | Retrieve the quantization type for a specific weight parameter. |
| 53 | Args: | 53 | Args: |
| @@ -28,7 +28,7 @@ from mindie_llm.runtime.utils.distributed.utils import even_divide | |||
| 28 | from mindie_llm.utils.log.logging import logger | 28 | from mindie_llm.utils.log.logging import logger |
| 29 | 29 | ||
| 30 | 30 | ||
| 31 | -SUPPORT_NZ_NPU_LIST = ("Ascend910B3", "Ascend910B4_1", "Ascend910_9381", "Ascend910_9372") | 31 | +SUPPORT_NZ_NPU_LIST = ("Ascend910B3", "Ascend910B4_1", "Ascend910_9382", "Ascend910_9362") |
| 32 | MXFP8_GROUP_SIZE = 32 | 32 | MXFP8_GROUP_SIZE = 32 |
| 33 | 33 | ||
| 34 | 34 | ||
| @@ -45,7 +45,7 @@ class W8A8PerTensorLinearMethod(LinearMethodBase): | |||
| 45 | weight_dtype: torch.dtype, | 45 | weight_dtype: torch.dtype, |
| 46 | bias_dtype: torch.dtype, | 46 | bias_dtype: torch.dtype, |
| 47 | **extra_weight_attrs, | 47 | **extra_weight_attrs, |
| 48 | - ): | 48 | + ) -> None: |
| 49 | """ | 49 | """ |
| 50 | Creates and registers quantized weights and scales. | 50 | Creates and registers quantized weights and scales. |
| 51 | 51 | ||
| @@ -80,6 +80,7 @@ class W8A8PerTensorLinearMethod(LinearMethodBase): | |||
| 80 | deq_scale_dtype = torch.float32 | 80 | deq_scale_dtype = torch.float32 |
| 81 | else: | 81 | else: |
| 82 | raise ValueError(f"Dtype {weight_dtype} is not supported in `W8A8PerTensorLinearMethod`.") | 82 | raise ValueError(f"Dtype {weight_dtype} is not supported in `W8A8PerTensorLinearMethod`.") |
| 83 | + | ||
| 83 | deq_scale = PerTensorScaleParameter(data=torch.empty(sum(output_partition_sizes), dtype=deq_scale_dtype)) | 84 | deq_scale = PerTensorScaleParameter(data=torch.empty(sum(output_partition_sizes), dtype=deq_scale_dtype)) |
| 84 | deq_scale.add_attrs({self.OUTPUT_DIM: 0, **extra_weight_attrs}) | 85 | deq_scale.add_attrs({self.OUTPUT_DIM: 0, **extra_weight_attrs}) |
| 85 | 86 | ||
| @@ -112,13 +113,15 @@ class W8A8PerTensorLinearMethod(LinearMethodBase): | |||
| 112 | # layer.input_scale.data: Scale factor for quantization (per-tensor) | 113 | # layer.input_scale.data: Scale factor for quantization (per-tensor) |
| 113 | # layer.input_offset.data: Zero-point offset for non-symmetric quantization | 114 | # layer.input_offset.data: Zero-point offset for non-symmetric quantization |
| 114 | # torch.qint8: Target quantization data type (8-bit signed integer) | 115 | # torch.qint8: Target quantization data type (8-bit signed integer) |
| 115 | - # -1: Quantization axis (whole tensor, no per-dimension quantization) | 116 | + # axis=-1: Quantize along the LAST dimension (last axis) of the input tensor |
| 116 | - # False: Non-symmetric quantization flag (uses offset; True would be symmetric) | 117 | + # div_mode=False: Use MULTIPLICATION (not division) for scale application in quantization. |
| 117 | input_tensor_quant = torch_npu.npu_quantize( | 118 | input_tensor_quant = torch_npu.npu_quantize( |
| 118 | - x, layer.input_scale.data, layer.input_offset.data, torch.qint8, -1, False) | 119 | + input=x, scales=layer.input_scale.data, |
| 120 | + zero_points=layer.input_offset.data, | ||
| 121 | + dtype=torch.qint8, axis=-1, div_mode=False) | ||
| 119 | out = torch_npu.npu_quant_matmul( | 122 | out = torch_npu.npu_quant_matmul( |
| 120 | input_tensor_quant, layer.weight.data, layer.deq_scale.data, | 123 | input_tensor_quant, layer.weight.data, layer.deq_scale.data, |
| 121 | - bias=layer.quant_bias.data, output_dtype=layer.weight_dtype) | 124 | + bias=layer.quant_bias.data, output_dtype=x.dtype) |
| 122 | return out | 125 | return out |
| 123 | 126 | ||
| 124 | def process_weights_after_loading(self, layer: nn.Module) -> None: | 127 | def process_weights_after_loading(self, layer: nn.Module) -> None: |
| @@ -202,12 +205,12 @@ class W8A8PerTokenLinearMethod(LinearMethodBase): | |||
| 202 | input_tensor_quant, pertoken_scale = torch_npu.npu_dynamic_quant(x) | 205 | input_tensor_quant, pertoken_scale = torch_npu.npu_dynamic_quant(x) |
| 203 | out = torch_npu.npu_quant_matmul( | 206 | out = torch_npu.npu_quant_matmul( |
| 204 | input_tensor_quant, layer.weight.data, layer.weight_scale.data, | 207 | input_tensor_quant, layer.weight.data, layer.weight_scale.data, |
| 205 | - pertoken_scale=pertoken_scale, bias=None, output_dtype=layer.weight_dtype) | 208 | + pertoken_scale=pertoken_scale, bias=None, output_dtype=x.dtype) |
| 206 | if layer.bias is not None: | 209 | if layer.bias is not None: |
| 207 | out = out + layer.bias.data | 210 | out = out + layer.bias.data |
| 208 | return out | 211 | return out |
| 209 | 212 | ||
| 210 | - def process_weights_after_loading(self, layer): | 213 | + def process_weights_after_loading(self, layer: nn.Module) -> None: |
| 211 | layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() | 214 | layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() |
| 212 | layer.weight_scale.data = layer.weight_scale.data.flatten() | 215 | layer.weight_scale.data = layer.weight_scale.data.flatten() |
| 213 | 216 | ||
| @@ -273,7 +276,7 @@ class W8A8MixLinearMethod(LinearMethodBase): | |||
| 273 | result = self.quant_method[InferenceMode.DECODE].apply(layer, x) | 276 | result = self.quant_method[InferenceMode.DECODE].apply(layer, x) |
| 274 | return result | 277 | return result |
| 275 | 278 | ||
| 276 | - def process_weights_after_loading(self, layer: nn.Module): | 279 | + def process_weights_after_loading(self, layer: nn.Module) -> None: |
| 277 | expanding_factor = layer.weight.data.shape[1] | 280 | expanding_factor = layer.weight.data.shape[1] |
| 278 | layer.input_scale.data = \ | 281 | layer.input_scale.data = \ |
| 279 | 1 / layer.input_scale.data.repeat(expanding_factor).to(layer.weight_dtype).contiguous().npu() | 282 | 1 / layer.input_scale.data.repeat(expanding_factor).to(layer.weight_dtype).contiguous().npu() |
| @@ -29,20 +29,24 @@ class QuantizationConfigBase(ABC): | |||
| 29 | 29 | ||
| 30 | 30 | ||
| 31 | def get_config_filenames() -> list[str]: | 31 | def get_config_filenames() -> list[str]: |
| 32 | + """List of quantization config filenames in the model directory.""" | ||
| 32 | raise NotImplementedError | 33 | raise NotImplementedError |
| 33 | 34 | ||
| 34 | 35 | ||
| 35 | 36 | ||
| 36 | def from_config(cls, config: dict[str, Any]) -> "QuantizationConfigBase": | 37 | def from_config(cls, config: dict[str, Any]) -> "QuantizationConfigBase": |
| 38 | + """Create quantization config instance from model's quantization config file.""" | ||
| 37 | raise NotImplementedError | 39 | raise NotImplementedError |
| 38 | 40 | ||
| 39 | 41 | ||
| 40 | def get_quant_method( | 42 | def get_quant_method( |
| 41 | self, layer: torch.nn.Module, prefix: str | 43 | self, layer: torch.nn.Module, prefix: str |
| 42 | ) -> QuantizationMethodBase | None: | 44 | ) -> QuantizationMethodBase | None: |
| 45 | + """Retrieve the appropriate quantization method for a given layer.""" | ||
| 43 | raise NotImplementedError | 46 | raise NotImplementedError |
| 44 | 47 | ||
| 45 | 48 | ||
| 46 | def get_quant_type_by_weight_name( | 49 | def get_quant_type_by_weight_name( |
| 47 | self, prefix: str | list[str], suffix: str) -> str: | 50 | self, prefix: str | list[str], suffix: str) -> str: |
| 51 | + """Retrieve the quantization type for a specific weight parameter.""" | ||
| 48 | raise NotImplementedError | 52 | raise NotImplementedError |
| @@ -19,27 +19,18 @@ from torch import nn | |||
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | class QuantizationMethodBase(ABC): | 21 | class QuantizationMethodBase(ABC): |
| 22 | - """Base class for different quantized methods.""" | 22 | + """Base class for quantized methods.""" |
| 23 | 23 | ||
| 24 | 24 | ||
| 25 | - def create_weights(self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs): | 25 | + def create_weights(self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs) -> None: |
| 26 | - """Create weights for a layer. | 26 | + """Create layer weights.""" |
| 27 | - | ||
| 28 | - The weights will be set as attributes of the layer. | ||
| 29 | - """ | ||
| 30 | raise NotImplementedError | 27 | raise NotImplementedError |
| 31 | 28 | ||
| 32 | 29 | ||
| 33 | def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: | 30 | def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: |
| 34 | - """Apply the weights in layer to the input tensor. | 31 | + """Apply layer weights to input tensor.""" |
| 35 | - | ||
| 36 | - Expects create_weights to have been called before on the layer. | ||
| 37 | - """ | ||
| 38 | raise NotImplementedError | 32 | raise NotImplementedError |
| 39 | 33 | ||
| 40 | def process_weights_after_loading(self, layer: nn.Module) -> None: | 34 | def process_weights_after_loading(self, layer: nn.Module) -> None: |
| 41 | - """Process the weight after loading. | 35 | + """Process weights after loading, e.g. transpose weights for computation.""" |
| 42 | - | ||
| 43 | - This can be used for example, to transpose weights for computation. | ||
| 44 | - """ | ||
| 45 | return | 36 | return |
| @@ -31,7 +31,7 @@ class UnquantizedLinearMethod(LinearMethodBase): | |||
| 31 | weight_dtype: torch.dtype, | 31 | weight_dtype: torch.dtype, |
| 32 | bias_dtype: torch.dtype, | 32 | bias_dtype: torch.dtype, |
| 33 | **extra_weight_attrs, | 33 | **extra_weight_attrs, |
| 34 | - ): | 34 | + ) -> None: |
| 35 | weight = ModelWeightParameter( | 35 | weight = ModelWeightParameter( |
| 36 | data=torch.empty( | 36 | data=torch.empty( |
| 37 | sum(output_partition_sizes), | 37 | sum(output_partition_sizes), |
| @@ -42,6 +42,9 @@ class UnquantizedLinearMethod(LinearMethodBase): | |||
| 42 | weight.add_attrs({"input_dim": 1, "output_dim": 0, **extra_weight_attrs}) | 42 | weight.add_attrs({"input_dim": 1, "output_dim": 0, **extra_weight_attrs}) |
| 43 | layer.register_parameter("weight", weight) | 43 | layer.register_parameter("weight", weight) |
| 44 | 44 | ||
| 45 | + # Determine if the anti-outlier feature is active by checking whether "norm.bias" parameters are present | ||
| 46 | + # in the weight files. If the bias tensor is added in the normalization module, it must be subtracted | ||
| 47 | + # from the subsequent linear module to ensure precision. | ||
| 45 | if layer.quant_config is not None: | 48 | if layer.quant_config is not None: |
| 46 | enable_anti_outlier = True | 49 | enable_anti_outlier = True |
| 47 | try: | 50 | try: |
| @@ -81,7 +84,7 @@ class UnquantizedEmbeddingMethod(QuantizationMethodBase): | |||
| 81 | output_size: int, | 84 | output_size: int, |
| 82 | params_dtype: torch.dtype, | 85 | params_dtype: torch.dtype, |
| 83 | **extra_weight_attrs, | 86 | **extra_weight_attrs, |
| 84 | - ): | 87 | + ) -> None: |
| 85 | weight = ModelWeightParameter( | 88 | weight = ModelWeightParameter( |
| 86 | torch.empty( | 89 | torch.empty( |
| 87 | input_size_per_partition, | 90 | input_size_per_partition, |
| @@ -107,7 +110,7 @@ class UnquantizedNormMethod(QuantizationMethodBase): | |||
| 107 | hidden_size: int, | 110 | hidden_size: int, |
| 108 | params_dtype: torch.dtype, | 111 | params_dtype: torch.dtype, |
| 109 | **extra_weight_attrs, | 112 | **extra_weight_attrs, |
| 110 | - ): | 113 | + ) -> None: |
| 111 | weight = BaseParameter(torch.ones(hidden_size, dtype=params_dtype)) | 114 | weight = BaseParameter(torch.ones(hidden_size, dtype=params_dtype)) |
| 112 | weight.add_attrs(extra_weight_attrs) | 115 | weight.add_attrs(extra_weight_attrs) |
| 113 | layer.register_parameter("weight", weight) | 116 | layer.register_parameter("weight", weight) |
| @@ -133,7 +136,7 @@ class UnquantizedLayerNormBiasMethod(QuantizationMethodBase): | |||
| 133 | hidden_size: int, | 136 | hidden_size: int, |
| 134 | params_dtype: torch.dtype, | 137 | params_dtype: torch.dtype, |
| 135 | **extra_weight_attrs, | 138 | **extra_weight_attrs, |
| 136 | - ): | 139 | + ) -> None: |
| 137 | weight = BaseParameter(torch.ones(hidden_size, dtype=params_dtype)) | 140 | weight = BaseParameter(torch.ones(hidden_size, dtype=params_dtype)) |
| 138 | weight.add_attrs(extra_weight_attrs) | 141 | weight.add_attrs(extra_weight_attrs) |
| 139 | layer.register_parameter("weight", weight) | 142 | layer.register_parameter("weight", weight) |
| @@ -147,5 +150,5 @@ class UnquantizedLayerNormBiasMethod(QuantizationMethodBase): | |||
| 147 | layer: torch.nn.Module, | 150 | layer: torch.nn.Module, |
| 148 | x: torch.Tensor, | 151 | x: torch.Tensor, |
| 149 | dim | 152 | dim |
| 150 | - ): | 153 | + ) -> torch.Tensor: |
| 151 | return torch.nn.functional.layer_norm(x, (dim,), layer.weight.data, layer.bias.data, layer.variance_epsilon) | 154 | return torch.nn.functional.layer_norm(x, (dim,), layer.weight.data, layer.bias.data, layer.variance_epsilon) |
| @@ -149,7 +149,7 @@ class ParallelInfoManager: | |||
| 149 | return list(range(num_layers)) | 149 | return list(range(num_layers)) |
| 150 | 150 | ||
| 151 | 151 | ||
| 152 | - def has_pp(): | 152 | + def has_pp() -> bool: |
| 153 | """Checks if pipeline parallelism is enabled (always False).""" | 153 | """Checks if pipeline parallelism is enabled (always False).""" |
| 154 | # (Note): depreciated, with change to parallelInfoManager.get(ParallelType.PP).is_enabled() | 154 | # (Note): depreciated, with change to parallelInfoManager.get(ParallelType.PP).is_enabled() |
| 155 | return False | 155 | return False |
| @@ -216,7 +216,7 @@ class ParallelInfoManager: | |||
| 216 | # (Note): depreciated, with change to parallelInfoManager.get(ParallelType.MOE_EP).is_enabled() | 216 | # (Note): depreciated, with change to parallelInfoManager.get(ParallelType.MOE_EP).is_enabled() |
| 217 | return self.get(ParallelType.MOE_EP).is_enabled() | 217 | return self.get(ParallelType.MOE_EP).is_enabled() |
| 218 | 218 | ||
| 219 | - def get(self, parallel_type: ParallelType): | 219 | + def get(self, parallel_type: ParallelType) -> ParallelInfo: |
| 220 | if parallel_type not in self._parallel_type_map: | 220 | if parallel_type not in self._parallel_type_map: |
| 221 | raise KeyError(f"Unsupported ParallelType: {parallel_type}") | 221 | raise KeyError(f"Unsupported ParallelType: {parallel_type}") |
| 222 | return self._parallel_type_map[parallel_type] | 222 | return self._parallel_type_map[parallel_type] |
| @@ -25,6 +25,7 @@ _BAR_FORMAT = "{desc}: {l_bar}{bar}| Completed | {n_fmt}/{total_fmt} [{elapsed}< | |||
| 25 | 25 | ||
| 26 | 26 | ||
| 27 | class DefaultModelLoader: | 27 | class DefaultModelLoader: |
| 28 | + """Model loader for safetensors checkpoint files.""" | ||
| 28 | def __init__(self): | 29 | def __init__(self): |
| 29 | self._counter_before_loading_weights: float = 0.0 | 30 | self._counter_before_loading_weights: float = 0.0 |
| 30 | self._counter_after_loading_weights: float = 0.0 | 31 | self._counter_after_loading_weights: float = 0.0 |
| @@ -32,6 +33,7 @@ class DefaultModelLoader: | |||
| 32 | self._weight_file_handler = None | 33 | self._weight_file_handler = None |
| 33 | 34 | ||
| 34 | def load_weights(self, model: nn.Module, model_path: str) -> None: | 35 | def load_weights(self, model: nn.Module, model_path: str) -> None: |
| 36 | + """Load model weights from checkpoint.""" | ||
| 35 | self._counter_before_loading_weights = time.perf_counter() | 37 | self._counter_before_loading_weights = time.perf_counter() |
| 36 | 38 | ||
| 37 | # Traverse module and map to corresponding weight | 39 | # Traverse module and map to corresponding weight |
| @@ -46,6 +48,7 @@ class DefaultModelLoader: | |||
| 46 | ) | 48 | ) |
| 47 | 49 | ||
| 48 | def _get_total_leaf_modules(self, module: nn.Module, prefix: str = "") -> dict[str, nn.Module]: | 50 | def _get_total_leaf_modules(self, module: nn.Module, prefix: str = "") -> dict[str, nn.Module]: |
| 51 | + """Get leaf modules with full names.""" | ||
| 49 | if len(list(module.children())) == 0: | 52 | if len(list(module.children())) == 0: |
| 50 | return {prefix: module} | 53 | return {prefix: module} |
| 51 | leaf_modules_dict = {} | 54 | leaf_modules_dict = {} |
| @@ -53,49 +56,52 @@ class DefaultModelLoader: | |||
| 53 | child_prefix = f"{prefix}.{name}" if prefix else name | 56 | child_prefix = f"{prefix}.{name}" if prefix else name |
| 54 | leaf_modules_dict.update(self._get_total_leaf_modules(child, child_prefix)) | 57 | leaf_modules_dict.update(self._get_total_leaf_modules(child, child_prefix)) |
| 55 | return leaf_modules_dict | 58 | return leaf_modules_dict |
| 59 | + | ||
| 60 | + def _load_multi_prefix_module(self, module: nn.Module) -> None: | ||
| 61 | + """Load weights for multi-prefix modules.(e.g., QKVParallelLinear).""" | ||
| 62 | + for shard_id, weight_prefix in enumerate(module.prefix): | ||
| 63 | + for weight_suffix, param in module.named_parameters(): | ||
| 64 | + if param is None: | ||
| 65 | + continue | ||
| 66 | + full_param_name = f"{weight_prefix}.{weight_suffix}" | ||
| 67 | + loaded_weight = self._weight_file_handler.get_tensor(full_param_name) | ||
| 68 | + param.weight_loader(param, loaded_weight, shard_id) | ||
| 69 | + | ||
| 70 | + def _load_single_prefix_module(self, module: nn.Module, prefix: str) -> None: | ||
| 71 | + """Load weights for single-prefix modules.""" | ||
| 72 | + for weight_suffix, param in module.named_parameters(): | ||
| 73 | + if param is None: | ||
| 74 | + continue | ||
| 75 | + full_param_name = f"{prefix}.{weight_suffix}" | ||
| 76 | + try: | ||
| 77 | + loaded_weight = self._weight_file_handler.get_tensor(full_param_name) | ||
| 78 | + except ValueError as e: | ||
| 79 | + # Try module-specific prefix for weight files(e.g., used for tie_word_embedding function) | ||
| 80 | + if "Weight file was not found" in str(e) and hasattr(module, "prefix"): | ||
| 81 | + full_param_name = f"{module.prefix}.{weight_suffix}" | ||
| 82 | + loaded_weight = self._weight_file_handler.get_tensor(full_param_name) | ||
| 83 | + param.weight_loader(param, loaded_weight) | ||
| 84 | + if isinstance(module, FusedMoE): | ||
| 85 | + module.weight_loader(loaded_weight, full_param_name) | ||
| 56 | 86 | ||
| 57 | - def _load_modules_with_progress(self, modules_dict: dict, pbar: tqdm,): | 87 | + def _load_modules_with_progress(self, modules_dict: dict, pbar: tqdm,) -> None: |
| 88 | + """Load weights for modules with progress.""" | ||
| 58 | for prefix, module in modules_dict.items(): | 89 | for prefix, module in modules_dict.items(): |
| 90 | + # Handling multi-prefix modules | ||
| 59 | if hasattr(module, "prefix") and isinstance(module.prefix, list): | 91 | if hasattr(module, "prefix") and isinstance(module.prefix, list): |
| 60 | - shard_id = 0 | 92 | + self._load_multi_prefix_module(module) |
| 61 | - for weight_prefix in module.prefix: | 93 | + else: # Processing single prefix Module |
| 62 | - for weight_suffix, param in module.named_parameters(): | 94 | + self._load_single_prefix_module(module, prefix) |
| 63 | - full_param_name = f"{weight_prefix}.{weight_suffix}" | ||
| 64 | - loaded_weight = self._weight_file_handler.get_tensor(full_param_name) | ||
| 65 | 95 | ||
| 66 | - if param is None: | 96 | + # Apply post-processing after pretrained-weights are loaded onto NPU device. |
| 67 | - continue | ||
| 68 | - | ||
| 69 | - param.weight_loader(param, loaded_weight, shard_id) | ||
| 70 | - | ||
| 71 | - shard_id += 1 | ||
| 72 | - else: | ||
| 73 | - for weight_suffix, param in module.named_parameters(): | ||
| 74 | - full_param_name = f"{prefix}.{weight_suffix}" | ||
| 75 | - try: | ||
| 76 | - loaded_weight = self._weight_file_handler.get_tensor(full_param_name) | ||
| 77 | - except ValueError as e: | ||
| 78 | - if "Weight file was not found" in str(e) and hasattr(module, "prefix"): | ||
| 79 | - loaded_weight = self._weight_file_handler.get_tensor(f"{module.prefix}.{weight_suffix}") | ||
| 80 | - if param is None: | ||
| 81 | - continue | ||
| 82 | - | ||
| 83 | - param.weight_loader(param, loaded_weight) | ||
| 84 | - | ||
| 85 | - if isinstance(module, FusedMoE): | ||
| 86 | - module.weight_loader(loaded_weight, full_param_name) | ||
| 87 | - | ||
| 88 | - # Process weights after loading weights | ||
| 89 | quant_method = getattr(module, "quant_method", None) | 97 | quant_method = getattr(module, "quant_method", None) |
| 90 | if isinstance(quant_method, QuantizationMethodBase): | 98 | if isinstance(quant_method, QuantizationMethodBase): |
| 91 | quant_method.process_weights_after_loading(module) | 99 | quant_method.process_weights_after_loading(module) |
| 92 | - | 100 | + |
| 93 | pbar.update(1) | 101 | pbar.update(1) |
| 94 | 102 | ||
| 95 | - def _load_modules(self, model: nn.Module): | 103 | + def _load_modules(self, model: nn.Module) -> None: |
| 96 | - """Traverse module and load from corresponding safetensors checkpoint. | 104 | + """Load model weights for leaf modules.""" |
| 97 | - return not loaded module names in model | ||
| 98 | - """ | ||
| 99 | leaf_modules_dict = self._get_total_leaf_modules(model) | 105 | leaf_modules_dict = self._get_total_leaf_modules(model) |
| 100 | 106 | ||
| 101 | pbar = tqdm(total=len(leaf_modules_dict), | 107 | pbar = tqdm(total=len(leaf_modules_dict), |
| @@ -10,7 +10,7 @@ | |||
| 10 | 10 | ||
| 11 | import os | 11 | import os |
| 12 | import json | 12 | import json |
| 13 | -from typing import List, Generator, Any, Tuple | 13 | +from typing import List, Any, Tuple |
| 14 | from pathlib import Path | 14 | from pathlib import Path |
| 15 | 15 | ||
| 16 | import torch | 16 | import torch |
| @@ -65,18 +65,21 @@ class WeightsFileHandler: | |||
| 65 | raise FileNotFoundError("The input model id is not exists or not a directory") | 65 | raise FileNotFoundError("The input model id is not exists or not a directory") |
| 66 | 66 | ||
| 67 | def release_file_handler(self) -> None: | 67 | def release_file_handler(self) -> None: |
| 68 | + """Release all file handlers""" | ||
| 68 | if self._handlers: | 69 | if self._handlers: |
| 69 | del self._handlers | 70 | del self._handlers |
| 70 | self._handlers = {} | 71 | self._handlers = {} |
| 71 | 72 | ||
| 72 | 73 | ||
| 73 | def get_tensor(self, tensor_name: str) -> Any: | 74 | def get_tensor(self, tensor_name: str) -> Any: |
| 75 | + """Get tensor by full name.""" | ||
| 74 | filename, tensor_name = self._get_filename(tensor_name) | 76 | filename, tensor_name = self._get_filename(tensor_name) |
| 75 | f = self._get_handler(filename) | 77 | f = self._get_handler(filename) |
| 76 | tensor = f.get_tensor(tensor_name) | 78 | tensor = f.get_tensor(tensor_name) |
| 77 | return tensor | 79 | return tensor |
| 78 | 80 | ||
| 79 | def _get_handler(self, filename: str) -> Any: | 81 | def _get_handler(self, filename: str) -> Any: |
| 82 | + """Get file handler by filename.""" | ||
| 80 | if filename not in self._handlers: | 83 | if filename not in self._handlers: |
| 81 | # Note: manually call the release_file_handler method after use. | 84 | # Note: manually call the release_file_handler method after use. |
| 82 | f = safetensors.safe_open(filename, framework="pytorch") | 85 | f = safetensors.safe_open(filename, framework="pytorch") |
| @@ -85,12 +88,14 @@ class WeightsFileHandler: | |||
| 85 | return self._handlers[filename] | 88 | return self._handlers[filename] |
| 86 | 89 | ||
| 87 | def _get_filename(self, tensor_name: str) -> Tuple[str, str]: | 90 | def _get_filename(self, tensor_name: str) -> Tuple[str, str]: |
| 91 | + """Get file name for tensor name.""" | ||
| 88 | filename = self._routing.get(tensor_name) | 92 | filename = self._routing.get(tensor_name) |
| 89 | if filename is None: | 93 | if filename is None: |
| 90 | raise ValueError(f"Weight file was not found for tensor named with {tensor_name}.") | 94 | raise ValueError(f"Weight file was not found for tensor named with {tensor_name}.") |
| 91 | return str(filename), tensor_name | 95 | return str(filename), tensor_name |
| 92 | 96 | ||
| 93 | def _load_weight_file_routing(self) -> dict: | 97 | def _load_weight_file_routing(self) -> dict: |
| 98 | + """Build routing of weight files.""" | ||
| 94 | routing = {} | 99 | routing = {} |
| 95 | for filename in self._filenames: | 100 | for filename in self._filenames: |
| 96 | filename = standardize_path(str(filename), check_link=False) | 101 | filename = standardize_path(str(filename), check_link=False) |
| @@ -9,7 +9,6 @@ | |||
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | 10 | ||
| 11 | from contextlib import contextmanager | 11 | from contextlib import contextmanager |
| 12 | -from typing import Generator | ||
| 13 | import torch | 12 | import torch |
| 14 | 13 | ||
| 15 | 14 | ||
| @@ -131,7 +131,7 @@ class TestW8A8PerTensorLinearMethod(unittest.TestCase): | |||
| 131 | layer.deq_scale.data = torch.randn(1024).to(torch.int64) | 131 | layer.deq_scale.data = torch.randn(1024).to(torch.int64) |
| 132 | layer.quant_bias.data = torch.randn(1024).to(torch.int32) | 132 | layer.quant_bias.data = torch.randn(1024).to(torch.int32) |
| 133 | 133 | ||
| 134 | - x = torch.randn(2, 3, 512) | 134 | + x = torch.randn(2, 3, 512, dtype=torch.float16) |
| 135 | 135 | ||
| 136 | # Mock npu functions | 136 | # Mock npu functions |
| 137 | mock_quantized_tensor = torch.randn(2, 3, 512) | 137 | mock_quantized_tensor = torch.randn(2, 3, 512) |
| @@ -143,10 +143,10 @@ class TestW8A8PerTensorLinearMethod(unittest.TestCase): | |||
| 143 | # Verify npu_quantize was called | 143 | # Verify npu_quantize was called |
| 144 | mock_npu_quantize.assert_called_once() | 144 | mock_npu_quantize.assert_called_once() |
| 145 | call_args = mock_npu_quantize.call_args | 145 | call_args = mock_npu_quantize.call_args |
| 146 | - self.assertIs(call_args[0][0], x) | 146 | + self.assertIs(call_args.kwargs['input'], x) |
| 147 | - self.assertTrue(torch.equal(call_args[0][1], layer.input_scale.data)) | 147 | + self.assertTrue(torch.equal(call_args.kwargs['scales'], layer.input_scale.data)) |
| 148 | - self.assertTrue(torch.equal(call_args[0][2], layer.input_offset.data)) | 148 | + self.assertTrue(torch.equal(call_args.kwargs['zero_points'], layer.input_offset.data)) |
| 149 | - self.assertEqual(call_args[0][3], torch.qint8) | 149 | + self.assertEqual(call_args.kwargs['dtype'], torch.qint8) |
| 150 | 150 | ||
| 151 | # Verify npu_quant_matmul was called | 151 | # Verify npu_quant_matmul was called |
| 152 | mock_npu_quant_matmul.assert_called_once() | 152 | mock_npu_quant_matmul.assert_called_once() |
| @@ -263,7 +263,7 @@ class TestW8A8PerTokenLinearMethod(unittest.TestCase): | |||
| 263 | layer.weight_scale.data = torch.randn(1024, 1, dtype=torch.float32) | 263 | layer.weight_scale.data = torch.randn(1024, 1, dtype=torch.float32) |
| 264 | layer.weight_offset.data = torch.randn(1024, 1, dtype=torch.float16) | 264 | layer.weight_offset.data = torch.randn(1024, 1, dtype=torch.float16) |
| 265 | 265 | ||
| 266 | - x = torch.randn(2, 3, 512) | 266 | + x = torch.randn(2, 3, 512, dtype=torch.float16) |
| 267 | 267 | ||
| 268 | # Mock npu functions | 268 | # Mock npu functions |
| 269 | mock_quantized_tensor = torch.randn(2, 3, 512) | 269 | mock_quantized_tensor = torch.randn(2, 3, 512) |