model factory

Author: shihaobai

lightllm/models/__init__.py

@@ -0,0 +1,28 @@
+import os
+import importlib
+import inspect
+from pathlib import Path
+
+
+def auto_import_models():
+    """
+    Automatically imports all classes from model.py files in model directories
+    """
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    models_dir = Path(base_dir)
+    for model_dir in models_dir.iterdir():
+        if not model_dir.is_dir():
+            continue
+        model_file = model_dir / "model.py"
+        if not model_file.exists():
+            continue
+        module_path = f"lightllm.models.{model_dir.name}.model"
+
+        try:
+            importlib.import_module(module_path)
+        except:
+            pass
+
+
+auto_import_models()
+from .registry import get_model

lightllm/models/bloom/model.py

@@ -1,6 +1,7 @@
 import os
 import json
 import torch
+from lightllm.models.registry import ModelRegistry
 from lightllm.models.bloom.layer_infer.pre_layer_infer import BloomPreLayerInfer
 from lightllm.models.bloom.layer_infer.post_layer_infer import BloomPostLayerInfer
 from lightllm.models.bloom.layer_infer.transformer_layer_infer import BloomTransformerLayerInfer
@@ -12,6 +13,7 @@
 from lightllm.common.build_utils import repair_config
 
 
+@ModelRegistry("bloom")
 class BloomTpPartModel(TpPartBaseModel):
     # weight class
     pre_and_post_weight_class = BloomPreAndPostLayerWeight

lightllm/models/chatglm2/model.py

@@ -2,6 +2,7 @@
 import json
 import torch
 
+from lightllm.models.registry import ModelRegistry
 from lightllm.models.chatglm2.layer_infer.transformer_layer_infer import ChatGLM2TransformerLayerInfer
 from lightllm.models.chatglm2.layer_weights.transformer_layer_weight import ChatGLM2TransformerLayerWeight
 from lightllm.models.chatglm2.layer_weights.pre_and_post_layer_weight import ChatGLM2PreAndPostLayerWeight
@@ -12,6 +13,7 @@
 logger = init_logger(__name__)
 
 
+@ModelRegistry("chatglm")
 class ChatGlm2TpPartModel(LlamaTpPartModel):
     # Please use the fast tokenizer from:
     # [THUDM/chatglm3-6b PR #12](https://huggingface.co/THUDM/chatglm3-6b/discussions/12).

lightllm/models/cohere/model.py

@@ -5,6 +5,7 @@
     TransformerLayerCohereInferTpl,
 )
 from lightllm.common.mem_manager import MemoryManager
+from lightllm.models.registry import ModelRegistry
 from lightllm.models.cohere.infer_struct import CohereInferStateInfo
 from lightllm.models.cohere.layer_infer.post_layer_infer import CoherePostLayerInfer
 from lightllm.models.cohere.layer_infer.transformer_layer_infer import CohereTransformerLayerInfer
@@ -17,6 +18,7 @@
 logger = init_logger(__name__)
 
 
+@ModelRegistry("cohere")
 class CohereTpPartModel(LlamaTpPartModel):
     pre_and_post_weight_class = CoherePreAndPostLayerWeight
     transformer_weight_class = CohereTransformerLayerWeight

lightllm/models/deepseek2/model.py

@@ -1,5 +1,6 @@
 import torch
 from typing import final
+from lightllm.models.registry import ModelRegistry
 from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer
 from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
@@ -22,6 +23,7 @@
 
 class DeepSeek2FlashInferStateExtraInfo:
     def __init__(self, model):
+        print(model)
         num_heads = model.config["num_attention_heads"]
         self.tp_q_head_num = num_heads // get_dp_world_size()
         self.qk_nope_head_dim = model.qk_nope_head_dim
@@ -49,6 +51,7 @@ def __init__(self, model):
                 self.softmax_scale = self.softmax_scale * mscale * mscale
 
 
+@ModelRegistry(["deepseek_v2", "deepseek_v3"])
 class Deepseek2TpPartModel(LlamaTpPartModel):
     # weight class
     transformer_weight_class = Deepseek2TransformerLayerWeight

lightllm/models/gemma3/model.py

@@ -3,6 +3,7 @@
 import json
 import numpy as np
 import torch
+from lightllm.models.registry import ModelRegistry
 from lightllm.common.basemodel.multimodal_tokenizer import BaseMultiModalTokenizer
 from lightllm.common.mem_utils import select_mem_manager_class
 from lightllm.models.gemma3.infer_struct import Gemma3InferStateInfo
@@ -22,6 +23,7 @@
 
 logger = init_logger(__name__)
 
+
 # Warp of the origal tokenizer
 class Gemma3Tokenizer(BaseMultiModalTokenizer):
     def __init__(self, tokenizer, model_cfg):
@@ -77,6 +79,7 @@ def encode(self, prompt, multimodal_params: MultimodalParams = None, add_special
         return input_ids
 
 
+@ModelRegistry("gemma3")
 class Gemma3TpPartModel(LlamaTpPartModel):
     # weight class
     pre_and_post_weight_class = Gemma3PreAndPostLayerWeight

lightllm/models/gemma_2b/model.py

@@ -1,3 +1,4 @@
+from lightllm.models.registry import ModelRegistry
 from lightllm.models.gemma_2b.layer_weights.transformer_layer_weight import Gemma_2bTransformerLayerWeight
 from lightllm.models.gemma_2b.layer_weights.pre_and_post_layer_weight import Gemma_2bPreAndPostLayerWeight
 from lightllm.models.gemma_2b.layer_infer.pre_layer_infer import Gemma_2bPreLayerInfer
@@ -8,6 +9,7 @@
 from lightllm.common.mem_utils import select_mem_manager_class
 
 
+@ModelRegistry("gemma")
 class Gemma_2bTpPartModel(LlamaTpPartModel):
     # weight class
     pre_and_post_weight_class = Gemma_2bPreAndPostLayerWeight

lightllm/models/internlm/model.py

@@ -1,11 +1,12 @@
 import os
 import json
 import torch
-
+from lightllm.models.registry import ModelRegistry
 from lightllm.models.internlm.layer_weights.transformer_layer_weight import InternlmTransformerLayerWeight
 from lightllm.models.llama.model import LlamaTpPartModel
 
 
+@ModelRegistry("internlm")
 class InternlmTpPartModel(LlamaTpPartModel):
     # weight class
     transformer_weight_class = InternlmTransformerLayerWeight

lightllm/models/internlm2/model.py

@@ -2,16 +2,17 @@
 import json
 import torch
 
+from lightllm.models.registry import ModelRegistry
 from lightllm.models.internlm2.layer_weights.transformer_layer_weight import Internlm2TransformerLayerWeight
-from lightllm.models.internlm2.layer_weights.pre_and_post_layer_weight import Internlm2PreAndPostLayerWeight 
+from lightllm.models.internlm2.layer_weights.pre_and_post_layer_weight import Internlm2PreAndPostLayerWeight
 from lightllm.models.internlm.model import InternlmTpPartModel
 
 
+@ModelRegistry("internlm2")
 class Internlm2TpPartModel(InternlmTpPartModel):
     # weight class
-    pre_and_post_weight_class = Internlm2PreAndPostLayerWeight 
+    pre_and_post_weight_class = Internlm2PreAndPostLayerWeight
     transformer_weight_class = Internlm2TransformerLayerWeight
 
     def __init__(self, kvargs):
         super().__init__(kvargs)
-    

lightllm/models/internlm2_reward/model.py

@@ -1,14 +1,15 @@
 import os
 import json
 import torch
-
+from lightllm.models.registry import ModelRegistry, is_reward_model
 from lightllm.models.internlm2_reward.layer_infer.post_layer_infer import Internlm2RewardPostLayerInfer
 from lightllm.models.internlm2_reward.layer_weights.pre_and_post_layer_weight import (
     Internlm2RewardPreAndPostLayerWeight,
 )
 from lightllm.models.internlm2.model import Internlm2TpPartModel
 
 
+@ModelRegistry("internlm2", condition=is_reward_model())
 class Internlm2RewardTpPartModel(Internlm2TpPartModel):
     # weight class
     pre_and_post_weight_class = Internlm2RewardPreAndPostLayerWeight