diff --git a/rtp_llm/test/hf_model_helper_test.py b/rtp_llm/test/hf_model_helper_test.py new file mode 100644 index 0000000000..3bf5bdb366 --- /dev/null +++ b/rtp_llm/test/hf_model_helper_test.py @@ -0,0 +1,106 @@ +import sys +import importlib.util +import unittest +from unittest.mock import MagicMock + +# Load model_factory_register directly to avoid __init__.py's .so dependency +_spec = importlib.util.spec_from_file_location( + "rtp_llm.model_factory_register", "rtp_llm/model_factory_register.py" +) +_mfr = importlib.util.module_from_spec(_spec) +sys.modules["rtp_llm.model_factory_register"] = _mfr +_spec.loader.exec_module(_mfr) + +# Register architectures to simulate production state +_mfr.register_model("deepseek2", object, ["DeepseekV2ForCausalLM"]) +_mfr.register_model("deepseek3", object, ["DeepseekV3ForCausalLM"]) +_mfr.register_model("qwen35_moe", object, ["Qwen3_5MoeForConditionalGeneration"]) +_mfr.register_model("qwen35_dense", object, ["Qwen3_5ForConditionalGeneration"]) + +# Mock huggingface_hub before importing hf_model_helper +sys.modules["huggingface_hub"] = MagicMock() +sys.modules["huggingface_hub.hf_api"] = MagicMock() + +_spec2 = importlib.util.spec_from_file_location( + "rtp_llm.tools.api.hf_model_helper", "rtp_llm/tools/api/hf_model_helper.py" +) +_hfh = importlib.util.module_from_spec(_spec2) +_spec2.loader.exec_module(_hfh) + +HfStyleModelInfo = _hfh.HfStyleModelInfo + + +class TestResolveFtModelType(unittest.TestCase): + def test_direct_match(self): + config = {"architectures": ["DeepseekV3ForCausalLM"]} + self.assertEqual(HfStyleModelInfo.resolve_ft_model_type(config), "deepseek3") + + def test_fuzzy_match_causal_to_conditional(self): + config = {"architectures": ["Qwen3_5MoeForCausalLM"]} + self.assertEqual(HfStyleModelInfo.resolve_ft_model_type(config), "qwen35_moe") + + def test_fuzzy_match_conditional_to_causal(self): + _mfr.register_model("test_causal", object, ["TestModelForCausalLM"]) + config = {"architectures": ["TestModelForConditionalGeneration"]} + self.assertEqual(HfStyleModelInfo.resolve_ft_model_type(config), "test_causal") + + def test_model_type_fallback_deepseek_v4(self): + config = {"architectures": ["DeepseekV4ForCausalLM"], "model_type": "deepseek_v4"} + self.assertEqual(HfStyleModelInfo.resolve_ft_model_type(config), "deepseek3") + + def test_model_type_fallback_qwen3_5(self): + config = {"architectures": ["UnknownArch"], "model_type": "qwen3_5"} + self.assertEqual(HfStyleModelInfo.resolve_ft_model_type(config), "qwen35_dense") + + def test_all_miss_returns_none(self): + config = {"architectures": ["CompletelyUnknownArch"], "model_type": "unknown"} + self.assertIsNone(HfStyleModelInfo.resolve_ft_model_type(config)) + + def test_empty_architectures(self): + config = {"architectures": [], "model_type": "deepseek_v3"} + self.assertEqual(HfStyleModelInfo.resolve_ft_model_type(config), "deepseek3") + + def test_no_architectures_key(self): + config = {"model_type": "qwen3_5_moe"} + self.assertEqual(HfStyleModelInfo.resolve_ft_model_type(config), "qwen35_moe") + + +class TestFuzzyMatchArchitecture(unittest.TestCase): + def test_empty_architectures_returns_none(self): + self.assertIsNone(HfStyleModelInfo._fuzzy_match_architecture({"architectures": []})) + + def test_no_architectures_key_returns_none(self): + self.assertIsNone(HfStyleModelInfo._fuzzy_match_architecture({})) + + def test_no_suffix_match_returns_none(self): + config = {"architectures": ["SomeRandomModel"]} + self.assertIsNone(HfStyleModelInfo._fuzzy_match_architecture(config)) + + def test_does_not_match_nextn_suffix(self): + _mfr.register_model("mtp", object, ["DeepseekV3ForCausalLMNextN"]) + config = {"architectures": ["DeepseekV3ForCausalLMNextN"]} + self.assertIsNone(HfStyleModelInfo._fuzzy_match_architecture(config)) + + +class TestDtypeBytes(unittest.TestCase): + def test_f32_is_4_bytes(self): + self.assertEqual(HfStyleModelInfo._DTYPE_BYTES["F32"], 4) + + def test_bf16_is_2_bytes(self): + self.assertEqual(HfStyleModelInfo._DTYPE_BYTES["BF16"], 2) + + def test_i8_is_1_byte(self): + self.assertEqual(HfStyleModelInfo._DTYPE_BYTES["I8"], 1) + + def test_f8_e4m3_is_1_byte(self): + self.assertEqual(HfStyleModelInfo._DTYPE_BYTES["F8_E4M3"], 1) + + def test_unknown_dtype_defaults_to_2(self): + self.assertEqual(HfStyleModelInfo._DTYPE_BYTES.get("UNKNOWN_TYPE", 2), 2) + + def test_bool_is_1_byte(self): + self.assertEqual(HfStyleModelInfo._DTYPE_BYTES["BOOL"], 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/rtp_llm/tools/api/hf_model_helper.py b/rtp_llm/tools/api/hf_model_helper.py index e6fe24c57f..735f0f5d27 100644 --- a/rtp_llm/tools/api/hf_model_helper.py +++ b/rtp_llm/tools/api/hf_model_helper.py @@ -110,6 +110,14 @@ def _get_auto_config_py(self, repo_or_link: str, revision: Optional[str]): ) return None + _DTYPE_BYTES = { + "F64": 8, "I64": 8, + "F32": 4, "U32": 4, "I32": 4, + "F16": 2, "BF16": 2, "U16": 2, "I16": 2, + "F8_E4M3": 1, "F8_E5M2": 1, "F8_E8M0": 1, + "I8": 1, "U8": 1, "BOOL": 1, + } + def _calculate_model_parameters(self) -> Tuple[Optional[int], Optional[int]]: param_count = None total_size = None @@ -117,11 +125,7 @@ def _calculate_model_parameters(self) -> Tuple[Optional[int], Optional[int]]: if self.model_info and self.model_info.safetensors: param_count = self.model_info.safetensors.total total_size = sum( - ( - count * 2 - if weight_type in ["FP16", "BF16", "FP32", "FP32", "INT8", "F16"] - else count - ) + count * self._DTYPE_BYTES.get(weight_type, 2) for weight_type, count in self.model_info.safetensors.parameters.items() ) elif self.meta_info_file and os.path.exists(self.meta_info_file): @@ -145,14 +149,52 @@ def _calculate_model_parameters(self) -> Tuple[Optional[int], Optional[int]]: ) return param_count, total_size + _MODEL_TYPE_MAP = { + "deepseek_v2": "deepseek2", + "deepseek_v3": "deepseek3", + "deepseek_v4": "deepseek3", + "qwen3_5_moe": "qwen35_moe", + "qwen3_5": "qwen35_dense", + } + + @staticmethod + def _fuzzy_match_architecture(config: dict) -> Optional[str]: + architectures = config.get("architectures", []) + if not architectures: + return None + architecture = architectures[0] + _SUFFIX_SWAPS = [ + ("ForCausalLM", "ForConditionalGeneration"), + ("ForConditionalGeneration", "ForCausalLM"), + ] + for old_suffix, new_suffix in _SUFFIX_SWAPS: + if architecture.endswith(old_suffix): + alt = architecture[: -len(old_suffix)] + new_suffix + ft_type = ModelDict.get_ft_model_type_by_hf_architectures(alt) + if ft_type: + return ft_type + return None + + @classmethod + def resolve_ft_model_type(cls, config: dict) -> Optional[str]: + ft_type = ModelDict.get_ft_model_type_by_config(config) + if ft_type: + return ft_type + ft_type = cls._fuzzy_match_architecture(config) + if ft_type: + return ft_type + hf_model_type = config.get("model_type", "") + if hf_model_type: + return cls._MODEL_TYPE_MAP.get(hf_model_type) + return None + @property def ft_model_type(self) -> Optional[str]: if self.model_info: - # Assume ModelDict.get_ft_model_type_by_hf_repo() is a valid method ft_type = ModelDict.get_ft_model_type_by_hf_repo(self.model_info.modelId) if ft_type is not None: return ft_type - return ModelDict.get_ft_model_type_by_config(self.model_config) + return self.resolve_ft_model_type(self.model_config) @staticmethod def is_from_hf(model_path: str) -> bool: diff --git a/rtp_llm/tools/api/model_basic_info_analyzer.py b/rtp_llm/tools/api/model_basic_info_analyzer.py index 6f8b53a4be..9fdb1720de 100644 --- a/rtp_llm/tools/api/model_basic_info_analyzer.py +++ b/rtp_llm/tools/api/model_basic_info_analyzer.py @@ -28,7 +28,7 @@ def _parse_hf_model_type(model_link): def parse_ft_model_type(model_path): # load config.json config = _get_raw_config(model_path) - ft_model_type = ModelDict.get_ft_model_type_by_config(config) + ft_model_type = HfStyleModelInfo.resolve_ft_model_type(config) return {"ft_model_type": ft_model_type} @@ -116,9 +116,12 @@ def _load_as_hf_style(model_path, ft_model_type, env_params) -> ModelBasicInfo: # config = PretrainedConfig.from_dict(config_dict) logging.info(f"config:{config}") hidden_size = config.hidden_size if hasattr(config, "hidden_size") else None - # num_hidden_layers = config.num_hidden_layers - # num_attention_heads = config.num_attention_heads - # vocab_size = config.vocab_size + if hidden_size is None and hasattr(config, "text_config"): + text_cfg = config.text_config + if isinstance(text_cfg, dict): + hidden_size = text_cfg.get("hidden_size") + elif hasattr(text_cfg, "hidden_size"): + hidden_size = text_cfg.hidden_size quant_config = None is_quant_weight = False if hasattr(config, "quantization_config"):