feat: 实现模型动态注册机制

astrai/model/__init__.py

@@ -6,5 +6,17 @@ from astrai.model.module import (
     RMSNorm,
 )
 from astrai.model.transformer import Transformer
+from astrai.model.automodel import AutoModel
 
-__all__ = ["Linear", "RMSNorm", "MLP", "GQA", "DecoderBlock", "Transformer"]
+
+__all__ = [
+    # Modules
+    "Linear",
+    "RMSNorm",
+    "MLP",
+    "GQA",
+    "DecoderBlock",
+    # Models
+    "Transformer",
+    "AutoModel",
+]

astrai/model/automodel.py

@@ -0,0 +1,134 @@
+"""
+AutoModel base class for model loading and saving.
+"""
+
+import torch.nn as nn
+import safetensors.torch as st
+
+from pathlib import Path
+from contextlib import contextmanager
+from typing import Self, Union, Dict, Type
+
+from astrai.config import ModelConfig
+
+
+@contextmanager
+def _disable_random_init(enable: bool = True):
+    init_functions = [
+        "xavier_normal_",
+        "xavier_uniform_",
+        "kaiming_normal_",
+        "kaiming_uniform_",
+        "zeros_",
+        "ones_",
+        "constant_",
+        "normal_",
+        "uniform_",
+    ]
+    original_funcs = {}
+    for name in init_functions:
+        if enable and hasattr(nn.init, name):
+            original_funcs[name] = getattr(nn.init, name)
+            setattr(nn.init, name, lambda *args, **kwargs: None)
+    try:
+        yield
+    finally:
+        if enable:
+            for name, orig_func in original_funcs.items():
+                setattr(nn.init, name, orig_func)
+
+
+class AutoModel(nn.Module):
+    """
+    Autoregressive language model base class.
+    Provides model loading/saving and generation capabilities.
+    """
+
+    # Model registry - stored as class attribute
+    _registry: Dict[str, Type["AutoModel"]] = {}
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+
+    @classmethod
+    def register(cls, model_type: str):
+        """
+        Class method decorator to register model type.
+
+        Usage:
+            @AutoModel.register('transformer')
+            class Transformer(AutoModel):
+                ...
+        """
+
+        def decorator(sub_cls: Type["AutoModel"]) -> Type["AutoModel"]:
+            cls._registry[model_type.lower()] = sub_cls
+            return sub_cls
+
+        return decorator
+
+    @classmethod
+    def get_model_class(cls, model_type: str) -> Type["AutoModel"]:
+        """Get model class by model_type string."""
+        model_type = model_type.lower()
+        if model_type not in cls._registry:
+            available = list(cls._registry.keys())
+            raise ValueError(
+                f"Unknown model_type: {model_type}. Available: {available}"
+            )
+        return cls._registry[model_type]
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        path: Union[str, Path],
+        disable_random_init: bool = True,
+    ) -> nn.Module:
+
+        model_path = Path(path)
+
+        # Load config
+        config = ModelConfig()
+        config_path = model_path / "config.json"
+        if config_path.exists():
+            config.load(str(config_path))
+        else:
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+
+        # If called from base class, use model_type to determine actual model class
+        if cls is AutoModel:
+            model_type = config.model_type or "transformer"
+            actual_cls = cls.get_model_class(model_type)
+        else:
+            raise ValueError(
+                f"Cannot call from_pretrained() on subclass {cls.__name__}"
+            )
+
+        with _disable_random_init(enable=disable_random_init):
+            model = actual_cls(config)
+
+        # Load weights
+        weights_path = model_path / "model.safetensors"
+        if weights_path.exists():
+            state_dict = st.load_file(str(weights_path))
+            model.load_state_dict(state_dict, strict=False)
+
+        return model
+
+    def save_pretrained(
+        self,
+        save_directory: Union[str, Path],
+    ) -> None:
+        save_path = Path(save_directory)
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        # Save config
+        self.config.save(str(save_path / "config.json"))
+
+        # Save weights
+        st.save_file(self.state_dict(), str(save_path / "model.safetensors"))
+
+    def to(self, *args, **kwargs) -> Self:
+        """Move model to device/dtype."""
+        return super().to(*args, **kwargs)

astrai/model/transformer.py

@@ -5,6 +5,7 @@ import torch.nn as nn
 from torch import Tensor
 
 from astrai.config.model_config import ModelConfig
+from astrai.model.automodel import AutoModel
 from astrai.model.module import (
     DecoderBlock,
     Embedding,
@@ -66,9 +67,14 @@ def process_attention_mask(
     return attention_mask
 
 
-class Transformer(nn.Module):
+@AutoModel.register("transformer")
+class Transformer(AutoModel):
+    """
+    Transformer language model.
+    """
+
     def __init__(self, config: ModelConfig):
-        super().__init__()
+        super().__init__(config)
         self.config = config
         self.rotary_embeding = RotaryEmbedding(
             config.dim // config.n_heads, config.max_len
@@ -97,16 +103,27 @@ class Transformer(nn.Module):
         if self.config.tie_weight:
             self.lm_head.weight = self.embed_tokens.weight
 
-        self._init_parameters()
+        self._init_weights()
+
+    def _init_weights(self):
+        for param in self.parameters():
+            if param.dim() > 1:
+                nn.init.normal_(param, mean=0.0, std=0.006)
 
     def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
         lm_head_key = "lm_head.weight"
         embed_key = "embed_tokens.weight"
 
+        # Make a copy to avoid modifying the original state_dict
+        state_dict = dict(state_dict)
+
         if self.config.tie_weight:
             # same tensor
-            state_dict[lm_head_key] = state_dict[embed_key]
+            if embed_key in state_dict:
+                state_dict[lm_head_key] = state_dict[embed_key]
         else:
+            # If lm_head.weight exists in checkpoint, use it directly
+            # If not, copy from embed_tokens.weight
             if lm_head_key not in state_dict and embed_key in state_dict:
                 # use clone to avoid sharing the same tensor
                 state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
@@ -125,11 +142,6 @@ class Transformer(nn.Module):
 
         return state_dict
 
-    def _init_parameters(self):
-        for param in self.parameters():
-            if param.dim() > 1:
-                nn.init.normal_(param, mean=0.0, std=0.006)
-
     def forward(
         self,
         input_ids: Tensor,