# Copyright (c) ModelScope Contributors. All rights reserved. import torch from transformers import PretrainedConfig from typing import Any, Dict, List, Optional, Tuple, Union from .utils import deep_getattr class HfConfigFactory: llm_keys = ['language_config', 'llm_config', 'text_config'] vision_keys = ['vit_config', 'vision_config', 'audio_config'] """This class is used to read config from config.json(maybe params.json also)""" @staticmethod def get_torch_dtype(config: Union[PretrainedConfig, Dict[str, Any]], quant_info: Dict[str, Any]) -> Optional[torch.dtype]: for key in ['torch_dtype', 'params_dtype']: torch_dtype = HfConfigFactory.get_config_attr(config, key) if torch_dtype is not None: break torch_dtype = HfConfigFactory.to_torch_dtype(torch_dtype) if torch_dtype is None: torch_dtype = quant_info.get('torch_dtype') return torch_dtype @staticmethod def get_text_config(config): for key in HfConfigFactory.llm_keys: value = getattr(config, key, None) if value is not None: return value return config @staticmethod def _get_config_attrs(config: Union[PretrainedConfig, Dict[str, Any]], attr_name: str, include_vit: bool = False, parent_key: Optional[str] = None) -> List[Tuple[PretrainedConfig, Any]]: res = [] if isinstance(config, dict): keys = config.keys() elif isinstance(config, PretrainedConfig): keys = dir(config) else: return [] config_keys = [None] + HfConfigFactory.llm_keys if include_vit: config_keys += HfConfigFactory.vision_keys if attr_name in keys and parent_key in config_keys: res.append((config, deep_getattr(config, attr_name))) for k in keys: if k.endswith('_config') and k != 'talker_config': if isinstance(config, dict): v = config[k] else: v = getattr(config, k) res += HfConfigFactory._get_config_attrs(v, attr_name, include_vit, k) return res @staticmethod def is_moe_model(config) -> bool: if 'Moe' in config.__class__.__name__: return True for key in ['num_experts', 'num_experts_per_tok', 'moe_intermediate_size']: if HfConfigFactory.get_config_attr(config, key): return True return False @staticmethod def is_multimodal(config) -> bool: if isinstance(config, dict): keys = config.keys() elif isinstance(config, PretrainedConfig): keys = dir(config) else: keys = [] keys = set(keys) for key in (HfConfigFactory.llm_keys + HfConfigFactory.vision_keys + ['thinker_config']): if key in keys: return True return False @staticmethod def get_config_attr(config: Union[PretrainedConfig, Dict[str, Any]], attr_name: str, include_vit: bool = False) -> Optional[Any]: """Get the value of the attribute named attr_name.""" attrs = HfConfigFactory._get_config_attrs(config, attr_name, include_vit) if len(attrs) == 0: return None else: return attrs[0][1] @staticmethod def set_config_attr(config: Union[PretrainedConfig, Dict[str, Any]], attr_name: str, value: Any, include_vit: bool = False, ensure_set: bool = True) -> int: """Set all the attr_name attributes to value.""" attrs = HfConfigFactory._get_config_attrs(config, attr_name, include_vit) if ensure_set and len(attrs) == 0: attrs.append((config, None)) for config, _ in attrs: if isinstance(config, dict): config[attr_name] = value else: setattr(config, attr_name, value) return len(attrs) @staticmethod def del_config_attr(config: Union[PretrainedConfig, Dict[str, Any]], attr_name: str, include_vit: bool = False) -> int: """Remove all the attr_name attributes.""" attrs = HfConfigFactory._get_config_attrs(config, attr_name, include_vit) for config, _ in attrs: if isinstance(config, dict): config.pop(attr_name, None) elif hasattr(config, attr_name): delattr(config, attr_name) return len(attrs) @staticmethod def get_max_model_len(config: Union[PretrainedConfig, Dict[str, Any]]) -> Optional[int]: """Get the max length supported by the model""" INF = int(1e9) max_model_len = INF possible_keys = [ 'seq_length', # qwen, chatglm 'max_position_embeddings', # qwen1.5, llama2 'n_positions', # polylm, phi-2 'model_max_length', # baichuan2 # others 'seq_len', 'max_seq_len', 'max_sequence_length', 'max_seq_length', ] for key in possible_keys: max_len_key = HfConfigFactory.get_config_attr(config, key) if max_len_key is not None: max_model_len = min(max_model_len, max_len_key) if max_model_len == INF: max_model_len = None return max_model_len @staticmethod def set_max_model_len(config: Union[PretrainedConfig, Dict[str, Any]], value: int): """Set the max length supported by the model""" possible_keys = [ 'seq_length', # qwen, chatglm 'max_position_embeddings', # qwen1.5, llama2 'n_positions', # polylm, phi-2 'model_max_length', # baichuan2 # others 'seq_len', 'max_seq_len', 'max_sequence_length', 'max_seq_length', ] for key in possible_keys: max_len_value = HfConfigFactory.get_config_attr(config, key) if max_len_value is not None: HfConfigFactory.set_config_attr(config, key, value) @staticmethod def compat_zero3(config: PretrainedConfig) -> None: value = HfConfigFactory.get_config_attr(config, 'hidden_size') try: # AttributeError: can't set attribute 'hidden_size' config.hidden_size = value except AttributeError: pass @staticmethod def to_torch_dtype(torch_dtype: Union[str, torch.dtype, None]) -> Optional[torch.dtype]: if torch_dtype is None: return None if isinstance(torch_dtype, str): torch_dtype = torch_dtype.replace('torch.', '') torch_dtype = getattr(torch, torch_dtype) return torch_dtype @staticmethod def get_quant_info(config: Union[PretrainedConfig, Dict[str, Any]]) -> Optional[Dict[str, Any]]: """Get quant_method, quant_bits, dtype. not support hqq/eetq now, support awq/gptq/bnb/aqlm""" if isinstance(config, dict): quantization_config = config.get('quantization_config') else: quantization_config = getattr(config, 'quantization_config', None) if quantization_config is None: return quantization_config = dict(quantization_config) quant_method = quantization_config.get('quant_method') res = {} if quant_method in {'gptq', 'awq', 'aqlm'}: res['quant_method'] = quant_method res['torch_dtype'] = torch.float16 quant_bits = quantization_config.get('bits') if quant_bits is not None: res['quant_bits'] = quant_bits elif quant_method == 'bitsandbytes': res['quant_method'] = 'bnb' load_in_4bit = quantization_config.get('_load_in_4bit') load_in_8bit = quantization_config.get('_load_in_8bit') bnb_4bit_compute_dtype = quantization_config.get('bnb_4bit_compute_dtype') if load_in_4bit: res['quant_bits'] = 4 elif load_in_8bit: res['quant_bits'] = 8 res['torch_dtype'] = HfConfigFactory.to_torch_dtype(bnb_4bit_compute_dtype) elif quant_method == 'hqq': res['quant_method'] = quant_method res['quant_bits'] = quantization_config['quant_config']['weight_quant_params']['nbits'] elif quant_method is not None: res['quant_method'] = quant_method return res or None